[Bug]: GSA and RWKV6 Occasionally Report Gradient=NAN when Backward

[WorldEditors]

Describe the bug

Running training for GSA and RWKV will result in NAN gradient occasionally, rare at the beginning stage, but getting more frequent as the training processes. I checked parameters and losses, all of which are reasonable and shows no sign of explosion. The NAN comes suddenly. By switching model back to Transformer, this never happens

by using with torch.detect_anomaly(), get the following log:

  File "/home/xxx/Codes/flash-linear-attention/fla/layers/gsa.py", line 181, in forward
    o, recurrent_state = chunk_gsa(q, k, v, s, f,
  File "/home/xxx/Codes/flash-linear-attention/fla/ops/gsa/chunk.py", line 1203, in chunk_gsa
    ov, *final_state = ChunkGSAFunction.apply(q, k, v, s, g, scale, hk0, hv0, output_final_state, checkpoint_level)
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/autograd/function.py", line 574, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
 (Triggered internally at ../torch/csrc/autograd/python_anomaly_mode.cpp:111.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
W1107 18:46:16.781683 139687942383424 torch/multiprocessing/spawn.py:146] Terminating process 2106869 via signal SIGTERM
W1107 18:46:16.783345 139687942383424 torch/multiprocessing/spawn.py:146] Terminating process 2106870 via signal SIGTERM
W1107 18:46:16.783725 139687942383424 torch/multiprocessing/spawn.py:146] Terminating process 2106871 via signal SIGTERM
Traceback (most recent call last):
  File "train.py", line 10, in <module>
    runner.start(AnyMDPRSA, AnyMDPEpoch, AnyMDPEpoch)
  File "/home/xxx/Codes/L3C_Baselines/l3c_baselines/utils/trainer.py", line 335, in start
    mp.spawn(dist_process,
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 282, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 238, in start_processes
    while not context.join():
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 189, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 76, in _wrap
    fn(i, *args)
  File "/home/xxx/Codes/L3C_Baselines/l3c_baselines/utils/trainer.py", line 288, in dist_process
    for need_evaluate in train_object.run(epoch, device, device_type):
  File "/home/xxx/Codes/L3C_Baselines/l3c_baselines/utils/trainer.py", line 151, in run
    self.computer.compute(
  File "/home/xxx/Codes/L3C_Baselines/projects/AnyMDP/anymdp_epoch.py", line 80, in compute
    syn_loss.backward()
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/_tensor.py", line 521, in backward
    torch.autograd.backward(
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/autograd/__init__.py", line 289, in backward
    _engine_run_backward(
  File "/home/xxx/.local/lib/python3.8/site-packages/torch/autograd/graph.py", line 769, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Function 'ChunkGSAFunctionBackward' returned nan values in its 3th output.

Steps to reproduce the bug

Can not provide a code sample, it did not happen.in specific model in specific steps.

Expected behavior

N/A

Environment info

1. torch: 2.4.1
2. triton: 3.0.0

CUDA: 12.6, NVIDIA A800 80GB PCIe

[yzhangcs]

Thank you for reporting this issue. Could you elaborate more on the input shapes so that I can do some simulation exps

[WorldEditors]

some debugs showing the following part in chunk.py causes the problem

1016     grid = (NV, NT * NC, B * HQ)
1017     chunk_gsa_bwd_kernel_intra_KV[grid](
1018         v, g, o, A, do, dv, dg,
1019         v.stride(1), v.stride(2),
1020         T=T, V=V, BT=BT, BC=BC, BV=BV, NC=NC, NG=NG,
1021         OVERWRITE_DG=overwrite_dg,
1022         num_warps=num_warps,
1023         num_stages=num_stages
1024     )   # After this function dv and dg suddenly goes to NAN
1025     return dq, dk, dv, dg, dh0

hyper parameters: hidden_size = 512, num_slots=64, nheads=4

Others are just default (GatedSlotAttention function default)

[yzhangcs]

@WorldEditors How about the sequence length

[WorldEditors]

tried sequence length with 3K, 12K, 24K, all possible to reproduce the problem

[WorldEditors]

In GSA, I tried restore training from a problematic checkpoint, and there are some interesting discoveries:

The training will be OK only if I reinitialize the f_proj parameters. While reinitialize the other parameters (q_proj, k_proj, v_proj, g_norm, o_proj) won't solve the problem.

So I guess there is something wrong with the f_proj parameter?

But why there is similar problems in RWKV6? I have no idea

[sustcsonglin]

Hi, Thanks for reporting it! Do you have input tensor and model weight such that we can reproduce it?

[WorldEditors]

Hi, Thanks for reporting it! Do you have input tensor and model weight such that we can reproduce it?

I'm afraid it is quite challenging to extract a minimal reproducible dataset, code, and model. But it seems to be very possible that this has something to do with the forgetting gate (f and g). Hopefully this will give some clues to you

[WorldEditors]

Update:

I tried bounding the forget gate value as

https://github.com/sustcsonglin/flash-linear-attention/pull/78/files

Running another 400 iterations does not yield any abnormalty

But this has the risk of impact exsiting models.

[yzhangcs]

@WorldEditors Hello, sry for my super late reply. Just refactored GSA layers/kernels to address potential indexing and gradient errors, could you try it again? Looking forward to your feedbacks.