perf(bagel): accelerate agent SFT replay with FlexAttention

[zzhuoxin1508]

Summary

• Add a replay_attention_backend switch for Bagel train replay: sdpa preserves the existing dense per-sample mask path, while flex builds a semantically equivalent BlockMask for FlexAttention.
• Thread the option through BagelPipeline and BagelARStage with validation; only the train-replay path consumes it.
• Import FlexAttention only when flex is selected, preserving a clear missing-version error without affecting the SDPA path.
• Enable Flex only in the long interleaved bagel_agent_sft recipe while retaining SDPA as the generic default; no other recipe or inference-replay path changes behavior.
• This change keeps the existing single-trajectory replay (micro_batch_size=1); it does not add multi-sample packing.

Benchmark results

Interleaved image-text agent SFT

50-step performance A/B

Clean 50-step A/B on 8x NVIDIA H20, global batch 8, micro-batch 1, same deterministic 72-example SearchGen manifest, with profiler/W&B/eval/checkpointing disabled:

• Steady steps 11-50: SDPA 9.4525s/step, Flex 5.9275s/step — 1.595x speedup.
• All 50 steps: SDPA 9.460s/step, Flex 6.434s/step — 1.470x speedup, including compilation.
• Flex first step: 23.1s; estimated one-time compile excess: 16.35s.
• Mean sampled GPU memory: SDPA 74,338 MiB, Flex 60,450 MiB; peaks: 88,065 MiB vs 68,587 MiB.
• Loss correlation: 0.99993; gradient-norm correlation: 0.97662; all values finite.

Conclusion: long SearchGen trajectories contain multiple interleaved full-image and causal-text regions, leaving enough fully invalid 128x128 tiles for Flex to skip substantial work.

1000-step quality parity

A 1000-step quality run compared this PR at b7db6c5 (Flex replay) with the merge-base main commit 25a2fa5 (SDPA replay), using the same cooked SearchGen manifests and the same training settings documented below.

• The train/sft_ppl curves overlap throughout the 1000 steps.
• Final validation loss: Flex 1.27332, SDPA 1.26963 — an absolute difference of 0.00369 (0.29% relative).
• Both runs completed all 1000 optimizer steps and wrote their step-1000 LoRA adapter checkpoints.
• These long runs are used for quality comparison only; wall-clock timing is intentionally not reported because the PR run had external GPU keepalive interference.

Conclusion: the long-run loss curves and final validation loss show no material SFT quality regression from using FlexAttention replay.

1000-step SFT test setup

Files and parameters

• Tested revision: b7db6c5cba373f0b872a71136a3819f32d3f9a1c; the current SFT path is behavior-equivalent, with only scope cleanup and a rebase onto latest main afterward.
• Entrypoint and config: python -m unirl.train_sft (unirl/train_sft.py), using examples/sft/bagel_agent_sft.yaml with the long-run overrides below. SearchGen manifests were prepared by datasets/searchgen/prepare_sft.py.
• Data: SearchGen-20K real interleaved image/text agent trajectories; max_traces=5000, candidates_per_query=4, max_image_px=448, max_target_chars=8000, val_fraction=0.02, and seed 42. The cooked manifests contain 8,247 train records and 132 validation records.
• Model and hardware: BAGEL-7B-MoT on 8x NVIDIA H20; understanding ViT enabled; BF16 parameter/compute precision with FP32 master weights, trajectory, and log-prob precision.
• Run shape: 1,000 optimizer steps, global batch size 8, micro-batch size 1, max prompt length 24576, max response length 2048, token-mean SFT loss, and gradient clipping at 1.0.
• Replay path: replay_mode=train and replay_attention_backend=flex.
• LoRA: rank 64, alpha 128, dropout 0; targets are the understanding-expert q/k/v/o projections and MLP gate/up/down projections.
• Optimizer and schedule: learning rate 1e-4, betas (0.9, 0.999), epsilon 1e-8, weight decay 1e-4, constant schedule, and no warmup.
• Distributed/checkpointing: full-shard FSDP, reshard-after-forward, activation checkpointing, no CPU offload or torch.compile; evaluation every 100 steps (eval_batch_size=8, eval_num_samples=64); adapter checkpoint saved at step 1,000.
• Logging: W&B enabled under project unirl-sft, run name pr383-bagel-flex-sft1000-searchgen5k-8xh20-20260824.

Test Plan

• Compile the changed Bagel Python modules.
• Resolve the agent SFT config and verify its default backend is flex.
• Verify lazy FlexAttention import and clear missing-version errors.
• Run the clean 50-step 8-GPU SFT performance A/B benchmark.
• Run the 1000-step 8-GPU SFT quality comparison against the main SDPA baseline.

[zzhuoxin1508]

sft test result