Repository navigation
(feat) Fold the CP axis into the GTP_remat weight-sharding group - #6886
fanshiqing wants to merge 5 commits into
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
/claude review |
|
/ok to test 5e799be |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
|
/claude strict-review |
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
CP ranks of a CP-folded GTP weight shared one checkpoint key, so their optimizer state silently aliased on save/load. Key now includes the CP rank for those buffers only. Co-authored-by: Dingqing Yang <dingqingy@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Signed-off-by: Dingqing Yang <dingqingy@nvidia.com>
CP folding into the GTP_remat weight group is now opt-in, at any GTP_remat size including 1. Also fix init seeds and checkpoint writer election for these modes, and restore a broken test_gtp_muon.py. Co-authored-by: Dingqing Yang <dingqingy@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Signed-off-by: Dingqing Yang <dingqingy@nvidia.com>
cda0ad5 to
9de4bef
Compare
|
/ok to test 9de4bef |
|
/claude review |
This comment was marked as resolved.
This comment was marked as resolved.
- LayerWiseDistributedOptimizer.count_zeros(): drop the gtp_remat_group / expt_gtp_remat_group kwargs that count_zeros_fp32 never accepted (TypeError on --log-num-zeros-in-grad), plus the now-dead attributes. - DistributedOptimizer: reject CP-folded GTP buffers in the dp_zero / fully_reshardable save and load paths. They gather over the global DP group keyed by gbuf_idx, which can't represent per-CP-rank distinct shards; point users to dp_reshardable. - gtp_replica_rank: fix the stale error message so it names the right replica group (dp / dp_cp / expt_dp). - Docs: fix stale TOC entries for section 3.8. - Add regression tests for count_zeros and the fully_reshardable rejection. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test cd54188 |
|
/claude review |
There was a problem hiding this comment.
Light review — read the diff, skills/mcore-testing/SKILL.md, skills/pr-review/references/light.md, and style-guide.md.
The design is coherent and unusually well-documented: the CP-free vs CP-folded group split, the per-param gtp_bucket_dp_divisor / excludes_cp_from_bucket stamps, and the .cp_rank_{i} optimizer-key suffix all line up with each other and with 3.8. Test coverage for the new behavior is genuinely good — grad parity against an exact fp32 reference, writer-election gap checks, rank-distinct optimizer round-trips, and a fully_reshardable rejection test.
One blocking issue.
--gtp-remat-fold-cp is never registered with argparse (critical). megatron/training/initialize.py:405 reads args.gtp_remat_fold_cp, but no add_argument for it exists in arguments.py or yaml_arguments.py. Since that line is unconditional, every pretrain_*.py run that reaches _initialize_distributed raises AttributeError — regardless of CP or GTP settings. The unit tests pass gtp_remat_fold_cp straight to ps.initialize_model_parallel() and never go through arguments.py, which is why the suite is green. Details and a suggested add_argument block are in the inline comment.
This also means the is_gtp_remat_active(args) extension in common_utils.py:607 and the validate_args gate at arguments.py:1545 can never observe a true gtp_remat_fold_cp, so none of the GTP validation (--fp4-param-gather, --ddp-average-in-collective off, --ckpt-format) currently applies to a fold-CP-only config once the flag does exist. Worth confirming those gates behave as intended at gtp_weight_remat_size == 1 with folding on — the one shape where folding alone activates dense GTP.
Notes, non-blocking, no action needed:
distributed_data_parallel.py:292calls theaverage_in_collectivebranch unreachable forexcludes_cp_from_bucket. That holds: the assert at line 245 rejectsaverage_in_collectivewheneveris_gtp_remat_active(process_group_dict), and a folded group always hassize() > 1. The comment is accurate.resolve_buffer_dp_world_sizedivides the caller-supplieddp_cpsize by the divisor to recoverdp, matching the group DDP assigns at line 321. The two stay consistent.- No new direct
parallel_state.get_*_group()reads outside the allowed compatibility points — the newget_gtp_weight_remat_group_no_cpaccessor lives inparallel_state.pyand is plumbed throughProcessGroupCollection, withfinalize_model_gradspreferringpg_collection.gtp_remat_no_cpand falling back only when absent. Therandom.pyandgeneralized_tensor_parallelism.pyMPU reads replace existing ones rather than adding new coupling. - No GPU kernel changes, so no
tests/unit_tests/determinism/kernels/entry is required.
The PR description already lists 8-GPU shapes and an end-to-end CP+GTP_remat training run as outstanding; the argparse gap is presumably why the latter has not run yet.
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
What does this PR do?
Adds GTP_remat + Context Parallelism (CP) support: optionally fold the CP axis into the
dense GTP_remat weight-sharding group, so CP ranks also shard dense weights.
--gtp-remat-fold-cp(default off): dense weight group becomescp × gtp_remat(at GTP_remat 1: the CP group). Off keeps the plain
gtp_rematgroup, giving a cleanwith/without-folding baseline on the same config. Expert weights never fold.
(
excludes_cp_from_bucket); checkpoint writer election excludes CP when folded, includes itwhen not; init seeds use the folded group.
dp_reshardablekeys of folded buffers get a.cp_rank_{i}suffix — previously CP ranks' FP32 masters / Adam moments aliased on one key. Other keys
unchanged;
fully_reshardableunaffected.Test plan
dp_reshardable+ model weights), folded and unfoldedConvergence Test
To be added.
Issue tracking
Linked issue: Related to #6844
Contribution process
Pre-checks