FSDP2 integration: torch.chunks(Params4bit)

[ved1beta]

Fix Params4bit tensor subclass handling

• Replace torch._C._dispatch_without_tensor_functionalization() with torch._C._DisableTorchDispatch()
• Add custom repr method to prevent .tolist() errors
• Implement chunk/cat operations that preserve Params4bit tensor type
• Ensure proper tensor subclass behavior in PyTorch 2.6.0
  Fix FSDP2 integration: torch.chunks(Params4bit) not returning Params4bit subclass  #1424

[ved1beta]

can i get a feed back on this :)

[Titus-von-Koeller]

Hey @ved1beta,

Thanks for the contribution. Could you tell us a bit more on why you deleted the lines that you deleted and how you validated your implementation?

[ved1beta]

When changing from torch.nn.Parameter to torch.Tensor, we're now using the parent class's initialization methods which already handle basic tensor properties.
for validation creating a test initialize 4-bit quantized model wraps it with FSDP2.

[Titus-von-Koeller]

Really cool, ultimately a high-level test making sure we also notice when this functionality breaks would be really cool. For now to reproduce your exact results it would already be helpful to include an example script in /examples.

[Titus-von-Koeller]

@ved1beta Just wondering why you closed this?

[ved1beta]

I was stuck, examples and tests were failing can you help me with it i amwilling to make it work

[djsaunde]

I'd like to have this for Axolotl, I might jump in here and help test @ved1beta