Proper way to combine Tensor subclass with FSDP

username: Irvingwangjr
cooked: <p>Currently FSDP replace param.data with a empty tensor to free memory.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/f/8f812f5fd32c1f486fb74a8818cae220cf7487a5.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/8f812f5fd32c1f486fb74a8818cae220cf7487a5" title="img_v3_02pn_d70c5b85-abad-4db2-911d-e2c2a6018ddg"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f812f5fd32c1f486fb74a8818cae220cf7487a5_2_690x398.jpeg" alt="img_v3_02pn_d70c5b85-abad-4db2-911d-e2c2a6018ddg" data-base62-sha1="ktv3tbBUnCwT2yXgFzZbEv1p1YN" width="690" height="398" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f812f5fd32c1f486fb74a8818cae220cf7487a5_2_690x398.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f812f5fd32c1f486fb74a8818cae220cf7487a5_2_1035x597.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f812f5fd32c1f486fb74a8818cae220cf7487a5_2_1380x796.jpeg 2x" data-dominant-color="292928"><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">img_v3_02pn_d70c5b85-abad-4db2-911d-e2c2a6018ddg</span><span class="informations">1920×1110 141 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>but when the param is a tensor sub-class, like this</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/81c733df329a52c6233963db672d58031a3bf58e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/81c733df329a52c6233963db672d58031a3bf58e" title="img_v3_02pn_de93d041-51a0-476c-8154-3a08b5ea8d0g"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81c733df329a52c6233963db672d58031a3bf58e_2_690x472.jpeg" alt="img_v3_02pn_de93d041-51a0-476c-8154-3a08b5ea8d0g" data-base62-sha1="iw4nLzg6m0nFbI0sHXZOD55LKwu" width="690" height="472" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81c733df329a52c6233963db672d58031a3bf58e_2_690x472.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81c733df329a52c6233963db672d58031a3bf58e_2_1035x708.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81c733df329a52c6233963db672d58031a3bf58e_2_1380x944.jpeg 2x" data-dominant-color="232424"><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">img_v3_02pn_de93d041-51a0-476c-8154-3a08b5ea8d0g</span><span class="informations">1398×958 105 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>this way won’t work as following</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b57edd4ad62ab48f663df526f7a732c2e1e704ce.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b57edd4ad62ab48f663df526f7a732c2e1e704ce" title="img_v3_02pn_3e7cc104-1b0c-4c27-b48c-51291074245g"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b57edd4ad62ab48f663df526f7a732c2e1e704ce_2_663x500.jpeg" alt="img_v3_02pn_3e7cc104-1b0c-4c27-b48c-51291074245g" data-base62-sha1="pTAe88KS4XDbUDCgUlZWx6s8H0y" width="663" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b57edd4ad62ab48f663df526f7a732c2e1e704ce_2_663x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b57edd4ad62ab48f663df526f7a732c2e1e704ce_2_994x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b57edd4ad62ab48f663df526f7a732c2e1e704ce_2_1326x1000.jpeg 2x" data-dominant-color="262525"><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">img_v3_02pn_3e7cc104-1b0c-4c27-b48c-51291074245g</span><span class="informations">1552×1170 175 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>what’s the best practice to integrate FSDP and tensor subclass? I saw d-tensor’s way by using FSDPExtension and unwrap the params, but that interface seems customized only for d-tensor</p>

username: ptrblck
cooked: <p>Based on <a href="https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html">this tutorial</a> for FSDP2 it seems one of the advantages between FSDP1 and FSDP2 is:</p>
<blockquote>
<p>Offering a tensor subclass extension point to customize the all-gather, e.g. for float8 all-gather for float8 linears (<a href="https://dev-discuss.pytorch.org/t/enabling-float8-all-gather-in-fsdp2/2359">doc</a>), and NF4 for QLoRA (<a href="https://github.com/pytorch/torchtune/blob/main/README.md">doc</a>)</p>
</blockquote>
<p>so maybe this extension point might be useful for your use case as well.</p>

username: Irvingwangjr
cooked: <p>Thx for the advice!</p>
<p>But I’m not able to use FSDP2 right now, so I switch another way using torch_dispatch to do so, but FSDP2 should be a better way.</p>