Fix layer-wise distributed optimizer MXFP8 wgrad failure by padding per-param start - #4889
Fix layer-wise distributed optimizer MXFP8 wgrad failure by padding per-param start#4889Wohox wants to merge 2 commits into
Conversation
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| for param in params[::-1]: | ||
| if pad_param_starts: |
There was a problem hiding this comment.
Nit: The existing TestDefaultParamLayout class in tests/unit_tests/distributed/test_param_layout.py tests this function but only exercises the default (no-padding) path. Since this is a bug fix, consider adding a regression test that calls _compute_default_per_buffer_param_layout(..., pad_param_starts=True) with params whose sizes aren't multiples of 64, and asserts that every param_start_index in the resulting layout is 64-element-aligned. This would prevent the alignment invariant from silently regressing.
There was a problem hiding this comment.
LGTM — the fix correctly reuses the existing pad_param_start helper, the condition in training.py precisely targets the affected code path, and the parameter threading follows established codebase patterns. One minor suggestion to add a regression test for the new pad_param_starts=True path (left as an inline comment).
|
Closed since already included in #4771 |
2 participants
What does this PR do ?
Pad each parameter's start index up to a 64-element boundary in the default
_ParamAndGradBufferlayout when the layer-wise distributed optimizer (Muon and other non-Adam/SGD optimizers) is active, so thatmain_gradslices reach the cuBLASLt-required D-pointer alignment for MXFP8 wgrad. Without this, mid-bucket params land at 16-byte D-pointer alignment andcublasLtMatmulAlgoGetHeuristicreturnsCUBLAS_STATUS_NOT_SUPPORTED, killing the training job.This is the
main-branch counterpart of #4840 (devbranch). The semantics are identical; the diff is rebased ontomainand the plumbing site is thewrap_model_chunks_with_ddphelper added in #4623, which centralises DDP wrapping onmain.Issue tracking
Linked issue: N/A (bug fix)
Details
Symptom. Training a model with
--use-precision-aware-optimizerMXFP8 wgrad + layer-wise distributed optimizer (use_layer_wise_distributed_optimizer=True) crashes inside cuBLASLt on the first wgrad call for a non-first param in any bucket:Root cause. Under the layer-wise distributed optimizer,
arguments.pysetsuse_distributed_optimizer=Falseand routes through the legacyallgather_paramssync path, sowrap_model_chunks_with_ddpdoes not supply a pre-paddedfull_param_layoutto DDP. DDP then falls back to_compute_default_per_buffer_param_layout, which packs params back-to-back with no padding. The resultingmain_gradslice for any param whose offset isn't a multiple of 64 elements (≥128 bytes for BF16/MXFP8) has only 16-byte D-pointer alignment, which cuBLASLt 12.8.x rejects for MXFP8 wgrad.The standard distributed optimizer is unaffected because it precomputes a shard-aligned
full_param_layoutthat already aligns per-param starts.Fix.
param_and_grad_buffer.py— addpad_param_starts: bool = Falseto_compute_default_per_buffer_param_layoutand_ParamAndGradBuffer.__init__. When set, each param'sparam_start_indexis rounded up to a 64-element boundary viapad_param_start(the same helper the distributed optimizer already uses). Cost: ≤ 126 bytes per param.distributed_data_parallel.py— forward the newpad_param_startskwarg fromDistributedDataParallel.__init__down to_ParamAndGradBuffer.training.py(wrap_model_chunks_with_ddp) — whenDP is DDP and use_layer_wise_distributed_optimizer=True and not use_layer_wise_param_layout(i.e. the path that supplies nofull_param_layout), passpad_param_starts=Trueto every wrapped chunk. The flag is ignored when afull_param_layoutis supplied, so callers that already provide a layout are unaffected.Contribution process
Pre-checks
Test plan
configs/benchmarking/recipes/deepseek_v4_flash_proxy/gb200/mxfp8_*muon*.yaml) onmainwithout this patch →CUBLAS_STATUS_NOT_SUPPORTEDon iter 1.--use-distributed-optimizer(standard) path is untouched (still supplies its own paddedfull_param_layout); BF16 layer-wise path also runs cleanly (padding is benign).Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
🤖 Generated with Claude Code