feat: HYBRID_SHARD draft training across colocated target islands - #784
Draft
maocheng23 wants to merge 1 commit into
Draft
feat: HYBRID_SHARD draft training across colocated target islands#784maocheng23 wants to merge 1 commit into
maocheng23 wants to merge 1 commit into
Conversation
K3-class colocated runs put one TP8 target island on each node. Sharding the draft across the whole world would make every FSDP parameter all-gather cross nodes; HYBRID_SHARD instead uses the target-TP group as the FSDP shard group and the target-DP group for replication, keeping parameter traffic node-local while only replica synchronization crosses nodes. ParallelConfig now carries explicit reduction and grad-norm groups: loss/metric reductions stay WORLD-wide (the HSDP fsdp_process_group is a (shard, replica) pair and cannot be used as a collective group), and the gradient norm reduces over the shard group only, counting each replicated parameter shard exactly once. Config validation accepts HYBRID_SHARD only for online local_colocated runs with training.tp_size > 1 and at least two target islands. Adds the Kimi K3 4x8 B300 TP8/HSDP starting recipe. HSDP is configuration- and process-group-tested; a full 4x8 B300 run is still pending. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
5 tasks
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
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.
Motivation
Stack 5/5 replacing draft #766. K3-class colocated runs put one TP8 target island on each node; sharding the draft across the whole world would make every FSDP parameter all-gather cross nodes.
HYBRID_SHARDuses the target-TP group as the FSDP shard group and the target-DP group for replication, keeping parameter traffic node-local while only replica synchronization crosses nodes.Modifications
training.fsdp_shardingacceptsHYBRID_SHARD, valid only for onlinelocal_colocatedwithtraining.tp_size > 1and at least two target islands.ParallelConfigcarries explicitreduction_process_group(WORLD — the HSDPfsdp_process_groupis a(shard, replica)pair and cannot serve as a collective group) andgrad_norm_process_group(the shard group, counting each replicated parameter shard exactly once).TrainerCorereduces losses/metrics through a single_reduction_group()helper with the FSDP-group fallback for every non-HSDP topology — feat: add scalable colocated online training #766's__dict__-based mock-dodging lookup is replaced by fixing the test's mock contract instead.Related Issues
Splits #766. Stack: #1 capture-rows ← #2 teardown-abort ← #3 rank0-tracker ← #4 colocated-core ← #5 (this).
Accuracy Test
ParallelConfig.from_distributedgroup wiring, grad-norm group selection, schema validation); the intended 4x8 B300 K3 run has not been executed yet — this PR stays draft until it has, matching feat: add scalable colocated online training #766's open checklist item.Checklist
black --checkandisort --check-only).