RWKV CUDA error: an illegal memory access was encountered during training from scratch
Author: npkanaka
Pull request: no

I’m encountering a runtime error while training a RWKV model (RwkvForCausalLM) with the Hugging Face Trainer. The error occurs during the backward pass and appears to be a CUDA illegal memory access. Below is the error traceback and a summary of my steps to reproduce.

### Error Traceback
```
Traceback (most recent call last):
  File "/users/PAS2581/kanaka/research/GrokkedTransformersarewang2024/trying_different_archs/mamba/main_v11_fsl.py", line 954, in <module>
    main()
  File "/users/PAS2581/kanaka/research/GrokkedTransformersarewang2024/trying_different_archs/mamba/main_v11_fsl.py", line 928, in main
    trainer.train()
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/transformers/trainer.py", line 2176, in train
    return inner_training_loop(
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/transformers/trainer.py", line 2536, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/transformers/trainer.py", line 3717, in training_step
    self.accelerator.backward(loss, **kwargs)
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/accelerate/accelerator.py", line 2244, in backward
    self.scaler.scale(loss).backward(**kwargs)
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/torch/_tensor.py", line 581, in backward
    torch.autograd.backward(
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/torch/autograd/function.py", line 307, in apply
    return user_fn(self, *args)
  File "/users/PAS2581/kanaka/miniconda3/envs/grokk/lib/python3.11/site-packages/transformers/models/rwkv/modeling_rwkv.py", line 181, in backward
    g_time_decay.to(input_dtype),
RuntimeError: CUDA error: an illegal memory access was encountered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

### Model
```
if(args.usemodel=="rwkv"):
    tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")
    tokenizer.pad_token = tokenizer.eos_token
    print(f"tokenizer.bos_token: {tokenizer.bos_token}, {tokenizer.bos_token_id}")
    with open(os.path.join(args.data_dir, "vocab.json"), "r") as vocab_file:
        new_tokens = json.load(vocab_file)
    tokenizer.add_tokens(new_tokens)
    print("Tokenizer size after adding vocab:", len(tokenizer.vocab))

    model = RwkvForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile", use_cache=False)
    model.resize_token_embeddings(len(tokenizer))
```

### Training loop

```
training_args = TrainingArguments(
    output_dir=args.output_dir,
    logging_dir='./logs',
    gradient_accumulation_steps=1,  
    save_steps=10000, 
    max_steps=1000000, 
    eval_strategy="steps",
    eval_steps=50000,
    logging_strategy="epoch",
    logging_steps=2000,
    save_total_limit=2, 
    fp16=True,
    dataloader_num_workers=4,
    per_device_train_batch_size=512,
    per_device_eval_batch_size=512,
    learning_rate=1e-4,
    lr_scheduler_type="constant_with_warmup",
    warmup_steps=2000,
    weight_decay=0.1,
    max_grad_norm=1.0,
    save_safetensors=False,
)
```

### What I’ve Tried

- Disabling mixed precision (fp16=False) did not resolve the issue.
- Checking token ranges: all token IDs are within vocab_size=51383.
- Removing attention_mask from the forward pass didn’t fix it (the error still occurred).
- Upgrading drivers / environment: still present.
- The code runs on an A100 GPU with ~80GB memory available, so it doesn’t appear to be a straightforward OOM.

Comments:
cc @SunMarc @muellerzr 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I wanted to follow up on this issue since it has been marked as stale, but I am still experiencing the CUDA illegal memory access error.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

same issue