Adapt SGDM layer learning rates from raw WeightWatcher alpha

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):
@@ -55,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,

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)

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)})