RuntimeError: Inference tensors cannot be saved for backward (with CUDA kernels installed)

[leoauri]

I was able to start training with the naive implementation, but having built/installed the CUDA kernels, it blows up like this:

Training:   0%|          | 0/12330 [00:00<?, ?it/s]
Epoch 0:   0%|          | 0/12330 [00:00<?, ?it/s] 
Traceback (most recent call last):
  File "/home/rave-sashimi/.venv/bin/rave", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/rave-sashimi/RAVE/scripts/main_cli.py", line 30, in main
    app.run(train.main)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/absl/app.py", line 316, in run
    _run_main(main, args)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/absl/app.py", line 261, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/home/rave-sashimi/RAVE/scripts/train.py", line 279, in main
    trainer.fit(model, train, val, ckpt_path=run)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit
    call._call_and_handle_interrupt(
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1103, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1182, in _run_stage
    self._run_train()
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1205, in _run_train
    self.fit_loop.run()
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/fit_loop.py", line 267, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 213, in advance
    batch_output = self.batch_loop.run(kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 90, in advance
    outputs = self.manual_loop.run(kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/optimization/manual_loop.py", line 110, in advance
    training_step_output = self.trainer._call_strategy_hook("training_step", *kwargs.values())
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1485, in _call_strategy_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/strategies/strategy.py", line 378, in training_step
    return self.model.training_step(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/RAVE/rave/model.py", line 306, in training_step
    y = self.decoder(z)
        ^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/RAVE/rave/sashimi_decoder.py", line 138, in forward
    x, _ = layer(x)
           ^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/s4/s4/src/models/sequence/backbones/block.py", line 105, in forward
    y_for, new_state = self.layer(y, state=state, **kwargs)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/s4/s4/src/models/sequence/modules/s4block.py", line 168, in forward
    y, state = self.layer(x, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/s4/s4/src/models/sequence/kernels/fftconv.py", line 93, in forward
    k, k_state =  self.kernel(L=l_kernel, rate=rate, state=state) # (C H L) (B C H L)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/s4/s4/src/models/sequence/kernels/ssm.py", line 846, in forward
    r = cauchy_mult(v, z, A)
        ^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/s4/s4/extensions/kernels/cauchy.py", line 76, in cauchy_mult
    y = _cauchy_mult(v.view(-1, N), z, w.view(-1, N))
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/s4/s4/extensions/kernels/cauchy.py", line 59, in _cauchy_mult
    return CauchyMultiplySymmetric.apply(v, z, w)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/autograd/function.py", line 576, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal tensor and use it in autograd.

Epoch 0:   0%|          | 0/12330 [00:19<?, ?it/s]

Any ideas?

[leoauri]

It would appear to be due to extra checks introduced in pytorch 2.x.

Since realising this, I have spent a couple of days in dependency hell, trying to get working versions of torch 1.13.1 in line with other stuff I am using, and I have hit a wall.

Therefore, I'm going to try: updating this repository for torch 2.x.

Any advice?

[leoauri]

So far all I had to do was clone a few tensors, everything seems to be running smoothly. See https://github.com/leoauri/s4/tree/up-torch2.x

[leoauri]

The custom CUDA kernels seem to be unstable in combination with multiprocessing, my dataloader threads keep dying:

Epoch 30:   0%|          | 0/10 [00:00<?, ?it/s, v_num=5]ERROR: Unexpected segmentation fault encountered in worker.
�
Epoch 30:  10%|█         | 1/10 [00:10<01:34, 10.46s/it, v_num=5]
Epoch 30:  10%|█         | 1/10 [00:10<01:34, 10.46s/it, v_num=5]Traceback (most recent call last):
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1103, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1182, in _run_stage
    self._run_train()
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1205, in _run_train
    self.fit_loop.run()
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/fit_loop.py", line 267, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 187, in advance
    batch = next(data_fetcher)
            ^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/utilities/fetching.py", line 184, in __next__
    return self.fetching_function()
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/utilities/fetching.py", line 265, in fetching_function
    self._fetch_next_batch(self.dataloader_iter)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/utilities/fetching.py", line 280, in _fetch_next_batch
    batch = next(iterator)
            ^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/supporters.py", line 569, in __next__
    return self.request_next_batch(self.loader_iters)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/supporters.py", line 581, in request_next_batch
    return apply_to_collection(loader_iters, Iterator, next)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_utilities/core/apply_func.py", line 64, in apply_to_collection
    return function(data, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 734, in __next__
    data = self._next_data()
           ^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1492, in _next_data
    idx, data = self._get_data()
                ^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1454, in _get_data
    success, data = self._try_get_data()
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1285, in _try_get_data
    data = self._data_queue.get(timeout=timeout)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/queues.py", line 122, in get
    return _ForkingPickler.loads(res)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/multiprocessing/reductions.py", line 541, in rebuild_storage_fd
    fd = df.detach()
         ^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/resource_sharer.py", line 57, in detach
    with _resource_sharer.get_connection(self._id) as conn:
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/resource_sharer.py", line 86, in get_connection
    c = Client(address, authkey=process.current_process().authkey)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/connection.py", line 525, in Client
    answer_challenge(c, authkey)
  File "/usr/lib/python3.12/multiprocessing/connection.py", line 953, in answer_challenge
    message = connection.recv_bytes(256)         # reject large message
              ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/connection.py", line 216, in recv_bytes
    buf = self._recv_bytes(maxlength)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/connection.py", line 430, in _recv_bytes
    buf = self._recv(4)
          ^^^^^^^^^^^^^
  File "/usr/lib/python3.12/multiprocessing/connection.py", line 395, in _recv
    chunk = read(handle, remaining)
            ^^^^^^^^^^^^^^^^^^^^^^^
ConnectionResetError: [Errno 104] Connection reset by peer

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/rave-sashimi/.venv/bin/rave", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/rave-sashimi/RAVE/scripts/main_cli.py", line 37, in main
    app.run(train.main)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/absl/app.py", line 316, in run
    _run_main(main, args)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/absl/app.py", line 261, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/home/rave-sashimi/RAVE/scripts/train.py", line 304, in main
    trainer.fit(model, train, val, ckpt_path=run)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit
    call._call_and_handle_interrupt(
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 63, in _call_and_handle_interrupt
    trainer._teardown()
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in _teardown
    self.strategy.teardown()
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/pytorch_lightning/strategies/strategy.py", line 492, in teardown
    _optimizers_to_device(self.optimizers, torch.device("cpu"))
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_fabric/utilities/optimizer.py", line 28, in _optimizers_to_device
    _optimizer_to_device(opt, device)
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_fabric/utilities/optimizer.py", line 34, in _optimizer_to_device
    optimizer.state[p] = apply_to_collection(v, Tensor, move_data_to_device, device)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_utilities/core/apply_func.py", line 70, in apply_to_collection
    return {k: function(v, *args, **kwargs) for k, v in data.items()}
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_fabric/utilities/apply_func.py", line 101, in move_data_to_device
    return apply_to_collection(batch, dtype=_TransferableDataType, function=batch_to)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_utilities/core/apply_func.py", line 64, in apply_to_collection
    return function(data, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/lightning_fabric/utilities/apply_func.py", line 95, in batch_to
    data_output = data.to(device, **kwargs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rave-sashimi/.venv/lib/python3.12/site-packages/torch/utils/data/_utils/signal_handling.py", line 73, in handler
    _error_if_any_worker_fails()
RuntimeError: DataLoader worker (pid 85777) is killed by signal: Segmentation fault. 

Apparently this backend

 #include <THC/THCAtomics.cuh>

is deprecated? In any case training seems to be working and just as fast with the pykeops kernels, so proceeding with that.