feat(moe): enable CPU optimizer offload for Kimi-K3 (cuda:nccl,cpu:gloo) - #2
Open
xiaoliang0601 wants to merge 3 commits into
Open
xiaoliang0601 wants to merge 3 commits into
xiaoliang0601 wants to merge 3 commits into
Conversation
…ckend Enable CPUOffloadPolicy for MoE/EP training (e.g. Kimi-K3 fp32-master) by making three components CPU-offload aware: - moe/experts: stream EP-sharded expert weights to the activation device (not just the activation dtype). Under CPUOffloadPolicy the local expert shard lives on CPU and is never FSDP-all-gathered back to GPU, so the grouped-mm operands must be moved to x.device before the CUDA kernel runs. - distributed/mesh_utils: when the default process group carries a cpu:gloo co-backend (offload signal), mirror it onto every per-axis DeviceMesh subgroup as cuda:nccl,cpu:gloo so DCP checkpoint save/load can run DTensor collectives on CPU-resident shards; NCCL-only subgroups reject them. The NCCL timeout is preserved for the CUDA collectives. - training/signal_handler: resolve get_device() from a device-typed combined backend string (cuda:nccl,cpu:gloo) via substring match, preferring CUDA. Adds focused unit tests for each behavior. Signed-off-by: root <xuezunyao@xiaohongshu.com>
dequantize_base_checkpoint describes the base HF checkpoint being MXFP4-packed, so requesting packed .weight_packed/.weight_scale load destinations is only correct when initializing from that base (is_init_step=True). Training checkpoints are always saved dequantized (save_model uses quantization=False), so a resume/restore (is_init_step=False) must read plain .weight keys; asking for packed destinations there looks for keys the checkpoint lacks and fails the DCP load plan. Gate the adapter's quantization flag on is_init_step. Signed-off-by: root <xuezunyao@xiaohongshu.com>
Under fp32-master weights (torch_dtype: float32) with bf16 mixed-precision compute, lm_head.weight is stored in fp32 but FSDP2 casts it to bf16 for the forward, so the emitted logits are bf16. The last-stage pipeline output meta must follow the pipeline compute dtype, not the weight storage dtype; otherwise PP shape inference expects fp32 while the real output is bf16 and raises PipeliningShapeError. Drop the weight-derived tensor_dtype so the meta inherits the compute dtype. Adds a unit test asserting the last-stage logits meta follows the pipeline dtype (bf16) rather than the fp32 weight storage dtype. Signed-off-by: root <xuezunyao@xiaohongshu.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.