How to train FSDP with unused parameters

VERSPD0: <p>My model has parameters that will not be used in some cases. This would cause some gradient problem while training.</p>
<p>In case of DDP, it can be solved by setting find_unused_parameters=True. But In FSDP, it seems no such a mechanism to do so.</p>
<p>The follows are my error messages:</p>
<pre data-code-wrap="shell"><code class="lang-plaintext">File "/opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 1108, in _post_backward_final_callback
    _catch_all_reshard(fsdp_state)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 1152, in _catch_all_reshard
    _p_assert(
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 145, in _p_assert
    traceback.print_stack()
    worker(args.local_rank, None, args, model_for_training, model_for_dataloader)
  File "trainer.py", line 1287, in worker
Traceback (most recent call last):
  File "pretrain.py", line 196, in &lt;module&gt;
    trainer.train(args, gpu_id, rank, train_loader, model_for_training, optimizer, scheduler)
  File "trainer.py", line 425, in train
    main()
  File "pretrain.py", line 192, in main
    main()
  File "pretrain.py", line 192, in main
    trainer.train(args, gpu_id, rank, train_loader, model_for_training, optimizer, scheduler)
  File "trainer.py", line 425, in train
    trainer.train(args, gpu_id, rank, train_loader, model_for_training, optimizer, scheduler)
  File "trainer.py", line 425, in train
    loss.backward()
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    trainer.train_and_validate(args)
  File "trainer.py", line 136, in train_and_validate
    loss.backward()
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    trainer.train_and_validate(args)
  File "trainer.py", line 136, in train_and_validate
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
SystemError: &lt;built-in method run_backward of torch._C._EngineBase object at 0x7ffbc34f7d30&gt; returned NULL without setting an exception
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
SystemError: &lt;built-in method run_backward of torch._C._EngineBase object at 0x7f39022dfd30&gt; returned NULL without setting an exception
    loss.backward()
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    worker(args.local_rank, None, args, model_for_training, model_for_dataloader)
  File "trainer.py", line 1287, in worker
    worker(args.local_rank, None, args, model_for_training, model_for_dataloader)
  File "trainer.py", line 1287, in worker
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
SystemError: &lt;built-in method run_backward of torch._C._EngineBase object at 0x7faca0037d30&gt; returned NULL without setting an exception
    trainer.train(args, gpu_id, rank, train_loader, model_for_training, optimizer, scheduler)
  File "trainer.py", line 425, in train
    trainer.train(args, gpu_id, rank, train_loader, model_for_training, optimizer, scheduler)
  File "trainer.py", line 425, in train
    loss.backward()
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
SystemError: &lt;built-in method run_backward of torch._C._EngineBase object at 0x7fe25d0c3d30&gt; returned NULL without setting an exception
    loss.backward()
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
SystemError: &lt;built-in method run_backward of torch._C._EngineBase object at 0x7fe62eaa3d30&gt; returned NULL without setting an exception
</code></pre>

S_Chopra: <p><a class="mention" href="/u/verspd0">@VERSPD0</a> Were you able to find a solution to this problem?</p>