[Quantization + FSDP] Support quantize_() for DTensor

[gau-nernst]

While trying out INT8 mixed precision pretraining (#748) with torchtitan, I came across an issue that if the model is FSDP-sharded, quantize_() won't work. The fix would be adding an extra logic to handle DTensor, similar to what FP8 is doing

ao/torchao/float8/float8_tensor.py

Lines 161 to 183 in f5703b0

	if isinstance(bits_fp8, DTensor):
	assert isinstance(
	scale, DTensor
	), "Expected Float8 scale to be a DTensor if bits_fp8 is a DTensor"
	bits_mesh = bits_fp8.device_mesh
	bits_placements = bits_fp8.placements
	local_bits = bits_fp8.to_local()
	local_scale = scale.to_local()
	inner_float8_tensor = Float8Tensor(
	local_bits,
	local_scale,
	tensor.dtype,
	linear_mm_config=linear_mm_config,
	gemm_input_role=gemm_input_role,
	)
	return DTensor.from_local(
	inner_float8_tensor,
	bits_mesh,
	bits_placements,
	run_check=False,
	shape=bits_fp8.size(),
	stride=bits_fp8.stride(),
	)

[msaroufim]

Yeah this came up in some discussions with inference providers like SGLang as well

[supriyar]

@jerryzh168 @kwen2501 is this addressed now with quantize + distributed inference composability work?

[jerryzh168]

this is not addressed yet, this is training use case I think, that we can explore in 2025 H1 together with @vkuzo, we do need a guide on how DTensor composes with quantization in both inference and training use cases