Skip to content

feat: HYBRID_SHARD draft training across colocated target islands - #784

Draft
maocheng23 wants to merge 1 commit into
maocheng/colocate-4-online-corefrom
maocheng/colocate-5-hybrid-shard
Draft

feat: HYBRID_SHARD draft training across colocated target islands#784
maocheng23 wants to merge 1 commit into
maocheng/colocate-4-online-corefrom
maocheng/colocate-5-hybrid-shard

Conversation

@maocheng23

Copy link
Copy Markdown
Collaborator

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_SHARD 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.

Modifications

  • training.fsdp_sharding accepts HYBRID_SHARD, valid only for online local_colocated with training.tp_size > 1 and at least two target islands.
  • ParallelConfig carries explicit reduction_process_group (WORLD — the HSDP fsdp_process_group is a (shard, replica) pair and cannot serve as a collective group) and grad_norm_process_group (the shard group, counting each replicated parameter shard exactly once). TrainerCore reduces 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.
  • Adds the Kimi K3 4x8 B300 TP8/HSDP starting recipe and the HSDP-across-islands documentation section.

Related Issues

Splits #766. Stack: #1 capture-rows ← #2 teardown-abort ← #3 rank0-tracker ← #4 colocated-core ← #5 (this).

Accuracy Test

  • HSDP is configuration- and process-group-tested (ParallelConfig.from_distributed group 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

  • Format your code according to the Code Formatting with Pre-Commit (black --check and isort --check-only).
  • Add unit tests.
  • Update documentation and example recipes.
  • Validate Kimi K3 with TP8/HSDP on 4x8 B300 before marking ready for review.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant