gpt-oss model enablement

[wwwjn]

Keep developing on top of #1559. Thanks @KhoomeiK for initial contribution!

Initialized by the same seed checkpoint, set seed=0 and deterministic = True.

GPT-oss
Run 1: dp_shard = 2

Run 2: dp_shard = 2, TP degree = 2 (NGPU=4)

Run 3: dp_shard = 2, TP degree =2, EP degree = 2 (NGPU=4)

Run 4: dp_shard = 2, TP degree = 2, EP degree = 2, ETP degree = 2 (NGPU=4)

Run 5: dp_shard=2, EP degree = 2 (NGPU=2)

[wwwjn]

Need to rebase onto #1776

[tianyu-l]

Looks great in general. Left some comments. May need some rebase on recent & near-future development.

[wwwjn]

Summary of current status:

There are some prerequisite PRs:

1. FlexAttn refactor Refactor attention and make attention mask an argument to the model #1776
2. EP refactor [EP] add initial support for NVSHMEM-based all-to-all #1569
3. refactor freqs_cis as a input of model.forward() [RFC] Lift freqs_cis as an input of models #1797

Once these PRs are landed, I will refactor:

1. FlexAttention, adding sliding_window attention mask, and make it orthogonal to block_causal mask.
2. ExpertParallel() and ExpertTensorParallel() class to reuse as much as possible, as keep aligns with main EP implementation

[fegin]

Please address all the comments before landing. I would appreciate that if you add the reason why we cannot do AuxOutput to the code. Thanks!

[wwwjn]

I explicitly leave the parameter here, to be compatible with https://github.com/pytorch/torchtitan/blob/refs/heads/main/torchtitan/train.py#L428 here, where we need to call get_attention_masks.

But I added a notes here to prevent user change this flag to false

[tianyu-l]

I think I found a tricky numerical bug in TP. Maybe we can disable it for now.

[tianyu-l]

need to raise ValueError if user didn't set window_size >= 1

[tianyu-l]

This might not be correct.

When they are dtensors, x * mlp2_weight + mlp2_bias will have placements Partial + Replicate, and sharding prop can automatically first make Replicate -> Partial then perform the addition.

However, when we do to_local, DTensor placement info is discarded, so instead of adding mlp2_bias, the net effect will be adding tp_degree * mlp2_bias.

I don't have clean way to solve this. For forward correctness, we can do mlp2_bias / tp_degree to cancel the extra reduction effect, but the backward will have an extra * tp_degree. Can we wrap mlp2_bias / tp_degree in torch.no_grad so the backward doesn't perform * tp_degree?

You can also disable TP / ETP altogether for gpt-oss for now and leave a TODO.

cc @ezyang @fmassa on difficulties of making TP correct in a local tensor region, when there is bias involved.

[tianyu-l]

This is different from main moe.py, where we init the weight params to have shape (num_experts, out_dim, in_dim) and do transpose before using them. The point is for hardware efficiency (mainly in low-precision case). We also need to change the TP / ETP plans to adapt.

See #1517

[drisspg]

yeah can you explain this?

[drisspg]

This API will be removed in a future release

[wwwjn]

Thanks! I also noticed the return_lse is being deprecated, the reason we use it here is we want to use TP annotation to change the lse tensor back to a DTensor with placement Shard(1) (in TP region, it's a plain tensor). https://github.com/pytorch/torchtitan/pull/1754/files#diff-3448dcaf6e8b68f3b66a8e1dd298273de3702f93de406569426cd9e03fd7f97bR222. We can not annotate an AuxOutput() object directly using TP APIs. And because we want to keep model code parallelism-free, we don't want to manually turn AuxOut.lse into a DTensor.

I think an alternative way is to handle it in FlexAttentionWarpper, if this is a better way, I will create another PR to fix.

[tianyu-l]

LGTM, some minor final comments

[tianyu-l]

comment not addressed

[tianyu-l]

Please add a comment to indicate which dim is input dim, which is output dim.