diff --git a/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py b/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py index f148795381..95602c731f 100644 --- a/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py +++ b/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -from typing import Optional, List +from typing import List import torch import flag_gems @@ -19,8 +19,6 @@ def multi_tensor_adam_fl( mode: int, bias_correction: int, weight_decay: float, - inv_scale: Optional[float] = 1.0, - out_dtype: Optional[torch.dtype] = None, ) -> None: num_lists = len(tensor_lists) @@ -50,9 +48,6 @@ def multi_tensor_adam_fl( if not g.is_contiguous(): g = g.contiguous() - 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) @@ -75,8 +70,6 @@ def multi_tensor_adam_fl( 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( @@ -91,27 +84,9 @@ def multi_tensor_adam_param_remainder_fl( 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 @@ -135,65 +110,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 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)) + # 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) - # 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) - ) + local_p = p.view(torch.int16).clone() + local_p_rem = p_remainder.clone() - # 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), - ) + # 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) + + # 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) + eps + denom = flag_gems.add(flag_gems.sqrt(v_corr), eps) # 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) - # Apply weight decay (if AdamW mode) - if is_adamw: - param_master = flag_gems.mul_(param_master, 1 - lr * weight_decay) + # 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 master weight - param_master = flag_gems.add_(param_master, update, alpha=-lr) + # Update master weight: p = p - lr * update + param_master = flag_gems.sub(param_master, flag_gems.mul(update, lr)) - # Split back into BF16 param + int16 remainder - # Convert to BF16 (this is the rounded version) - param_bf16 = param_master.to(dtype=p.dtype) + # 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) - # 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 - ) + 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) + + # 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) + flag_gems.copy_(p, new_p.view(torch.bfloat16)) + flag_gems.copy_(p_remainder, new_p_rem) diff --git a/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py b/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py index 4421487ff1..d728a76242 100644 --- a/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py +++ b/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py @@ -2,25 +2,67 @@ # # 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, noop_flag, tensor_lists, per_tensor, *args): +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) tensors = tensor_lists[0] + # Compute per-tensor norms + per_tensor_norms = [] + total_norm_sq = torch.tensor(0.0, device=device) + + for tensor in tensors: + t_float = tensor.float() + norm_sq = flag_gems.sum(flag_gems.mul(t_float, t_float)) + # Check for inf/nan (matches CUDA behavior) + if not torch.isfinite(norm_sq): + noop_flag.fill_(1) + total_norm_sq = flag_gems.add(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: - norms = [torch.norm(t.float(), p=2) for t in tensors] - return norms, None + per_tensor_result = torch.stack(per_tensor_norms) else: - 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 + per_tensor_result = torch.tensor(0.0, device=device) + + return total_norm, per_tensor_result -def multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale): +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 for src, dst in zip(tensor_lists[0], tensor_lists[1]): - flag_gems.copy_(dst, src * scale) + # 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, flag_gems.mul(src, scale)) diff --git a/transformer_engine/plugin/core/backends/reference/impl/optimizer.py b/transformer_engine/plugin/core/backends/reference/impl/optimizer.py index ceac199837..890ae9a563 100644 --- a/transformer_engine/plugin/core/backends/reference/impl/optimizer.py +++ b/transformer_engine/plugin/core/backends/reference/impl/optimizer.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -from typing import List, Union +from typing import List, Tuple, Union import torch __all__ = [ @@ -33,6 +33,9 @@ def multi_tensor_scale_torch( raise ValueError("Output and input tensor lists must have the same length") for in_tensor, out_tensor in zip(input_tensors, output_tensors): + # Check for inf/nan (matches CUDA behavior for AMP gradient scaling) + if not torch.isfinite(in_tensor).all(): + noop_flag.fill_(1) out_tensor.copy_(in_tensor * scale) @@ -41,26 +44,43 @@ def multi_tensor_l2norm_torch( noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], per_tensor: bool = False, -) -> Union[torch.Tensor, List[torch.Tensor]]: +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute L2 norm of tensors. + + 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: - if per_tensor: - return [torch.tensor(0.0, device=t.device) for t in tensor_lists[0]] - else: - return torch.tensor(0.0, device=tensor_lists[0][0].device) + return torch.tensor(0.0, device=device), torch.tensor(0.0, device=device) 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 = torch.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(torch.sqrt(norm_sq)) + + total_norm = torch.sqrt(total_norm_sq) + if per_tensor: - norms = [] - for tensor in tensors: - norm = torch.norm(tensor.float(), p=2) - norms.append(norm) - return norms + per_tensor_result = torch.stack(per_tensor_norms) else: - total_norm_sq = torch.tensor(0.0, device=tensors[0].device) - for tensor in tensors: - total_norm_sq += torch.sum(tensor.float() ** 2) - return torch.sqrt(total_norm_sq) + per_tensor_result = torch.tensor(0.0, device=device) + + return total_norm, per_tensor_result def multi_tensor_adam_torch( @@ -70,12 +90,18 @@ def multi_tensor_adam_torch( lr: float, beta1: float, beta2: float, - eps: float, + epsilon: float, step: int, mode: int, bias_correction: int, weight_decay: float, ) -> 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) + """ if noop_flag.item() != 0: return @@ -98,18 +124,43 @@ def multi_tensor_adam_torch( if grad is None: continue - if mode == 1 and weight_decay != 0: - param.mul_(1 - lr * weight_decay) + # Convert to float for computation (matches CUDA's MATH_T = float) + g = grad.float() + p = param.float() + + if mode == 0: # L2 regularization + # Add weight decay to gradient before moment update + g = g + weight_decay * p + + # Update moments with modified gradient + exp_avg.mul_(beta1).add_(g, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2) + + # Bias correction + m_corr = exp_avg / bias_correction1 + v_corr = exp_avg_sq / bias_correction2 - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + # Compute update + denom = v_corr.sqrt().add_(epsilon) + update = m_corr / denom - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + # Update parameter + param.add_(update, alpha=-lr) + else: # mode == 1, AdamW (decoupled weight decay) + # Update moments with original gradient + exp_avg.mul_(beta1).add_(g, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2) - corrected_exp_avg = exp_avg / bias_correction1 - corrected_exp_avg_sq = exp_avg_sq / bias_correction2 + # Bias correction + m_corr = exp_avg / bias_correction1 + v_corr = exp_avg_sq / bias_correction2 - denom = corrected_exp_avg_sq.sqrt().add_(eps) - param.addcdiv_(corrected_exp_avg, denom, value=-lr) + # Compute update with weight decay added (matches CUDA exactly) + denom = v_corr.sqrt().add_(epsilon) + update = (m_corr / denom) + (weight_decay * p) + + # Update parameter + param.add_(update, alpha=-lr) def multi_tensor_adam_param_remainder_torch( @@ -119,7 +170,7 @@ def multi_tensor_adam_param_remainder_torch( lr: float, beta1: float, beta2: float, - eps: float, + epsilon: float, step: int, mode: int, bias_correction: int, @@ -128,17 +179,30 @@ def multi_tensor_adam_param_remainder_torch( """ Adam optimizer with parameter remainders for BF16 precision. - 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. + This variant stores BF16 parameters + int16 remainders to reconstruct FP32 master weights + using bit manipulation, matching the CUDA implementation exactly. + + The CUDA implementation stores: + - p: int16 representing the high 16 bits of FP32 (viewed as BF16) + - p_remainder: int16 representing the low 16 bits of FP32 + + To reconstruct FP32: + - If p_remainder < 0, decrement p (undo rounding) + - Combine: fp32.int16[1] = p, fp32.int16[0] = p_remainder + + To split FP32 back: + - p = fp32.int16[1] (high 16 bits) + - p_remainder = fp32.int16[0] (low 16 bits) + - If p_remainder < 0, increment p (round up) Args: chunk_size: Chunk size for processing (unused in PyTorch implementation) noop_flag: If non-zero, skip computation - tensor_lists: [grads, params (bf16), exp_avgs (fp32), exp_avg_sqs (fp32), param_remainders (int16)] + tensor_lists: [grads, params (int16/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 + epsilon: 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) @@ -166,61 +230,76 @@ def multi_tensor_adam_param_remainder_torch( bias_correction1 = 1.0 bias_correction2 = 1.0 + is_adamw = mode == 1 + for grad, param, exp_avg, exp_avg_sq, param_remainder in zip( grads, params, exp_avgs, exp_avg_sqs, param_remainders ): - if grad is None: - continue + # Convert gradient to float + g_float = grad.float() - # Reconstruct FP32 master weight from BF16 param + int16 remainder - # The CUDA implementation uses bit manipulation to combine them - # In PyTorch, we approximate this by: - # 1. Convert param (bf16) to fp32 - this gives us the high-precision bits - # 2. Add the remainder scaled appropriately - param_fp32 = param.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 - # The remainder represents the lower 16 bits lost in BF16 conversion - # We need to scale it back to the proper magnitude - # BF16 has 16 bits total (1 sign, 8 exponent, 7 mantissa) - # The remainder compensates for the lost precision - param_master = param_fp32 + param_remainder.float() * (2.0**-16) + local_p = param.view(torch.int16).clone() + local_p_rem = param_remainder.clone() - # Standard Adam update on FP32 master weight - if mode == 0: # L2 regularization - grad_with_decay = grad.float() + weight_decay * param_master - else: # mode == 1, AdamW - grad_with_decay = grad.float() + # 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) + + # L2 mode: add weight decay to gradient before updating moments + if not is_adamw and weight_decay != 0: + g_float = g_float + weight_decay * param_master - # Update moments - exp_avg.mul_(beta1).add_(grad_with_decay, alpha=1 - beta1) - exp_avg_sq.mul_(beta2).addcmul_(grad_with_decay, grad_with_decay, value=1 - beta2) + # Update first moment: m = beta1 * m + (1 - beta1) * g + exp_avg.mul_(beta1).add_(g_float, alpha=1 - beta1) + + # Update second moment: v = beta2 * v + (1 - beta2) * g^2 + exp_avg_sq.mul_(beta2).addcmul_(g_float, g_float, value=1 - beta2) # Apply bias correction - corrected_exp_avg = exp_avg / bias_correction1 - corrected_exp_avg_sq = exp_avg_sq / bias_correction2 + m_corr = exp_avg / bias_correction1 + v_corr = exp_avg_sq / bias_correction2 + + # Compute denominator: sqrt(v_corr) + epsilon + denom = torch.sqrt(v_corr) + epsilon # Compute update - denom = corrected_exp_avg_sq.sqrt().add_(eps) - update = corrected_exp_avg / denom + update = m_corr / denom - if mode == 1: # AdamW: apply weight decay directly + # AdamW mode: add decoupled weight decay to update + if is_adamw and weight_decay != 0: update = update + weight_decay * param_master - # Update master weight - param_master.add_(update, alpha=-lr) + # Update master weight: p = p - lr * update + param_master = param_master - lr * update + + # 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) - # Split back into BF16 param + int16 remainder - # Convert to BF16 (this is the rounded version) - param_bf16 = param_master.to(dtype=param.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 = (param_master - param_bf16.float()) * (2.0**16) - remainder_int16 = remainder_fp32.round().clamp(-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 - param.copy_(param_bf16) - param_remainder.copy_(remainder_int16) + param.view(torch.int16).copy_(new_p) + param_remainder.copy_(new_p_rem) def multi_tensor_sgd_torch( diff --git a/transformer_engine/pytorch/optimizers/__init__.py b/transformer_engine/pytorch/optimizers/__init__.py index a19c797dea..c76f75743d 100644 --- a/transformer_engine/pytorch/optimizers/__init__.py +++ b/transformer_engine/pytorch/optimizers/__init__.py @@ -4,6 +4,8 @@ """Fused optimizers and multi-tensor kernels.""" from transformer_engine_torch import ( + multi_tensor_scale, + multi_tensor_l2norm, multi_tensor_unscale_l2norm, multi_tensor_adam, multi_tensor_adam_fp8,