transformers
[`distributed`] Add context parallelism: `DistributedConfig(cp_size=N)`
#50496

Open

[`distributed`] Add context parallelism: `DistributedConfig(cp_size=N)` #50496

3outeille wants to merge 8 commits into cp-1-implementations from cp-3-cp-size

[distributed] Add context parallelism sizes to DistributedConfig

faf2d3e7

[distributed] Add the cp axis to the device mesh and shard FSDP2 over…

931c06b8

[distributed] Refuse models whose layers or head context parallelism …

cb227c6d

[distributed] Split the sequences across the cp group when the model …

8051af49

[distributed] Refuse context parallelism in Trainer

95e96cb3

[distributed] Document context parallelism in the load and config doc…

5964f5d8

[distributed] Test context parallelism against the same weights on fu…

635bb98e

[distributed] Refuse context parallelism for attention outside the at…

6b2c9674

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone