Problem with RWKV training when autocast and GradScaler both enabled
Author: nanjiangwill
Pull request: no

### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.15.0-79-generic-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.2
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

run the following script

``` 
import os
import platform
import torch
import torch.nn as nn
import torch.optim as optim
import transformers
import time

input_size = 32
seq_len = 256
hidden_size = input_size
num_layers = 2
num_heads = 2
batch_size = 32
num_epochs = 1000

# set seed
torch.manual_seed(0)


def count_paramters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


rwkv_config = transformers.RwkvConfig(
    context_length=seq_len,
    hidden_size=hidden_size,
    num_hidden_layers=num_layers,
    intermediate_size=hidden_size * 4,
)


def train(model, device="cuda"):
    ptdtype = {
        "float32": torch.float32,
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
    }["float32"]
    ctx = torch.amp.autocast(device_type="cuda", dtype=ptdtype, enabled=True)

    scaler = torch.cuda.amp.GradScaler(enabled=True)

    head = nn.Linear(hidden_size, 1)

    model = model.to(device)
    head = head.to(device)
    model.train()
    head.train()

    # Define loss function and optimizer
    criterion = nn.MSELoss()  # Mean Squared Error loss
    optimizer = optim.SGD(list(model.parameters()) + list(head.parameters()), lr=0.01)

    # Training loop
    start = time.time()
    for epoch in range(num_epochs):
        inputs = torch.rand(batch_size, seq_len, input_size, device=device)
        labels = torch.rand(batch_size, seq_len, 1, device=device)

        # Zero the gradients
        optimizer.zero_grad()

        # Forward pass
        with ctx:
            outputs = model(inputs_embeds=inputs).last_hidden_state
            outputs = head(outputs)

            # Compute the loss
            loss = criterion(outputs, labels)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        scaler.step(optimizer)
        scaler.update()

        # Print loss every 100 epochs
        if (epoch + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
    end = time.time()
    return end - start


for device in ["cuda"]:
    print(f"Training RWKV on {device}")
    rwkv_model = transformers.RwkvModel(rwkv_config)
    print(f"RWKV parameters: {count_paramters(rwkv_model)}\n")
    rwkv_train_time = train(rwkv_model, device=device)
    print(f"RWKV training time: {rwkv_train_time:.2f} seconds\n")

# print cuda version
print(f"cuda version: {torch.version.cuda}")
# print torch version
print(f"torch version: {torch.__version__}")
# print what cuda driver is being used
print(f"cuda driver: {torch.backends.cudnn.version()}")
# print huggingface version
print(f"huggingface version: {transformers.__version__}")
# print system information like python version, operating system, etc.
print(f"operating system: {os.name}")
print(platform.system(), platform.release())
```

### Expected behavior

expect no NaN value
current output is the following
```
Training RWKV on cuda
RWKV parameters: 1636320

Epoch [100/1000], Loss: nan
Epoch [200/1000], Loss: nan
Epoch [300/1000], Loss: nan
Epoch [400/1000], Loss: nan
Epoch [500/1000], Loss: nan
Epoch [600/1000], Loss: nan
Epoch [700/1000], Loss: nan
Epoch [800/1000], Loss: nan
Epoch [900/1000], Loss: nan
Epoch [1000/1000], Loss: nan
RWKV training time: 8.87 seconds

cuda version: 11.7
torch version: 2.0.1+cu117
cuda driver: 8500
huggingface version: 4.31.0
operating system: posix
Linux 5.15.0-79-generic
```

Comments:
Also, when I use Adam/AdamW/RMSProp, loss also goes to NaN
I find out during first training step, it is fine, but second forward pass will have some layers with gradient like `1e-15` which is a underflow 

when GradScaler is disabled, the above works

