[kimi2.7] pass correct fsdp mesh for backend=spmd_types - #4063
Conversation
meta-cla
Bot
added
the
CLA Signed
weifengpy
changed the title
weifengpy
changed the title
| reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce], | ||
| reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward, | ||
| pp_enabled=parallel_dims.pp_enabled, | ||
| dp_mesh_dims=None, |
There was a problem hiding this comment.
why this change, maybe we can do this in a separate commit for qwen 3.5?
There was a problem hiding this comment.
from below, seems you are comfortable with dp_mesh_dims: DataParallelMeshDims | None, I can delete this qwen change
| reshard_after_forward_policy: str = "default", | ||
| pp_enabled: bool = False, | ||
| *, | ||
| dp_mesh_dims: DataParallelMeshDims | None, |
There was a problem hiding this comment.
can we make it optional to be consistent with apply_fsdp_to_decoder, i.e., use default value None?
Use the shared dense and sparse storage-mesh resolvers for the spmd_types backend and pass their explicit data-parallel dimensions to FSDP2. Preserve the existing mesh-selection path for the default backend. Test Plan: - Run Moonlight 16B-A3B for 6 steps on 8 B200 GPUs with DP8, EP4, spmd_types, local batch size 4, and sequence length 1024. ghstack-source-id: 3d7cad7 Pull-Request: #4063
Use the shared dense and sparse storage-mesh resolvers for the spmd_types backend and pass their explicit data-parallel dimensions to FSDP2. Preserve the existing mesh-selection path for the default backend. Test Plan: - Run Moonlight 16B-A3B for 6 steps on 8 B200 GPUs with DP8, EP4, spmd_types, local batch size 4, and sequence length 1024. ghstack-source-id: 6ca5512 Pull-Request: #4063
Use the shared dense and sparse storage-mesh resolvers for the spmd_types backend and pass their explicit data-parallel dimensions to FSDP2. Preserve the existing mesh-selection path for the default backend. Test Plan: - Run Moonlight 16B-A3B for 6 steps on 8 B200 GPUs with DP8, EP4, spmd_types, local batch size 4, and sequence length 1024. ghstack-source-id: 3ab807e Pull-Request: #4063
| state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM}, | ||
| in_src_shardings={"input": _REPLICATE_ACT}, | ||
| in_dst_shardings={"input": _REPLICATE_ACT}, | ||
| out_src_shardings=_REPLICATE_ACT, |
There was a problem hiding this comment.
required once model.parallelize() runs under DP-only spmd_types
fixed. added integration test to catch this but let me know if you want such test |
hmm, does this pass with spmd typechecking? or is it just enabling the spmd backend. I can take the typechecking portion if not. |
no, spmd type checking is deliberately ignored, see here: so this is only for enablement. |
| OverrideDefinitions( | ||
| [ | ||
| [ | ||
| # Do not enable --debug.spmd_typechecking: multimodal pixel |
There was a problem hiding this comment.
@pianpwk sorry, no spmd_typechecking this time
clone of a previous approval #4063 before: spmd_types exposes the dense DP axis as dp_shard, but Kimi 2.7 requested the legacy fsdp axis, causing Invalid mesh dim: 'fsdp' during initialization after: explicitly tells FSDP that dp_shard or efsdp is the data-parallel shard axis
3 participants
Stack from ghstack (oldest at bottom):
before: spmd_types exposes the dense DP axis as dp_shard, but Kimi 2.7 requested the legacy fsdp axis, causing Invalid mesh dim: 'fsdp' during initialization
after: explicitly tells FSDP that dp_shard or efsdp is the data-parallel shard axis