Crash training video when offload active

[mcDandy]

This is for bugs only

Did you already ask in the discord?

Yes

You verified that this is a bug and not a feature request or question by asking in the discord?

Yes

Describe the bug

Version: 0e4b6e8

I am trying to train a videos (added at step 1000). It fails when offloading is enabled and hangs GPU when offloading is disabled.

The videos are FHD 226 frames and FHD 107 frames

PLFUN4-minimax_H3_1:   3%|▎         | 1021/30000 [02:42<62:16:05,  7.74s/it, lr: 1.5e-04 loss: 4.073e+00]Error running job: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
========================================
Result:
 - 0 completed jobs
 - 1 failure
========================================
Traceback (most recent call last):
  File "D:\stability\Data\Packages\AI-Toolkit\run.py", line 147, in <module>
    main()
  File "D:\stability\Data\Packages\AI-Toolkit\run.py", line 132, in main
    raise e
  File "D:\stability\Data\Packages\AI-Toolkit\run.py", line 120, in main
    job.run()
  File "D:\stability\Data\Packages\AI-Toolkit\jobs\ExtensionJob.py", line 22, in run
    process.run()
  File "D:\stability\Data\Packages\AI-Toolkit\jobs\process\BaseSDTrainProcess.py", line 2670, in run
    loss_dict = self.hook_train_loop(batch_list)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\sd_trainer\SDTrainer.py", line 2257, in hook_train_loop
    loss = self.train_single_accumulation(batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\sd_trainer\SDTrainer.py", line 2146, in train_single_accumulation
    noise_pred = self.predict_noise(
                 ^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\sd_trainer\SDTrainer.py", line 1358, in predict_noise
    return self.sd.predict_noise(
           ^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\toolkit\models\base_model.py", line 947, in predict_noise
    noise_pred = self.get_noise_prediction(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\diffusion_models\minimax_h3\minimax_h3.py", line 1025, in get_noise_prediction
    video_pred, audio_pred = self.model(
                             ^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\nn\modules\module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\nn\modules\module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\diffusion_models\minimax_h3\src\transformer.py", line 489, in forward
    x = checkpoint(
        ^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 1044, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\utils\checkpoint.py", line 503, in checkpoint
    ret = function(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\nn\modules\module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\nn\modules\module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\diffusion_models\minimax_h3\src\transformer.py", line 304, in forward
    self.adaln_proj(temb)
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\nn\modules\module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\nn\modules\module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\stability\Data\Packages\AI-Toolkit\extensions_built_in\diffusion_models\minimax_h3\src\transformer.py", line 229, in forward
    weight = self.linear.weight.to(device=temb.device, dtype=torch.float32)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.