Hey @nanjiangwill, I am not really sure we can say the issue is in `transformers` as you are using both autocast and GradScaler, and you custom training loop 😅. We usually recommend to ask these questions on [the forum](https://discuss.huggingface.co/) ! 

Hi! same problem, and training loop is extremely simple:
```py
train_dataloader = torch.utils.data.DataLoader(dataset["train"],batch_size=8)
optimizer = torch.optim.AdamW(rwkv.parameters(), lr=0.001)
from tqdm.notebook import tqdm
epochs_tqdm = tqdm(range(12))
for epoch in epochs_tqdm:
    for batch in tqdm(train_dataloader):
        batch = {k: v.to("cuda") for k, v in batch.items()}
        outputs = rwkv(**batch)
        loss = outputs.loss
        epochs_tqdm.set_description(f"{loss}")
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    generate()
```
I train the model from scratch, I also get NaN loss and logits
```py
RwkvCausalLMOutput(loss=tensor(nan, device='cuda:0', grad_fn=<NllLossBackward0>), logits=tensor([[[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         ...,
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]]], device='cuda:0',
       grad_fn=<UnsafeViewBackward0>), state=None, hidden_states=None, attentions=None)
```

cc @younesbelkada I think we saw something similar and the model is just unstable for certain sizes no? 

Hi @nanjiangwill @drimeF0 can you share with us how do you load your models? might be a precision issue 

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.

EDIT: I will create a new issue to track this.

I realise this was auto-closed 2 years ago, but I am experiencing the same issue when trying to train a small RWKV. After a deep dive I have identified that under a specific set of conditions (happens on my remote distributed training, but not on local single GPU training) the value `GradScaler._scale` is changed from 65536.0 to 0.035331886261701584 during the backward pass of the first iteration. This means that: 1. during the backward pass half the network has  the incorrect scaling applied; and 2. during rescaling half the network is rescaled incorrectly. 

The logs below are created by registering a hook on the backward pass. Somehow the `GradScaler._scale` changes between backward for `module.temporal.rwkv.blocks.0.attention` and `module.temporal.rwkv.blocks.0.attention.receptance`.
  
I assume this is not the expected behaviour. How is the backward pass able to modify the GradScalers scale value?
```
Backward pass for module.temporal.rwkv.blocks.0.attention.output:
	
2025-12-09T16:02:25.855Z
	
Grad Input
	
2025-12-09T16:02:25.856Z
	
tensor([[-1.2789e-03, 6.6137e-04, 7.4692e-03, ..., -9.2506e-04, 5.3635e-03, -3.7537e-03], [ 1.1024e-03, -3.2349e-03, 1.1444e-02, ..., 3.9043e-03, 3.7823e-03, -3.5877e-03], [-6.8378e-04, -1.8013e-04, 9.0256e-03, ..., 6.9962e-03, 2.7275e-03, -2.7103e-03], ..., [-1.4688e+00, 1.9844e+00, 1.1230e-02, ..., -1.8726e-01, -3.0334e-02, -2.3694e-01], [-3.1787e-01, 6.3135e-01, -7.5879e-01, ..., -1.2520e+00, -1.5264e+00, -1.9824e+00], [ 0.0000e+00, 0.0000e+00, 0.0000e+00, ..., 0.0000e+00, 0.0000e+00, 0.0000e+00]], device='cuda:0', dtype=torch.float16)
	
2025-12-09T16:02:25.856Z
	
Gradient w.r.t. input: torch.Size([6400, 384]) False
	
2025-12-09T16:02:25.856Z
	
input_mean=0.000e+00, input_std=1.483e-01, input_max=4.270e+00, input_min=-3.824e+00
	
2025-12-09T16:02:25.856Z
	
Scaler scale: 65536.0
	
2025-12-09T16:02:25.856Z
	
--------------------------------
	
2025-12-09T16:02:25.856Z
	
Backward pass for module.temporal.rwkv.blocks.0.attention:
	
2025-12-09T16:02:25.856Z
	
Grad Input
	
2025-12-09T16:02:25.856Z
	
tensor([[-1.2789e-03, 6.6137e-04, 7.4692e-03, ..., -9.2506e-04, 5.3635e-03, -3.7537e-03], [ 1.1024e-03, -3.2349e-03, 1.1444e-02, ..., 3.9043e-03, 3.7823e-03, -3.5877e-03], [-6.8378e-04, -1.8013e-04, 9.0256e-03, ..., 6.9962e-03, 2.7275e-03, -2.7103e-03], ..., [-1.4688e+00, 1.9844e+00, 1.1230e-02, ..., -1.8726e-01, -3.0334e-02, -2.3694e-01], [-3.1787e-01, 6.3135e-01, -7.5879e-01, ..., -1.2520e+00, -1.5264e+00, -1.9824e+00], [ 0.0000e+00, 0.0000e+00, 0.0000e+00, ..., 0.0000e+00, 0.0000e+00, 0.0000e+00]], device='cuda:0', dtype=torch.float16)
	
2025-12-09T16:02:25.856Z
	
Gradient w.r.t. input: torch.Size([6400, 384]) False
	
2025-12-09T16:02:25.856Z
	
input_mean=0.000e+00, input_std=1.483e-01, input_max=4.270e+00, input_min=-3.824e+00
	
2025-12-09T16:02:25.856Z
	
Scaler scale: 65536.0
	
2025-12-09T16:02:25.856Z
	
--------------------------------
	
2025-12-09T16:02:25.856Z
	
Backward pass for module.temporal.rwkv.blocks.0.attention.receptance:
	
2025-12-09T16:02:25.856Z
	
Grad Input
	
2025-12-09T16:02:25.856Z
	
tensor([[ 4.2748e-04, -4.5419e-05, -1.4007e-05, ..., 7.7152e-04, 7.2670e-04, 1.3065e-03], [ 1.9312e-05, 4.6492e-05, -5.3406e-05, ..., 8.4877e-04, 8.2636e-04, 1.4095e-03], [-1.5438e-04, 8.3327e-05, -1.8585e-04, ..., 4.8113e-04, 4.7350e-04, 1.1053e-03], ..., [-2.6169e-02, -6.6071e-03, 1.5312e-02, ..., -1.2573e-01, 1.2497e-02, 3.1204e-02], [ 3.0859e-01, -2.0275e-03, 7.3059e-02, ..., -1.1890e-01, 3.7537e-03, -2.0477e-02], [ 0.0000e+00, -0.0000e+00, -0.0000e+00, ..., -0.0000e+00, -0.0000e+00, 0.0000e+00]], device='cuda:0', dtype=torch.float16)
	
2025-12-09T16:02:25.856Z
	
Gradient w.r.t. input: torch.Size([6400, 384]) False
	
2025-12-09T16:02:25.856Z
	
input_mean=4.289e-04, input_std=3.041e-02, input_max=1.535e+00, input_min=-1.653e+00
	
2025-12-09T16:02:25.856Z
	
Scaler scale: 0.035331886261701584
	
2025-12-09T16:02:25.856Z
	
--------------------------------
	
2025-12-09T16:02:25.856Z
	
Backward pass for module.temporal.rwkv.blocks.0.attention.value:
	
2025-12-09T16:02:25.856Z
	
Grad Input
	
2025-12-09T16:02:25.856Z
	
tensor([[-0.0028, -0.0067, -0.0088, ..., -0.0054, -0.0072, 0.0046], [ 0.0008, -0.0067, -0.0135, ..., -0.0023, -0.0038, 0.0038], [ 0.0027, -0.0053, -0.0118, ..., -0.0022, -0.0010, 0.0029], ..., [ 0.0064, 0.0166, -0.0150, ..., 0.6323, -0.0299, 0.0739], [ 0.0187, 0.0070, -0.0186, ..., 0.4993, -0.0066, -0.0747], [ 0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000]], device='cuda:0', dtype=torch.float16)
	
2025-12-09T16:02:25.856Z
	
Gradient w.r.t. input: torch.Size([6400, 384]) False
	
2025-12-09T16:02:25.856Z
	
input_mean=2.594e-04, input_std=5.133e-02, input_max=2.002e+00, input_min=-2.266e+00
	
2025-12-09T16:02:25.856Z
	
Scaler scale: 0.035331886261701584
	
2025-12-09T16:02:25.856Z
	
--------------------------------
	
2025-12-09T16:02:25.856Z
	
Backward pass for module.temporal.rwkv.blocks.0.attention.key:
	
2025-12-09T16:02:25.856Z
	
Grad Input
	
2025-12-09T16:02:25.856Z
	
tensor([[-6.3934e-03, -1.3838e-03, -1.9569e-03, ..., -3.3545e-04, -1.2231e-04, -1.3459e-04], [-2.1420e-03, -9.5463e-04, -1.5831e-04, ..., 4.0007e-04, 7.0930e-05, 4.4179e-04], [-2.5177e-03, -4.7135e-04, 1.9703e-03, ..., 4.8220e-05, 4.5657e-05, -3.7384e-04], ..., [-3.0365e-03, -4.6844e-03, -1.9531e-03, ..., -5.7892e-02, -5.5838e-04, 1.3065e-03], [-4.6883e-03, 1.0128e-03, -4.8752e-03, ..., 2.2781e-02, 3.4666e-04, 1.1902e-03], [ 0.0000e+00, 0.0000e+00, 0.0000e+00, ..., 0.0000e+00, 0.0000e+00, 0.0000e+00]], device='cuda:0', dtype=torch.float16)
	
2025-12-09T16:02:25.856Z
	
Gradient w.r.t. input: torch.Size([6400, 384]) False
	
2025-12-09T16:02:25.856Z
	
input_mean=0.000e+00, input_std=5.814e-03, input_max=4.038e-01, input_min=-4.341e-01
	
2025-12-09T16:02:25.856Z
	
Scaler scale: 0.035331886261701584
	
```