transformers - 💡(How to fix) Fix [Bug] GlmMoeDsa crashes on second forward pass — stale indexer cache [2 comments, 2 participants]

  File ".../transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py", line 414, in forward
    index_mask.scatter_(-1, topk_indices, 0.0)  # [B, S, T]
torch.AcceleratorError: CUDA error: device-side assert triggered

k_cached = torch.cat([self._cached_keys, k], dim=1)  # [B, T, D]

# In modeling_glm_moe_dsa.py
class GlmMoeDsaModel(...):
    def forward(self, ...):
        # Reset the DSA indexer's cache
        self.dsa_indexer._cached_keys = None
        self.dsa_indexer._cached_indices = None
        
        # Rest of the forward method remains the same
        ...

# In dsa_indexer.py
class DSAIndexer(...):
    def __call__(self, ...):
        # Reset the cache
        self._cached_keys = None
        self._cached_indices = None
        
        # Rest of the __call__ method remains the same
        ...

model = AutoModelForCausalLM.from_pretrained("yujiepan/glm-5-tiny-random", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("yujiepan/glm-5-tiny-random")

inputs = tokenizer("Hello", return_tensors="pt").to(model.device)

# First forward: OK
out1 = model(**inputs)
print(out1.logits.shape)  # torch.Size([1, 1, 154880])

# Second forward: Should not crash
out2 = model(**inputs)
print(out2.logits.shape)  # torch.Size([1, 1, 154880])