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.