[QUESTION] How does veScale FSDP handles 2D FP8 blockwise scaling for RaggedShard

[zhongbozhu]

For deepseek FP8 blockwise recipe, the scaling block size for weights is 128x128. Having 2D block scaling means sharding by blocks will make memory incontiguous. Does it mean that we have to go through the following function to "permute" the original [N, K] weight tensor. And then after the parameter AG, do we also need to "unpermute" it such that we can launch a regular FP8 GEMM?

def to_blocked_128x128(a: torch.Tensor) -> torch.Tensor:
    N, K = a.shape
    return a.view(N//128, 128, K//128, 128).permute(0, 2, 1, 3).contiguous()

[asd123www]

Hi,

Thank you for the question!

No, you don't need to permute/unpermute the weight in this way.

In RaggedShard, we support customized sharding granularity in contiguous memory, and RaggedShard itself does not introduce extra memory movement. Take DeepSeek-V3 as an example, to support 128*128 quantization block, we set the sharding granularity to be 128 rows, combining several quantization blocks. This keeps each device’s local shard contiguous in the global layout, and each 128×128 quantization block remains local to a single device.

And the fsdp buffer padding ratio is shown in Figure 11(a).