Partially-sharded-data-parallel

[blahBlahhhJ]

An explanation:

Here I will ignore the difference between DP and FSDP and assume there're only 2 axis (DP, TP) because I handled all DP, FSDP axis together DP_FSDP_index = DP_index * FSDP_size + FSDP_index

Suppose we have this mesh of DP=n, TP=4: (d -> device. p-> process)

         TP0    TP1     TP2      TP3
DP0    (d0 p0) (d1 p0) (d4 p1) (d5 p1)
DP1    (d8 p2) (d9 p2) (d12 p3) (d13 p3)
  ... a bunch of other DP groups in between so that devices are non-contiguous
DPk     (d2 p0) (d3 p0) (d6 p1) (d7 p1)

The 4 devices (e.g. d0d1d4d5) in every rows gets the same data and performs TP (shards the model). Each column receives different data.

We first take care of each process individually and ensure

• devices in the same DP group (with different TP index) receives the same data
• devices in different DP groups receive different data

This is handled by local_device_mapping().

From each device, we can extract its DP index by its position in the mesh and map it to a uid. In this example, we will have a mapping of

p0: {0: 0, k: 1}
p1: {0: 0, k: 1}
p2: {1: 0, ...}
p3: {1: 0, ...}

When we call make_array_from_callback() with the mesh, each device gets a slice of size per_device_batch_size and thus the slice's start will be global_begin = DP_index * per_device_batch_size. Thus we can extract DP_index from the slice

DP_index = global_begin // per_device_batch_size

See batch_callback() in loader.py, devices will get

local_batch[uid*per_device_batch_size : (uid+1)*per_device_batch_size]

For p0, this means devices with DP index 0 (d0&d1) will get the first half of local_batch, and devices with DP index k (d2&d3) will get the second half. This satisfied the two bullet points above.

Next, we take care of the process-level stuff and ensure

• different processes' devices that are in the same DP group (different TP index) receive the same data.
• different processes' devices that are in different DP groups receive different data.

This is handled by process_mesh_mapping().

Each process:

1. maps to its upper left device in the mesh
2. maps to the location of the device
3. maps to the location of the device with TP0 in the same DP group
4. maps to the device in that TP0 location
5. maps to the process of that device
6. maps to a uid

p0 -> d0 -> (DP0 TP0) -> (DP0 TP0) -> d0 -> p0 -> 0
p1 -> d4 -> (DP0 TP2) -> (DP0 TP0) -> d0 -> p0 -> 0
p2 -> d8 -> (DP1 TP0) -> (DP1 TP0) -> d8 -> p2 -> 1
p3 -> d12-> (DP1 TP2) -> (DP1 TP0) -> d8 -> p2 -> 1

Thus we get this mapping

{p0: 0, p1: 0, p2: 1, p3: 1}

The uid becomes the shard_idx of the dataloader. Thus, p0 and p1 as processes will get the same shard_idx because and will get the same local_batch.

When looking at device-level, we already ensure that d0&d1 gets the first half, d2&d3 gets the second half. This is also true for p1: d4&d5 gets the first half and d6&d7 gets the second half. Now because p0 and p1 have the same shard_idx, we further ensure that d0&d1&d4&d5 gets the same data, d2&d3&d6&d7 gets the same data.

Therefore, these 2 mapping function works for any DP/TP configuration :)

[dlwh]

I'm confused. Can you explain (as a comment) why the data uid stuff is correct?

[dlwh]

could you add a note explaining why this is correct (also so I can be sure I understand)

[blahBlahhhJ]

updated the PR description.

[dlwh]

i don't actually understand how process_mesh is still a valid abstraction/idea in the world with "non-contiguous" devices meshes

[blahBlahhhJ]

updated the PR description.

[dlwh]

also can you merge main so that the TPU tests run

[dlwh]

ok i think i'm convinced!