DeepSpeed
918b6904 - fix(rollout): fix CUDA graph capture producing incorrect logits

Commit
39 days ago
fix(rollout): fix CUDA graph capture producing incorrect logits Transformers' StaticLayer.update() uses an internal cumulative_length counter to determine KV write positions. During CUDA graph capture this counter is frozen and never advances on replay, causing all decode KV writes to land at wrong positions and producing silently incorrect logits (observed as ~62% higher distillation loss in GC ON vs OFF). Introduce DeepSpeedStaticCache whose write position is controlled by an externally-managed tensor (static_cache_position) that the caller updates in-place before each graph replay. The cache also reports the correct sequence length to HF's mask generation via get_seq_length(). Validated with: - Co-simulation: 63/64 decode steps show exact match (diff=0) between graph replay and eager forward - 50-run x 200-step benchmark: GC ON mean loss 0.001466 vs GC OFF 0.001475 (-0.65%, statistically indistinguishable; 26/50 runs higher) Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Author
Committer
Parents
Loading