Remove optimizer changes for single-purpose PR

Author: lihongyang1990

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

@@ -159,12 +159,9 @@ def multi_tensor_adam(
         bias_correction: int,
         weight_decay: float,
     ) -> None:
-        if chunk_size is None:
-            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, epsilon=epsilon,
-            step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
+            chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, epsilon,
+            step, mode, bias_correction, weight_decay,
         )
     def multi_tensor_adam_param_remainder(
         self,
@@ -180,12 +177,10 @@ def multi_tensor_adam_param_remainder(
         bias_correction: int,
         weight_decay: float,
     ) -> None:
-        if chunk_size is None:
-            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, epsilon=epsilon,
-            step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
+            chunk_size, noop_flag, tensor_lists,
+            lr, beta1, beta2, epsilon,
+            step, mode, bias_correction, weight_decay,
         )
 
     # Misc

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

@@ -14,20 +14,15 @@ def multi_tensor_adam_fl(
     lr: float,
     beta1: float,
     beta2: float,
-    epsilon: float,
+    eps: 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}"
 
@@ -55,67 +50,66 @@ 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)
 
-        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 = 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))
 
-            # Compute update
-            denom = flag_gems.add(flag_gems.sqrt(v_corr), epsilon)
-            update = flag_gems.true_divide(m_corr, denom)
+        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)
 
-            # 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)
+        update = flag_gems.true_divide(m_corr, flag_gems.add(flag_gems.sqrt(v_corr), eps))
 
-            # Bias correction
-            m_corr = flag_gems.true_divide(m, bias_correction1)
-            v_corr = flag_gems.true_divide(v, bias_correction2)
+        if is_adamw:
+            p = flag_gems.mul_(p, 1 - lr * weight_decay)
+        else:
+            update = flag_gems.add_(update, p, alpha=weight_decay)
 
-            # 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)
+        p = flag_gems.add_(p, 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,
-    epsilon: float,
+    eps: 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
@@ -139,78 +133,58 @@ def multi_tensor_adam_param_remainder_fl(
 
     for i in range(num_tensors):
         g = tensor_lists[0][i]
-        p = tensor_lists[1][i]  # int16 parameter (high 16 bits of FP32)
+        p = tensor_lists[1][i]  # BF16 parameter
         m = tensor_lists[2][i]  # FP32 first moment
         v = tensor_lists[3][i]  # FP32 second moment
-        p_remainder = tensor_lists[4][i]  # int16 remainder (low 16 bits of FP32)
+        p_remainder = tensor_lists[4][i]  # int16 remainder
 
         if not g.is_contiguous():
             g = g.contiguous()
 
-        # 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()
-
-        # Undo rounding: if remainder < 0, decrement p
-        local_p = torch.where(local_p_rem < 0, local_p - 1, local_p)
-
-        # 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)
+        # Apply gradient unscaling if needed
+        if inv_scale is not None and inv_scale != 1.0:
+            g = flag_gems.mul(g, inv_scale)
 
-        # 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)
+        # 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))
 
-        # Update first moment: m = beta1 * m + (1 - beta1) * g
-        flag_gems.add_(flag_gems.mul_(m, beta1), g_float, alpha=1 - beta1)
+        # 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))
 
-        # 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)
+        # 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))
 
         # Apply bias correction
-        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)
+        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
-        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)
+        update = flag_gems.true_divide(m_corr, flag_gems.add(flag_gems.sqrt(v_corr), eps))
 
-        # Update master weight: p = p - lr * update
-        param_master = flag_gems.sub(param_master, flag_gems.mul(update, lr))
+        # Apply weight decay (if AdamW mode)
+        if is_adamw:
+            param_master = flag_gems.mul_(param_master, 1 - lr * weight_decay)
 
-        # 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)
+        # Update master weight
+        param_master = flag_gems.add_(param_master, update, alpha=-lr)
 
-        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)
+        # Split back into BF16 param + int16 remainder
+        # Convert to BF16 (this is the rounded version)
+        param_bf16 = param_master.to(dtype=p.dtype)
 
-        # Round up: if remainder < 0, increment p
-        new_p = torch.where(new_p_rem < 0, new_p + 1, new_p)
+        # 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)
 
         # Write back
-        flag_gems.copy_(p, new_p.view(torch.bfloat16))
-        flag_gems.copy_(p_remainder, new_p_rem)
+        flag_gems.copy_(p, param_bf16)
+        flag_gems.copy_(p_remainder, remainder_int16)

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

@@ -2,65 +2,25 @@
 #
 # See LICENSE for license information.
 
-from typing import List, Tuple
 import torch
+from torch.distributed._tensor import DTensor
 import flag_gems
 
 
-def multi_tensor_l2_norm_fl(
-    _chunk_size: int,
-    noop_flag: torch.Tensor,
-    tensor_lists: List[List[torch.Tensor]],
-    per_tensor: bool = False,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Compute L2 norm of tensors using flag_gems.
-
-    Returns:
-        Tuple of (total_norm, per_tensor_norms_or_dummy)
-        - total_norm: The combined L2 norm of all tensors
-        - per_tensor_norms_or_dummy: Per-tensor norms stacked if per_tensor=True, else dummy tensor
-    """
-    device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu'
-
-    if noop_flag.item() != 0:
-        return torch.tensor(0.0, device=device), torch.tensor(0.0, device=device)
+def multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor, *args):
 
     tensors = tensor_lists[0]
 
-    # Compute per-tensor norms
-    per_tensor_norms = []
-    total_norm_sq = torch.tensor(0.0, device=device)
-
-    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))
-
-    total_norm = flag_gems.sqrt(total_norm_sq)
-
     if per_tensor:
-        per_tensor_result = torch.stack(per_tensor_norms)
+        norms = [torch.norm(t.float(), p=2) for t in tensors]
+        return norms, None
     else:
-        per_tensor_result = torch.tensor(0.0, device=device)
+        total_norm_sq = sum(flag_gems.sum(flag_gems.pow_func(t.float(), 2)) for t in tensors)
+        total_norm = flag_gems.sqrt(total_norm_sq)
+        return total_norm, None
 
-    return total_norm, per_tensor_result
 
-def multi_tensor_scale_fl(
-    _chunk_size: int,
-    noop_flag: torch.Tensor,
-    tensor_lists: List[List[torch.Tensor]],
-    scale: float,
-) -> None:
-    if noop_flag.item() != 0:
-        return
+def multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale):
 
     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)
+        flag_gems.copy_(dst, src * scale)