Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BE] replace the extra DeviceMesh _flatten with mesh access #666

Merged
merged 2 commits into from
Oct 31, 2024
Merged

[BE] replace the extra DeviceMesh _flatten with mesh access #666

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
47356dc
[BE] replace the extra DeviceMesh _flatten with mesh access
XilunWu Oct 30, 2024
e35e17d
Update on "[BE] replace the extra DeviceMesh _flatten with mesh access"
XilunWu Oct 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 1 addition & 15 deletions torchtitan/parallelisms/parallelize_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,21 +81,7 @@ def parallelize_llama(
parallel_dims.dp_shard_enabled
): # apply FSDP or HSDP, potentially with Context Parallel

# TODO: instead of flattening the mesh twice, we could've done in a batter way:
# dp_mesh = world_mesh["dp_cp"] if parallel_dims.cp_enabled else world_mesh["dp"]
# However, this leads to an error in `DeviceMesh.__get_item__` which I believe is
# a bug in DeviceMesh. We should fix it and then use the above line.
dp_mesh_dim_names = (
("dp_replicate", "dp_shard")
if parallel_dims.dp_replicate_enabled
else ("dp",)
)
# note that mesh can only be flattened from the finest-grained mesh dimensions
dp_mesh = (
world_mesh[(*dp_mesh_dim_names, "cp")]._flatten("dp_cp")
if parallel_dims.cp_enabled
else world_mesh[dp_mesh_dim_names]
)
dp_mesh = world_mesh["dp_cp"] if parallel_dims.cp_enabled else world_mesh["dp"]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a new DeviceMesh functionality that reacts specifically to <name1>_<name2>?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not new. DeviceMesh supports world_mesh[<name1>_<name2>] when the _flatten behavior was implemented. However, it has a bug -- if the flattened mesh is constructed from 3+ mesh dimensions (e.g. dp_cp is flattened using dp_shard, dp_replicate, and cp. Accessing world_mesh[dp_cp] throws error which breaks 3D/4D/5D composability).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was flattened here
https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallel_dims.py#L82

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we catch the error and ask users to update to some version?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For my understanding, for dp, if hsdp is enabled, "dp" is the flatten mesh for "dp_replicate", "dp_shard", right? Otherwise, "dp" is just "dp_shard".

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@wz337 , that's right. To summarize:

  1. FSDP: the only dp dimension in mesh is "dp"
  2. DDP: the only dp dimension in mesh is "dp"
  3. HSDP: the basic dp dimensions in mesh are "dp_shard" and "dp_replicate", which are later on flattened into "dp"


apply_fsdp(
model,
Expand Down
Loading