Optimization fused adam multi tensor impl for flagos and torch

Author: lihongyang1990

transformer_engine/plugin/core/backends/flagos/flagos.py

@@ -163,7 +163,7 @@ def multi_tensor_adam(
             return multi_tensor_adam_fl
         return multi_tensor_adam_fl(
             chunk_size=chunk_size, noop_flag=noop_flag, tensor_lists=tensor_lists,
-            lr=lr, beta1=beta1, beta2=beta2, eps=epsilon,
+            lr=lr, beta1=beta1, beta2=beta2, epsilon=epsilon,
             step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
         )
     def multi_tensor_adam_param_remainder(
@@ -184,7 +184,7 @@ def multi_tensor_adam_param_remainder(
             return multi_tensor_adam_param_remainder_fl
         return multi_tensor_adam_param_remainder_fl(
             chunk_size=chunk_size, noop_flag=noop_flag, tensor_lists=tensor_lists,
-            lr=lr, beta1=beta1, beta2=beta2, eps=epsilon,
+            lr=lr, beta1=beta1, beta2=beta2, epsilon=epsilon,
             step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
         )
 

transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py

@@ -14,15 +14,20 @@ def multi_tensor_adam_fl(
     lr: float,
     beta1: float,
     beta2: float,
-    eps: float,
+    epsilon: float,
     step: int,
     mode: int,
     bias_correction: int,
     weight_decay: float,
     inv_scale: Optional[float] = 1.0,
     out_dtype: Optional[torch.dtype] = None,
 ) -> None:
+    """
+    Adam optimizer implementation matching CUDA exactly.
 
+    mode == 0: L2 regularization (add weight_decay * param to gradient before moment update)
+    mode == 1: AdamW (add weight_decay * param to update after moment computation)
+    """
     num_lists = len(tensor_lists)
     assert num_lists in [4, 5], f"Expected 4 or 5 tensor lists, got {num_lists}"
 
@@ -50,66 +55,67 @@ def multi_tensor_adam_fl(
         if not g.is_contiguous():
             g = g.contiguous()
 
+        # Convert to float for computation (matches CUDA's MATH_T = float)
+        g = g.float()
+        p_float = p.float()
+
         if inv_scale is not None and inv_scale != 1.0:
             g = flag_gems.mul(g, inv_scale)
 
-        m = flag_gems.add_(flag_gems.mul_(m, beta1), g, alpha=1-beta1)
-        v = flag_gems.add_(flag_gems.mul_(v, beta2), flag_gems.mul_(flag_gems.mul_(g, g), 1 - beta2))
+        if mode == 0:  # L2 regularization
+            # Add weight decay to gradient before moment update
+            g = flag_gems.add(g, p_float, alpha=weight_decay)
+
+            # Update moments with modified gradient
+            flag_gems.add_(flag_gems.mul_(m, beta1), g, alpha=1-beta1)
+            flag_gems.add_(flag_gems.mul_(v, beta2), flag_gems.mul(g, g), alpha=1-beta2)
+
+            # Bias correction
+            m_corr = flag_gems.true_divide(m, bias_correction1)
+            v_corr = flag_gems.true_divide(v, bias_correction2)
 
-        m_corr = m.clone()
-        v_corr = v.clone()
-        if bias_correction == 1:
-            m_corr = flag_gems.true_divide(m_corr, bias_correction1)
-            v_corr = flag_gems.true_divide(v_corr, bias_correction2)
+            # Compute update
+            denom = flag_gems.add(flag_gems.sqrt(v_corr), epsilon)
+            update = flag_gems.true_divide(m_corr, denom)
 
-        update = flag_gems.true_divide(m_corr, flag_gems.add(flag_gems.sqrt(v_corr), eps))
+            # Update parameter
+            p.add_(update, alpha=-lr)
+        else:  # mode == 1, AdamW (decoupled weight decay)
+            # Update moments with original gradient
+            flag_gems.add_(flag_gems.mul_(m, beta1), g, alpha=1-beta1)
+            flag_gems.add_(flag_gems.mul_(v, beta2), flag_gems.mul(g, g), alpha=1-beta2)
 
-        if is_adamw:
-            p = flag_gems.mul_(p, 1 - lr * weight_decay)
-        else:
-            update = flag_gems.add_(update, p, alpha=weight_decay)
+            # Bias correction
+            m_corr = flag_gems.true_divide(m, bias_correction1)
+            v_corr = flag_gems.true_divide(v, bias_correction2)
 
-        p = flag_gems.add_(p, update, alpha=-lr)
+            # Compute update with weight decay added (matches CUDA exactly)
+            denom = flag_gems.add(flag_gems.sqrt(v_corr), epsilon)
+            update = flag_gems.add(flag_gems.true_divide(m_corr, denom), p_float, alpha=weight_decay)
+
+            # Update parameter
+            p.add_(update, alpha=-lr)
 
         if p_master is not None:
             flag_gems.copy_(p_master, p)
             out_dtype = p_master.dtype if out_dtype is None else out_dtype
             p.data = p.data.to(out_dtype)
 
-
 def multi_tensor_adam_param_remainder_fl(
     chunk_size: int,
     noop_flag: torch.Tensor,
     tensor_lists: List[List[torch.Tensor]],
     lr: float,
     beta1: float,
     beta2: float,
-    eps: float,
+    epsilon: float,
     step: int,
     mode: int,
     bias_correction: int,
     weight_decay: float,
-    inv_scale: Optional[float] = 1.0,
 ) -> None:
     """
     Adam optimizer with parameter remainders for BF16 precision (FlagOS implementation).
-
-    This variant stores BF16 parameters + int16 remainders to reconstruct FP32 master weights.
-    Used when you have BF16 params and need FP32 master params without storing full FP32 copies.
-
-    Args:
-        chunk_size: Chunk size for processing (unused in this implementation)
-        noop_flag: If non-zero, skip computation
-        tensor_lists: [grads, params (bf16), exp_avgs (fp32), exp_avg_sqs (fp32), param_remainders (int16)]
-        lr: Learning rate
-        beta1: First moment decay rate
-        beta2: Second moment decay rate
-        eps: Epsilon for numerical stability
-        step: Current optimization step
-        mode: 0 = L2 regularization, 1 = AdamW (decoupled weight decay)
-        bias_correction: Whether to apply bias correction (1 = yes, 0 = no)
-        weight_decay: Weight decay coefficient
-        inv_scale: Inverse gradient scale for mixed precision training
     """
     if noop_flag.item() != 0:
         return
@@ -133,58 +139,78 @@ def multi_tensor_adam_param_remainder_fl(
 
     for i in range(num_tensors):
         g = tensor_lists[0][i]
-        p = tensor_lists[1][i]  # BF16 parameter
+        p = tensor_lists[1][i]  # int16 parameter (high 16 bits of FP32)
         m = tensor_lists[2][i]  # FP32 first moment
         v = tensor_lists[3][i]  # FP32 second moment
-        p_remainder = tensor_lists[4][i]  # int16 remainder
+        p_remainder = tensor_lists[4][i]  # int16 remainder (low 16 bits of FP32)
 
         if not g.is_contiguous():
             g = g.contiguous()
 
-        # Apply gradient unscaling if needed
-        if inv_scale is not None and inv_scale != 1.0:
-            g = flag_gems.mul(g, inv_scale)
+        # Convert gradient to float
+        g_float = g.float()
+
+        # Reconstruct FP32 master weight from int16 param + int16 remainder using bit manipulation
+        # This matches the CUDA implementation exactly:
+        # 1. If p_remainder < 0, decrement p (undo rounding)
+        # 2. Combine high 16 bits (p) and low 16 bits (p_remainder) into FP32
+        # Note: Use PyTorch native ops for bit manipulation (int16/int32 operations)
+
+        local_p = p.view(torch.int16).clone()
+        local_p_rem = p_remainder.clone()
 
-        # Reconstruct FP32 master weight from BF16 param + int16 remainder
-        # The remainder represents the lower 16 bits lost in BF16 conversion
-        param_fp32 = p.float()
-        param_master = flag_gems.add(param_fp32, flag_gems.mul(p_remainder.float(), 2.0 ** -16))
+        # Undo rounding: if remainder < 0, decrement p
+        local_p = torch.where(local_p_rem < 0, local_p - 1, local_p)
 
-        # Compute gradient with weight decay (if L2 mode)
-        grad_with_decay = g.float()
-        if not is_adamw:  # L2 regularization mode
-            grad_with_decay = flag_gems.add(grad_with_decay, flag_gems.mul(param_master, weight_decay))
+        # Combine into FP32 using bit shift operations
+        # local_p is high 16 bits, local_p_rem is low 16 bits
+        high_bits = local_p.to(torch.int32) << 16
+        low_bits = local_p_rem.to(torch.int32) & 0xFFFF  # Mask off sign extension
+        param_int32 = high_bits | low_bits
+        param_master = param_int32.view(torch.float32)
 
-        # Update moments
-        m = flag_gems.add_(flag_gems.mul_(m, beta1), grad_with_decay, alpha=1 - beta1)
-        v = flag_gems.add_(flag_gems.mul_(v, beta2), flag_gems.mul_(flag_gems.mul_(grad_with_decay, grad_with_decay), 1 - beta2))
+        # L2 mode: add weight decay to gradient before updating moments
+        if not is_adamw and weight_decay != 0:
+            g_float = flag_gems.add(g_float, param_master, alpha=weight_decay)
+
+        # Update first moment: m = beta1 * m + (1 - beta1) * g
+        flag_gems.add_(flag_gems.mul_(m, beta1), g_float, alpha=1 - beta1)
+
+        # Update second moment: v = beta2 * v + (1 - beta2) * g^2
+        flag_gems.add_(flag_gems.mul_(v, beta2), flag_gems.mul(g_float, g_float), alpha=1 - beta2)
 
         # Apply bias correction
-        m_corr = m.clone()
-        v_corr = v.clone()
-        if bias_correction == 1:
-            m_corr = flag_gems.true_divide(m_corr, bias_correction1)
-            v_corr = flag_gems.true_divide(v_corr, bias_correction2)
+        m_corr = flag_gems.true_divide(m, bias_correction1)
+        v_corr = flag_gems.true_divide(v, bias_correction2)
+
+        # Compute denominator: sqrt(v_corr) + epsilon
+        denom = flag_gems.add(flag_gems.sqrt(v_corr), epsilon)
 
         # Compute update
-        update = flag_gems.true_divide(m_corr, flag_gems.add(flag_gems.sqrt(v_corr), eps))
+        update = flag_gems.true_divide(m_corr, denom)
+
+        # AdamW mode: add decoupled weight decay to update
+        if is_adamw and weight_decay != 0:
+            update = flag_gems.add(update, param_master, alpha=weight_decay)
 
-        # Apply weight decay (if AdamW mode)
-        if is_adamw:
-            param_master = flag_gems.mul_(param_master, 1 - lr * weight_decay)
+        # Update master weight: p = p - lr * update
+        param_master = flag_gems.sub(param_master, flag_gems.mul(update, lr))
 
-        # Update master weight
-        param_master = flag_gems.add_(param_master, update, alpha=-lr)
+        # Split FP32 back into int16 param + int16 remainder using bit manipulation
+        # This matches the CUDA implementation exactly:
+        # 1. Extract high 16 bits as p
+        # 2. Extract low 16 bits as p_remainder
+        # 3. If p_remainder < 0, increment p (round up)
+        # Note: Use PyTorch native ops for bit manipulation (int32 operations)
 
-        # Split back into BF16 param + int16 remainder
-        # Convert to BF16 (this is the rounded version)
-        param_bf16 = param_master.to(dtype=p.dtype)
+        param_int32 = param_master.view(torch.int32)
+        # Extract low 16 bits (remainder) and high 16 bits (param)
+        new_p_rem = (param_int32 & 0xFFFF).to(torch.int16)
+        new_p = ((param_int32 >> 16) & 0xFFFF).to(torch.int16)
 
-        # Compute remainder: difference between FP32 master and BF16 representation
-        # Scale and quantize to int16 range
-        remainder_fp32 = flag_gems.mul(flag_gems.sub(param_master, param_bf16.float()), 2.0 ** 16)
-        remainder_int16 = flag_gems.clamp(torch.round(remainder_fp32), -32768, 32767).to(dtype=torch.int16)
+        # Round up: if remainder < 0, increment p
+        new_p = torch.where(new_p_rem < 0, new_p + 1, new_p)
 
         # Write back
-        flag_gems.copy_(p, param_bf16)
-        flag_gems.copy_(p_remainder, remainder_int16)
+        p.view(torch.int16).copy_(new_p)
+        p_remainder.copy_(new_p_rem)

transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py

@@ -34,6 +34,9 @@ def multi_tensor_l2_norm_fl(
 
     for tensor in tensors:
         norm_sq = flag_gems.sum(tensor.float() ** 2)
+        # Check for inf/nan (matches CUDA behavior)
+        if not torch.isfinite(norm_sq):
+            noop_flag.fill_(1)
         total_norm_sq = total_norm_sq + norm_sq
         if per_tensor:
             per_tensor_norms.append(flag_gems.sqrt(norm_sq))
@@ -47,7 +50,6 @@ def multi_tensor_l2_norm_fl(
 
     return total_norm, per_tensor_result
 
-
 def multi_tensor_scale_fl(
     _chunk_size: int,
     noop_flag: torch.Tensor,
@@ -58,4 +60,7 @@ def multi_tensor_scale_fl(
         return
 
     for src, dst in zip(tensor_lists[0], tensor_lists[1]):
+        # Check for inf/nan (matches CUDA behavior for AMP gradient scaling)
+        if not torch.isfinite(src).all():
+            noop_flag.fill_(1)
         flag_gems.copy_(dst, src * scale)