Fix: Handle FusedLAMB Fails with LowLevelZeroPlugin When Using CPU Offload

colossalai/nn/optimizer/cpu_adam.py

@@ -190,17 +190,21 @@ def step(self, closure=None, div_scale: float = -1):
                         )
                     self._post_update(p, "exp_avg", "exp_avg_sq")
                 elif target_device.type == "cuda":
-                    assert div_scale == -1, "div_scale should remain default"
                     assert state["exp_avg"].device.type == "cuda", "exp_avg should stay on cuda"
                     assert state["exp_avg_sq"].device.type == "cuda", "exp_avg should stay on cuda"
 
                     bias_correction1 = 1 - beta1 ** state["step"]
                     bias_correction2 = 1 - beta2 ** state["step"]
 
+                    # scale gradient if div_scale is provided
+                    grad = p.grad.data
+                    if div_scale != -1:
+                        grad = grad / div_scale
+
                     # adam on cuda
                     self.torch_adam_update(
                         p.data,
-                        p.grad.data,
+                        grad,
                         state["exp_avg"],
                         state["exp_avg_sq"],
                         group["lr"],

colossalai/quantization/fp8.py

@@ -797,37 +797,32 @@ def forward(
         ctx.w_fp8_t = w_fp8.t()
         ctx.inv_scale_x = inv_scale_x
         ctx.inv_scale_w = inv_scale_w
-        out = torch._scaled_mm(
-            x_fp8,
-            ctx.w_fp8_t,
-            bias=bias,
-            out_dtype=ctx.out_dtype,
-            scale_a=inv_scale_x,
-            scale_b=inv_scale_w,
-            use_fast_accum=True,
-        )[0]
+
+        # Dequantize and compute matrix multiplication (compatible with TorchDynamo)
+        x_deq = x_fp8.to(ctx.out_dtype) * inv_scale_x
+        w_t_deq = ctx.w_fp8_t.to(ctx.out_dtype) * inv_scale_w
+
+        out = x_deq @ w_t_deq
+        if bias is not None:
+            out = out + bias.to(ctx.out_dtype)
+
+        out = out.to(ctx.out_dtype)
         return out.reshape(*ctx.x_shape[:-1], w.shape[0])
 
     @staticmethod
     def backward(ctx: Any, out_grad) -> Any:
         out_grad = out_grad.reshape(-1, out_grad.shape[-1])
         out_grad_fp8, out_grad_scale = cast_to_fp8(out_grad, fp8_format="e5m2")
-        x_grad = torch._scaled_mm(
-            out_grad_fp8,
-            ctx.w_fp8_t.contiguous().t(),
-            out_dtype=ctx.out_dtype,
-            scale_a=out_grad_scale,
-            scale_b=ctx.inv_scale_w,
-            use_fast_accum=True,
-        )[0]
-        w_grad = torch._scaled_mm(
-            out_grad_fp8.t().contiguous(),
-            ctx.x_fp8.t().contiguous().t(),
-            out_dtype=ctx.out_dtype,
-            scale_a=out_grad_scale,
-            scale_b=ctx.inv_scale_x,
-            use_fast_accum=True,
-        )[0]
+
+        # Dequantize (force contiguous after cast)
+        out_grad_deq = (out_grad_fp8.to(ctx.out_dtype) * out_grad_scale).contiguous()
+        w_t_deq = (ctx.w_fp8_t.to(ctx.out_dtype) * ctx.inv_scale_w).contiguous()
+        x_deq = (ctx.x_fp8.to(ctx.out_dtype) * ctx.inv_scale_x).contiguous()
+
+        # Compute gradients
+        x_grad = out_grad_deq @ w_t_deq.t()
+        w_grad = out_grad_deq.t() @ x_deq
+
         bias_grad = None
         if ctx.has_bias:
             bias_grad = out_grad.sum(0)

colossalai/utils/multi_tensor_apply/multi_tensor_apply.py

@@ -4,6 +4,7 @@
 class MultiTensorApply(object):
     """
     Apply an operation to a list of tensors efficiently.
+    Move tensors to CUDA if they are on CPU.
 
     Args:
         chunk_size (int): Size of a chunk.
@@ -32,4 +33,10 @@ def check_avail(self):
     def __call__(self, op, noop_flag_buffer, tensor_lists, *args):
         self.check_avail()
 
+        # Move tensors to GPU if not already on GPU
+        for i, tensor_list in enumerate(tensor_lists):
+            for j, tensor in enumerate(tensor_list):
+                if tensor.device.type == "cpu":
+                    tensor_lists[i][j] = tensor.to("cuda")
+
         return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)