Layer offload + W8A8/Fp8 quantization: DoRA / weight-decomposed LoKr crash during initialization (device mismatch)

[dxqb]

Summary

Training a model quantized with W8A8 (int8 or fp8) or Fp8, with layer offloading enabled and a weight-decomposition adapter (DoRA, or LoKr with weight decomposition), raises a device-mismatch RuntimeError while the adapters are being initialized, before the first training step:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
  File "modules/util/quantization_util.py", line 76, in dequantize
    return q.float() * scale

Mechanism

GenericTrainer runs setup_train_device before setup_model. With layer_offload_fraction > 0, setup_train_device activates the LayerOffloadConductor, which parks a fraction of the quantized layers on the temp device. The offload path (get_offload_tensors / offload_quantized) relocates each layer's weight (and bias), while the scale buffer of LinearW8A8 / LinearFp8 is held on the train device. An offloaded quantized layer therefore ends up with weight on the temp device and scale on the train device.

setup_model then builds the adapters. DoRAModule.initialize_weights and the weight-decompose branch of LoKrModule.initialize_weights reconstruct each layer's base weight with get_unquantized_weight(orig_module, torch.float, train_device). For an offloaded layer this reaches LinearW8A8.unquantized_weight / LinearFp8.unquantized_weight, which dequantize weight against scale without moving them to the requested device. With weight on the temp device and scale on the train device, dequantize throws.

LinearNf4.unquantized_weight moves both the weight and its quant state to the requested device before dequantizing, which is why NF4 handles the same situation cleanly.

Affected

• Quantization: W8A8 int8, W8A8 fp8, Fp8.
• Gradient checkpointing with CPU offload and layer_offload_fraction > 0 (any fraction that offloads at least one layer).
• Adapters that reconstruct the base weight at initialization: DoRAModule, and LoKrModule with weight decomposition. Plain LoRA, LoHa, and OFT add a delta on top of orig_forward and are unaffected.

Fix

1. LinearW8A8.unquantized_weight and LinearFp8.unquantized_weight should honor the device argument by moving weight and scale to device before dequantizing, matching LinearNf4.unquantized_weight.
2. DoRAModule.initialize_weights and LoKrModule.initialize_weights place the decomposition scale (dora_scale) on self.orig_module.weight.device; for an offloaded layer this is the temp device, so it should use the train device instead.

LinearFp8 additionally keeps self._scale as a plain attribute aliasing the registered scale buffer. nn.Module._apply / .to() rebind the buffer but not the attribute, so the two can diverge onto different devices independently of the above; folding _scale into the buffer would remove that.

---

Drafted by Claude

[dxqb]

@Koratahiu FYI: whenever tensors are created during initialization of pefts, they can't copy orig_module.weight.device, because half of layers are on CPU if offloading is enabled. the layers are moved to cuda later, but the peft tensors never are