From 9922568211544de3d848667f50e5fa8f9df7e807 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 20:40:40 -0700 Subject: [PATCH 1/4] Configure adaptive alpha control --- baseline/gpt2_small/muon_speedrun/benchmark_config.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index 7f62454..3ee056d 100644 --- a/baseline/gpt2_small/muon_speedrun/benchmark_config.py +++ b/baseline/gpt2_small/muon_speedrun/benchmark_config.py @@ -9,7 +9,7 @@ LLMC_PROFILE = 'llmc_reference' SGD_SPEEDRUN_PROFILE = 'sgd_speedrun' PROFILE = os.environ.get('RG_GPT2_PROFILE', LLMC_PROFILE) -BENCHMARK = ('stock-gpt2-124m-fineweb10b-sgdm-layerlr-v1' if PROFILE == SGD_SPEEDRUN_PROFILE else 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2') +BENCHMARK = ('stock-gpt2-124m-fineweb10b-sgdm-alpha-control-v1' if PROFILE == SGD_SPEEDRUN_PROFILE else 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2') SOURCE_COMMIT = '7ecd8906afe6ed7a2b2cdb731c042f26d525b820' TOTAL_STEPS = 19560 BATCH_TOKENS = 524288 @@ -30,6 +30,13 @@ SGD_OTHER_MULTIPLIER = 0.75 SGD_BLOCK_MULTIPLIERS = (0.75, 0.80, 0.85, 0.90, 1.20, 1.35, 1.45, 1.40, 1.30, 1.15, 1.00, 0.90) RESUMABLE_CHECKPOINT_INTERVAL = 2000 +ALPHA_TARGET = 2.0 +ALPHA_CONTROL_GAIN = 0.08 +ALPHA_CONTROL_EMA = 0.5 +ALPHA_MULTIPLIER_MIN = 0.5 +ALPHA_MULTIPLIER_MAX = 2.0 +ALPHA_MAX_LOG_CHANGE = 0.12 +ALPHA_CONTROL_MIN_STEP = 1000 def lr_factor(step): From e6ea048794b01f618be0bd62f4afdc24f117d86d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 20:40:45 -0700 Subject: [PATCH 2/4] Publish atomic raw-alpha control signal --- baseline/gpt2_small/muon_speedrun/tracking.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/baseline/gpt2_small/muon_speedrun/tracking.py b/baseline/gpt2_small/muon_speedrun/tracking.py index 893bb33..a58ba23 100644 --- a/baseline/gpt2_small/muon_speedrun/tracking.py +++ b/baseline/gpt2_small/muon_speedrun/tracking.py @@ -159,6 +159,24 @@ def measure(path): return {'layers':rows, 'summary':summary(rows, identity), 'options':WW_OPTIONS} +def publish_alpha_control(root, result): + """Publish completed raw-alpha block means atomically for the trainer.""" + by_block = {} + for row in result['layers']: + alpha = row.get('alpha_raw') + if alpha is not None: + by_block.setdefault(int(row['block']), []).append(float(alpha)) + blocks = {} + for block in range(12): + values = by_block.get(block, []) + if len(values) >= 4: + blocks[str(block)] = {'raw_alpha_mean':sum(values)/len(values), + 'valid_count':len(values)} + write_json(root/'ALPHA_CONTROL.json', + {'schema':1, 'step':result['summary']['step'], 'source':'raw_alpha', + 'blocks':blocks}) + + def csv_write(path, rows): temp = path.with_suffix('.csv.tmp') fields = list(dict.fromkeys(key for row in rows for key in row)) @@ -203,6 +221,7 @@ def watch(root, deadline): result = measure(path) write_json(output, result) collect(root) + publish_alpha_control(root, result) print('[weightwatcher] '+json.dumps(result['summary']), flush=True) except Exception as exc: write_json(failure, {'step':int(path.stem), 'error':repr(exc)}) From d388360a3726bfde27711d55f11713a8f550f6cf Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 20:40:54 -0700 Subject: [PATCH 3/4] Adapt SGDM block LRs from raw alpha feedback --- baseline/gpt2_small/muon_speedrun/run.py | 49 +++++++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py index c1bef65..fe261d0 100644 --- a/baseline/gpt2_small/muon_speedrun/run.py +++ b/baseline/gpt2_small/muon_speedrun/run.py @@ -21,7 +21,9 @@ from benchmark_config import (BENCHMARK, TOTAL_STEPS, BATCH_TOKENS, VAL_TOKENS, MEASUREMENT_INTERVAL, WARMUP, WARMDOWN, LR_SCHEDULE, GRAD_CLIP, SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, SGD_NESTEROV, - lr_factor, protocol, RESUMABLE_CHECKPOINT_INTERVAL) + lr_factor, protocol, RESUMABLE_CHECKPOINT_INTERVAL, ALPHA_TARGET, ALPHA_CONTROL_GAIN, + ALPHA_CONTROL_EMA, ALPHA_MULTIPLIER_MIN, ALPHA_MULTIPLIER_MAX, + ALPHA_MAX_LOG_CHANGE, ALPHA_CONTROL_MIN_STEP) TARGET = 3.28 HERE = Path(__file__).resolve().parent @@ -86,6 +88,46 @@ def save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validat print('Checkpoint saved at update '+str(step), flush=True) +def alpha_controller(root, adam, state, step): + """Consume only completed atomic WeightWatcher measurements; adjust block peak LRs.""" + path = root/'ALPHA_CONTROL.json' + if step < ALPHA_CONTROL_MIN_STEP or not path.exists(): + return state, None + try: + signal = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return state, None + measured_step = int(signal.get('step', -1)) + if measured_step <= state.get('last_measurement_step', -1) or measured_step > step: + return state, None + event = {'kind':'alpha_lr_control', 'step':step, 'measurement_step':measured_step, 'blocks':{}} + for block in range(12): + item = signal.get('blocks', {}).get(str(block)) + if not item or int(item.get('valid_count', 0)) < 4: + continue + alpha = float(item['raw_alpha_mean']) + old_ema = state.setdefault('ema', {}).get(str(block), alpha) + ema = ALPHA_CONTROL_EMA*old_ema + (1-ALPHA_CONTROL_EMA)*alpha + state['ema'][str(block)] = ema + current = state.setdefault('multipliers', {}).get(str(block)) + if current is None: + candidates = [g.get('lr_multiplier') for g in adam.param_groups + if g['role'].startswith(f'block_{block}_')] + current = candidates[0] + log_delta = max(-ALPHA_MAX_LOG_CHANGE, min(ALPHA_MAX_LOG_CHANGE, + ALPHA_CONTROL_GAIN*(ema-ALPHA_TARGET))) + new = max(ALPHA_MULTIPLIER_MIN, min(ALPHA_MULTIPLIER_MAX, current*math.exp(log_delta))) + state['multipliers'][str(block)] = new + for group in adam.param_groups: + if group['role'].startswith(f'block_{block}_'): + group['lr_multiplier'] = new + group['peak_lr'] = SGD_PEAK_LR*new + event['blocks'][str(block)] = {'raw_alpha_mean':alpha, 'ema':ema, + 'old_multiplier':current, 'new_multiplier':new} + state['last_measurement_step'] = measured_step + return state, event + + def target_met(validation): return bool(validation and validation.get('full_benchmark_evaluation') and validation.get('evaluation_tokens') == VAL_TOKENS @@ -252,7 +294,12 @@ def train(a): rt.step(wait=True) save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, None, best) timings = [] + alpha_state = {'last_measurement_step':-1, 'ema':{}, 'multipliers':{}} while step < TOTAL_STEPS and time.time() < a.deadline-180 and not stopped[0] and not (root/'STOP').exists(): + alpha_state, control_event = alpha_controller(root, adam, alpha_state, step) + if control_event and control_event['blocks']: + record(root, control_event) + write_json(root/'ALPHA_CONTROLLER_STATE.json', alpha_state) began = time.monotonic() model.zero_grad(set_to_none=False) loss_sum = torch.zeros((), device=rt.device, dtype=torch.float32) From 42466d46474435a8008192a98863e4c22f17dac2 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 20:41:00 -0700 Subject: [PATCH 4/4] Record adaptive alpha protocol --- baseline/gpt2_small/muon_speedrun/benchmark_config.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index 3ee056d..6653ac8 100644 --- a/baseline/gpt2_small/muon_speedrun/benchmark_config.py +++ b/baseline/gpt2_small/muon_speedrun/benchmark_config.py @@ -62,6 +62,10 @@ def protocol(): sgd_block_multipliers=list(SGD_BLOCK_MULTIPLIERS), sgd_cosine_floor=SGD_COSINE_FLOOR, resumable_checkpoint_interval=RESUMABLE_CHECKPOINT_INTERVAL, + alpha_control=dict(target=ALPHA_TARGET, gain=ALPHA_CONTROL_GAIN, ema=ALPHA_CONTROL_EMA, + multiplier_min=ALPHA_MULTIPLIER_MIN, multiplier_max=ALPHA_MULTIPLIER_MAX, + max_log_change=ALPHA_MAX_LOG_CHANGE, min_step=ALPHA_CONTROL_MIN_STEP, + source='raw_alpha_block_mean'), adam_betas=[0.9, 0.95], adam_eps=1e-8, matrix_weight_decay=0.1, gradient_clip_norm=GRAD_CLIP,