What's recommended way to intergrate FSDP with Customize Tensor Unit

Irvingwangjr: <p>I’m trying to wrap a model which contains special tensors with FSDP, but some error happens. I’ll first show the specific error in case someone else might also run into this.<br>
Refence Repo:</p>
<ul>
<li><a href="https://github.com/microsoft/TRELLIS/blob/eeacb0bf6a7d25058232d746bef4e5e880b130ff/trellis/modules/sparse/linear.py#L10:" class="inline-onebox" rel="noopener nofollow ugc">TRELLIS/trellis/modules/sparse/linear.py at eeacb0bf6a7d25058232d746bef4e5e880b130ff · microsoft/TRELLIS · GitHub</a> use this SparseTensor as the input and output.</li>
<li>sparse backends: <a href="https://github.com/mit-han-lab/torchsparse" class="inline-onebox" rel="noopener nofollow ugc">GitHub - mit-han-lab/torchsparse: [MICRO'23, MLSys'22] TorchSparse: Efficient Training and Inference Framework for Sparse Convolution on GPUs.</a>.</li>
</ul>
<p>firstly I wrap this SparseLinear as a FSDP Unit, using FULLSHARD, this error will happen:</p>
<pre><code class="lang-auto">RuntimeError: setStorage: sizes [1024, 4096], strides [4096, 1], storage offset 12597248, and itemsize 2 requiring a storage sizeof 33583104 are out of bounds for storage of size 0
</code></pre>
<p>Then I switch FULLSHARD to SHARD_GRAD_OP, another error show up:</p>
<pre><code class="lang-auto">ValueError: expected to be in states [&lt;TrainingState.FORWARD_BACKWARD: 2&gt;] but current state is TrainingSttate.IDLE
</code></pre>
<p>eventually, it turns out that, when register hook to the FSDP Unit’s output tensors, if it’s not a torch.Tensor, the pre-backward hook will miss.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/3/e3e984f7517f52631d72d097b7be74ee3481740f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e3e984f7517f52631d72d097b7be74ee3481740f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/3/e3e984f7517f52631d72d097b7be74ee3481740f_2_690x480.jpeg" alt="image" data-base62-sha1="wwcHkaKo9Gba9gdQry9iOQPXm5F" width="690" height="480" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/3/e3e984f7517f52631d72d097b7be74ee3481740f_2_690x480.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/3/e3e984f7517f52631d72d097b7be74ee3481740f_2_1035x720.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/3/e3e984f7517f52631d72d097b7be74ee3481740f.jpeg 2x" data-dominant-color="242524"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1280×892 155 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>now I just make sure that the SparseLinear will output a normal tensor and I do the conversions else where. But I’m seeking for better advice here!</p>

Irvingwangjr: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/0/10643a12130ca30a11c5ca182f6fe93f5ad2160e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/10643a12130ca30a11c5ca182f6fe93f5ad2160e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/10643a12130ca30a11c5ca182f6fe93f5ad2160e_2_690x116.jpeg" alt="image" data-base62-sha1="2l0ndIRAgwxlsunK3UDNVviLzJk" width="690" height="116" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/10643a12130ca30a11c5ca182f6fe93f5ad2160e_2_690x116.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/10643a12130ca30a11c5ca182f6fe93f5ad2160e_2_1035x174.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/10643a12130ca30a11c5ca182f6fe93f5ad2160e_2_1380x232.jpeg 2x" data-dominant-color="10100F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2886×488 201 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>