[samsja]

couple of thing that I noticed are different between us and torchtitan:

• they do torch compile with full graph on each transformer block https://github.com/pytorch/torchtitan/blob/40a87254ff15dbabebcfe8f70828a872cc6fa009/torchtitan/models/llama3/infra/parallelize.py#L243

• they do ac - > compile -> fsdp wheras we do ac-> fsdp -> compile which could create graph break https://github.com/pytorch/torchtitan/blob/40a87254ff15dbabebcfe8f70828a872cc6fa009/torchtitan/models/llama3/infra/parallelize.py#L106

• They shard explicitly the lm head and embedding, we don't, potentially meaning we don't shard this layers ? https://github.com/pytorch/torchtitan/blob/40a87254ff15dbabebcfe8f70828a872cc6fa009/torchtitan/models/llama3/infra/parallelize.py#L308

• They don't to the "do not reshard after forward for the last layer" trick, which they used to do (bc we stole it from there) My guess is that it breaks full graph for torchtitan https://github.com/pytorch/torchtitan/blob/40a87254ff15dbabebcfe8f70828a872cc6fa009/torchtitan/models/llama3/infra/parallelize.py#L301

┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃    Step ┃         MFU          ┃      Throughput      ┃      Step Time       ┃
┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│       1 │        39.50%        │        63.74K        │        2.05s         │
│       2 │        39.39%        │        63.56K        │        2.05s         │
│       3 │        39.32%        │        63.45K        │        2.05s         │
│         │                      │                      │                      │
│ Overall │    39.40% ± 0.09%    │   63.58K ± 147.78    │    2.05s ± 0.00s     │
│         │   [39.32%, 39.50%]   │   [63.45K, 63.74K]   │    [2.05s, 2.05s]    │
└─────────┴──────────────────────┴──────────────────────┴──────────────────────┘

branch sami/update_fsdp_orderring

┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃    Step ┃         MFU          ┃      Throughput      ┃      Step Time       ┃
┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│       1 │        27.29%        │        44.03K        │        2.97s         │
│       2 │        27.32%        │        44.08K        │        2.96s         │
│       3 │        27.32%        │        44.08K        │        2.97s         │
│         │                      │                      │                      │
│ Overall │    27.31% ± 0.02%    │    44.07K ± 27.75    │    2.96s ± 0.00s     │
│         │   [27.29%, 27.32%]   │   [44.03K, 44.08K]   │    [2.96s, 2.97s]    │
└─────────┴──────────────────────┴──────────────────────┴──────────────────────┘

reproduce with

srun bash -c '
    source /data/.env
    cd /data/prime-rl
    
    uv run torchrun \
    --nnodes=1 \
    --nproc_per_node=8 \
    --node-rank $SLURM_PROCID \
    src/prime_rl/trainer/sft/train.py \
    --model.name Qwen/Qwen3-8B\
    --model.attn sdpa \
    --model.ac \
    --model.compile \
    --data.type fake \
    --data.seq-len 2048 \
    --data.batch-size 64 \
    --data.micro-batch-size 8 \
    --data.pack-function stack \
    --log.level debug \
    --bench 2>&1 | tee /data/prime-rl/outputs/logs/latest_rank${SLURM_PROCID}.log /data/prime-rl/outputs/logs/torchrun_${SLURM_JOB_ID}_rank${SLURM_PROCID}.log'
    

[mikasenghaas]

banger