FSDP2 integration: torch.chunks(Params4bit) not returning Params4bit subclass

[mreso]

System Info

Hi, I'am trying to make FSDP2 work with a llama model quantized with bitsandbytes but it seems that bitsandbytes' tensor subclasses like Params4bit are not compatible with the way FSDP2 shards the model.
When creating the DTensors to shard the model FSDP2 applies torch.chunk to the parameters which get returned by torch.chunk as ordinary Tensors instead of the original subclass (like Params4bit) which leads to errors down the line.

Is this a known issue and are there plans to make bitsandbytes composable with FSDP2?

Reproduction

Created a simple repro:

import torch
import torch.nn as nn

import bitsandbytes as bnb
from bitsandbytes.nn import Params4bit
        
blocksize=64
compress_statistics = True
quant_type = "fp4"
quant_storage=torch.uint8

w = torch.ones(4).to("cuda")

w_4bit, quant_state = bnb.functional.quantize_4bit(
    w,
    blocksize=blocksize,
    compress_statistics=compress_statistics,
    quant_type=quant_type,
    quant_storage=quant_storage,
    )

b = Params4bit.from_prequantized(w_4bit, quant_state.as_dict(packed=True))
print(f"{b=}")

chunks = torch.chunk(b, 2, dim=0)

print(f"{chunks=}")

Output:

b=Parameter containing:
Parameter(Params4bit([[51],
            [51]], device='cuda:0', dtype=torch.uint8))
chunks=(tensor([[51]], device='cuda:0', dtype=torch.uint8), tensor([[51]], device='cuda:0', dtype=torch.uint8))

Expected behavior

Expecting the output of torch.chunk to be a a tuple of Params4bits instead of a Tensors.

[Titus-von-Koeller]

We’re prioritizing this with high prio, but our current focus is on tasks with even greater impact. Contributions are highly appreciated, and we’re happy to assist with anything needed along the way.

Let’s collaborate to address this as soon as possible. Thanks for taking the initiative and highlighting its importance, along with the limitations you’ve encountered.