Use new DeviceMesh unflatten to rewrite parallel_dims #1660
Conversation
meta-cla
bot
added
the
CLA Signed
fegin
force-pushed
the
chienchin/new_device_mesh
branch
7 times, most recently
from
12eca61 to
19e4a23
Compare
|
|
||
| return mesh | ||
| if self._meshes[dim].size() == 1: | ||
| return None |
There was a problem hiding this comment.
Not sure if this will break user expectation. We got asks that DTensor redistribute running on a mesh of size 1 should perform no op.
There was a problem hiding this comment.
But even for current TorchTitan, we won't create any DeviceMesh if the parallelism degree is 1. So it is unclear to me how DeviceMesh with size 1 exists?
There was a problem hiding this comment.
not in torchtitan, in internal
There was a problem hiding this comment.
PyTorch? Then it is okay, right? DeviceMesh still supports the case but TorchTitan makes a stronger assumption in our use case.
fegin
force-pushed
the
chienchin/new_device_mesh
branch
from
19e4a23 to
178bc11
Compare
| fsdp = self.dp_shard * self.cp | ||
| efsdp = fsdp * self.tp // (self.etp * self.ep) | ||
|
|
||
| self._world_mesh = init_device_mesh( |
There was a problem hiding this comment.
does this initialize a world PG?
it may be fine to just ignore this for now in torchtitan, but, i am wondering if users want control over world group creation what would that look like?
There was a problem hiding this comment.
cc., @fduwjj are we able to disable the global PG initialization?
There was a problem hiding this comment.
I think so right now we don't use split, so we can make it a fake pg. But if split is needed then we need to materialize the world PG anyway.
There was a problem hiding this comment.
We should modify FLUX train.py as it's in core now.
@ruisizhang123 let's adapt SimpleFSDP after this PR is merged.
oh it seems being fixed in #1959
|
|
||
| return mesh | ||
| if self._meshes[dim].size() == 1: | ||
| return None |
There was a problem hiding this comment.
not in torchtitan, in internal
| ) | ||
|
|
||
| self._meshes = { | ||
| "pp": dataloading_mesh["pp"], |
There was a problem hiding this comment.
to confirm if things match my expected behavior:
- PG will be created for each sub dimension during
unflatten, unlessbackend_overrideis specified on some dimension with the"fake"backend. flattenwill create a new mesh and a new PG.- slicing will create a new mesh, but reuse the PG created in parent mesh.
| self._world_mesh = init_device_mesh( | ||
| device_type, (self.world_size,), mesh_dim_names=("world",) | ||
| ) | ||
| dataloading_mesh = unflatten_mesh( |
There was a problem hiding this comment.
Curious what will happen if self.pp * batch * self.cp * self.tp != world_size? Will the _unflatten() fail?
This is a demonstration of how parallel_dims will be when using pytorch/pytorch#161224 stack. ghstack-source-id: d29d2e2 Pull-Request: #1885
fegin
force-pushed
the
chienchin/new_device_mesh
branch
from
20910ef to
a67e87a
Compare
fegin
requested review from
allenwang28,
ebsmothers,
joecummings and
pbontrager
as code owners
Summary
This PR utilizes the latest APIs provided by DeviceMesh to simplify the creation of all different meshes.
The design philosophy is as follow: