`current_max_attn_logits` 会导致 CUDA 图形回放失败,因为其 Python 级别的赋值没有重新执行
作者: lk137095576创建于 2026年9月18日更新于 2026年9月18日
标签bugcommunity-request
在 megatron/core/extensions/transformer_engine.py (大约 L2479) 中,使用以下模式来累积前向调用中的注意力 logits: self.current_max_attn_logits = torch.max( self.current_max_attn_logits, batch_max_attention_logits ) 这与在包含 forward() 的情况下捕获 CUDA 图形不兼容。赋值是 Python 级别的状态更新,在图形回放中不会重新执行,因此在捕获时记录的张量地址与预期的跨调用累加语义不一致或过时。
内容来源: NVIDIA/Megatron-LM