From 31cf84cd8e0e8ce8dc2f2d0a6bdbf8c9370827da Mon Sep 17 00:00:00 2001 From: dxqb Date: Tue, 17 Mar 2026 19:29:29 +0100 Subject: [PATCH 1/3] SR momentum --- adv_optm/optim/AdamW_adv.py | 32 ++++++++++++++++++++++++++++++-- adv_optm/util/param_update.py | 4 ++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/adv_optm/optim/AdamW_adv.py b/adv_optm/optim/AdamW_adv.py index b569711..ac31712 100644 --- a/adv_optm/optim/AdamW_adv.py +++ b/adv_optm/optim/AdamW_adv.py @@ -145,6 +145,7 @@ def __init__( "nnmf_factor": nnmf_factor } self.stochastic_rounding = stochastic_rounding + self.momentum_stochastic_rounding = True self.cautious_mask = cautious_mask self.grams_moment = grams_moment self.use_AdEMAMix = use_AdEMAMix @@ -336,7 +337,21 @@ def _step_parameter(self, p, grad, state, group, step_size, beta1, beta2, sqrt_b else: # Standard AdamW logic for non-factored tensors if beta1 > 0: exp_avg = state['exp_avg'] - exp_avg.lerp_(grad, 1.0 - beta1) + if exp_avg.dtype == torch.bfloat16 and self.momentum_stochastic_rounding: + exp_avg_fp32 = exp_avg.float() + grad_fp32 = grad.float() + #TODO duplicated code + exp_avg_fp32.lerp_(grad_fp32, 1.0 - beta1) + if random_int_tensor is not None: + # Compiled path: use the pre-computed random tensor + #TODO use the same random int tensor? + param_update._copy_stochastic_core_(exp_avg, exp_avg_fp32, random_int_tensor) + else: + assert False + # Uncompiled path: generate randoms inside + param_update.copy_stochastic_(exp_avg, exp_avg_fp32) + else: + exp_avg.lerp_(grad, 1.0 - beta1) if self.grams_moment: update_mt = _grams_update(exp_avg, grad) @@ -357,7 +372,20 @@ def _step_parameter(self, p, grad, state, group, step_size, beta1, beta2, sqrt_b update = update_mt if beta1 > 0 else grad.clone() exp_avg_sq = state['exp_avg_sq'] - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + if exp_avg_sq.dtype == torch.bfloat16 and self.momentum_stochastic_rounding: + exp_avg_sq_fp32 = exp_avg_sq.float() + grad_fp32 = grad.float() + #TODO duplicated code + exp_avg_sq_fp32.mul_(beta2).addcmul_(grad_fp32, grad_fp32, value=1 - beta2) + if random_int_tensor is not None: + # Compiled path: use the pre-computed random tensor + param_update._copy_stochastic_core_(exp_avg_sq, exp_avg_sq_fp32, random_int_tensor) + else: + assert False + # Uncompiled path: generate randoms inside + param_update.copy_stochastic_(exp_avg_sq, exp_avg_sq_fp32) + else: + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) if group['use_atan2']: denom = exp_avg_sq.sqrt() diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index a116b7d..dd8e8b2 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -133,7 +133,7 @@ def _copy_stochastic_core_(target: Tensor, source: Tensor, random_int_tensor: Te Core logic for stochastic rounding using a pre-computed random integer tensor. This version is designed to be torch.compile-friendly. """ - result = random_int_tensor + result = random_int_tensor.clone() #TODO reused for now # add the random number to the lower 16 bit of the mantissa result.add_(source.view(dtype=torch.int32)) @@ -159,7 +159,7 @@ def copy_stochastic_(target: Tensor, source: Tensor): """ random_int_tensor = _get_random_int_for_sr(source) _copy_stochastic_core_(target, source, random_int_tensor) - del random_int_tensor + #del random_int_tensor FIXME def add_stochastic_(input: Tensor, other: Tensor, alpha: float = 1.0): From 715e9f9e0cbac073c7ef71b3d32c4d7e4157724b Mon Sep 17 00:00:00 2001 From: dxqb Date: Tue, 17 Mar 2026 19:30:39 +0100 Subject: [PATCH 2/3] SR momentum --- adv_optm/optim/AdamW_adv.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/adv_optm/optim/AdamW_adv.py b/adv_optm/optim/AdamW_adv.py index ac31712..d7c490a 100644 --- a/adv_optm/optim/AdamW_adv.py +++ b/adv_optm/optim/AdamW_adv.py @@ -347,7 +347,6 @@ def _step_parameter(self, p, grad, state, group, step_size, beta1, beta2, sqrt_b #TODO use the same random int tensor? param_update._copy_stochastic_core_(exp_avg, exp_avg_fp32, random_int_tensor) else: - assert False # Uncompiled path: generate randoms inside param_update.copy_stochastic_(exp_avg, exp_avg_fp32) else: @@ -381,7 +380,6 @@ def _step_parameter(self, p, grad, state, group, step_size, beta1, beta2, sqrt_b # Compiled path: use the pre-computed random tensor param_update._copy_stochastic_core_(exp_avg_sq, exp_avg_sq_fp32, random_int_tensor) else: - assert False # Uncompiled path: generate randoms inside param_update.copy_stochastic_(exp_avg_sq, exp_avg_sq_fp32) else: From 8c87709fd16dcf2c294d59ed723d6880ecb0e03e Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 28 Mar 2026 12:18:17 +0100 Subject: [PATCH 3/3] debug message print debug message to make sure it's enabled, until there is a UI to config it --- adv_optm/optim/AdamW_adv.py | 1 + 1 file changed, 1 insertion(+) diff --git a/adv_optm/optim/AdamW_adv.py b/adv_optm/optim/AdamW_adv.py index d7c490a..00c4066 100644 --- a/adv_optm/optim/AdamW_adv.py +++ b/adv_optm/optim/AdamW_adv.py @@ -145,6 +145,7 @@ def __init__( "nnmf_factor": nnmf_factor } self.stochastic_rounding = stochastic_rounding + print("Stochastic rounding of momentums enabled!") self.momentum_stochastic_rounding = True self.cautious_mask = cautious_mask self.grams_moment = grams_moment