Conversation
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces end-to-end context parallelism (CP) support for the GPT-2 path by (a) enabling CP-aware SDPA dispatch for PyTorch flash attention, (b) sharding trainer/evaluator batches across the CP mesh (including global position_ids for RoPE correctness), and (c) registering/configuring a new gpt2_cp model variant with supporting tests and example configs.
Changes:
- Add CP SDPA dispatcher patching + CP input/target sharding utilities.
- Extend GPT-2 forward path to accept
position_idsand forward them into RoPE (RotaryTransform) under CP. - Add registration/configuration for
gpt2_cpplus unit + e2e parity test coverage and CP example configs.
Reviewed changes
Copilot reviewed 25 out of 25 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_trainer_context_parallel.py | Unit tests for trainer-side CP sharding behavior and load balancer propagation. |
| tests/test_registry_components.py | Verifies gpt2_cp model variant registration. |
| tests/models/parallelism/test_context_parallel.py | Unit tests for CP SDPA patching, sharding helper, and CP seq-len validation. |
| tests/fsdp2_parallelization/test_context_parallelism.py | Multi-GPU parity tests comparing CP(/TP/PP) outputs/losses to FSDP2 baseline. |
| tests/fsdp2_parallelization/cp_test_configs/fsdp2_config.yaml | Baseline FSDP2 config for parity tests. |
| tests/fsdp2_parallelization/cp_test_configs/fsdp2_8gpu_nope_config.yaml | 8-GPU baseline config variant used in parity suite. |
| tests/fsdp2_parallelization/cp_test_configs/fsdp2_4gpu_nope_config.yaml | 4-GPU baseline config variant used in parity suite. |
| tests/fsdp2_parallelization/cp_test_configs/fsdp2_4gpu_config.yaml | 4-GPU baseline config variant used in parity suite. |
| tests/fsdp2_parallelization/cp_test_configs/cp_tp_pp_config.yaml | CP+TP+PP test config wiring context_parallel_load_balancer and staged pipeline. |
| tests/fsdp2_parallelization/cp_test_configs/cp_tp_config.yaml | CP+TP test config wiring context_parallel_load_balancer. |
| tests/fsdp2_parallelization/cp_test_configs/cp_pp_config.yaml | CP+PP test config for loss parity (noting RoPE/position_ids PP limitation). |
| tests/fsdp2_parallelization/cp_test_configs/cp_config.yaml | CP-only test config for logit parity. |
| tests/end2end_tests/system_tests/test_context_parallel_parity.py | End-to-end CP vs non-CP training loss parity regression guard. |
| src/modalities/trainer.py | Adds trainer-side CP sequence sharding (including derived position_ids). |
| src/modalities/registry/components.py | Registers new gpt2_cp component variant. |
| src/modalities/models/parallelism/context_parallel.py | Implements CP SDPA dispatch patching and tensor buffer sharding helper. |
| src/modalities/models/model_factory.py | Adds GPT-2 CP model factory path and CP-aware seq-len validation; threads CP LB config. |
| src/modalities/models/gpt2/gpt2_model.py | Adds optional position_ids plumbing through GPT-2 → blocks → attention → RoPE. |
| src/modalities/evaluator.py | Applies the same CP batch sharding path during evaluation. |
| src/modalities/config/config.py | Adds GPT2ModelCPConfig and context_parallel_load_balancer for TP/CP configs. |
| config_files/training/config_lorem_ipsum_long_fsdp2_pp.yaml | Adds/aligns experiments_root_path wiring. |
| config_files/training/config_lorem_ipsum_long_fsdp2_pp_tp.yaml | Adds/aligns experiments_root_path wiring. |
| config_files/training/config_lorem_ipsum_long_fsdp2_pp_tp_cp.yaml | New example config for combined PP+TP+CP setup. |
| config_files/training/config_lorem_ipsum_long_fsdp2_pp_cp.yaml | New example config for PP+CP setup. |
| config_files/training/config_lorem_ipsum_long_fsdp2_cp.yaml | New example config for CP setup. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+85
to
+86
| if batch.samples[sample_key].device.type != "cuda": | ||
| batch.samples[sample_key] = batch.samples[sample_key].to(torch.cuda.current_device(), non_blocking=True) |
Comment on lines
+103
to
+104
| if batch.targets[target_key].device.type != "cuda": | ||
| batch.targets[target_key] = batch.targets[target_key].to(torch.cuda.current_device(), non_blocking=True) |
Comment on lines
+100
to
+105
| def _cuda_batch(seq_len: int = 8) -> DatasetBatch: | ||
| dev = torch.cuda.current_device() | ||
| return DatasetBatch( | ||
| samples={"input_ids": torch.ones(1, seq_len, device=dev, dtype=torch.long)}, | ||
| targets={"target_ids": torch.ones(1, seq_len, device=dev, dtype=torch.long)}, | ||
| ) |
Comment on lines
+35
to
+39
| original_method = CausalSelfAttention.execute_attention | ||
| saved_wrapped = getattr(CausalSelfAttention, "_cp_execute_attention_wrapped", _UNSET) | ||
| saved_mesh = getattr(CausalSelfAttention, "_cp_mesh", _UNSET) | ||
| yield | ||
| CausalSelfAttention.execute_attention = original_method |
Comment on lines
+155
to
+163
| args=( | ||
| world_size, | ||
| 24831, | ||
| base_config_path, | ||
| tmp_path, | ||
| "parity_non_cp", | ||
| False, | ||
| non_cp_losses_path, | ||
| ), |
Comment on lines
+169
to
+178
| TestContextParallelParity._run_training_and_write_losses, | ||
| args=( | ||
| world_size, | ||
| 24832, | ||
| base_config_path, | ||
| tmp_path, | ||
| "parity_cp", | ||
| True, | ||
| cp_losses_path, | ||
| ), |
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.
What does this PR do?
This PR adds end-to-end context parallel support for the GPT2 path. The current implementation requires the usage of pytorch's flash attention when enabling context parallelism
General Changes
Breaking Changes
Checklist before submitting final PR
python tests/tests.py)CHANGELOG_DEV.md)