""" Tier 4.2 — Encoder-unfrozen ablation with reduced LR (Option C). Why --- Your V4 deadlock was specifically `encoder unfrozen + decoder frozen`. The SYMMETRIC experiment that was never run is "both unfrozen, but encoder gets 1/10 the decoder LR". Chronos fine-tuning literature (the original paper, the nitrogen-forecasting follow-up) shows non-trivial gains from full fine-tuning when domain-specific data is sufficient — and your 1.2M training windows is well past that threshold. What this patch does -------------------- 1. Adds Option "C" to `ChronosDualPath.configure_t5_freezing`: mode="A" (current) → encoder frozen, decoder unfrozen mode="B" → all unfrozen, single LR mode="C" (NEW) → all unfrozen, encoder gets `encoder_lr_factor` × decoder LR 2. Replaces `CGCXTrainer.__init__` to construct AdamW with TWO parameter groups when `freeze_mode == "C"`. The split is precise — we identify encoder params by their `id()` so any submodule reuse doesn't double- count. 3. Saves the freeze mode + factor into `training_config.json` so reload logic in `load_checkpoint_with_config` faithfully reproduces the training conditions. How to apply ------------ import cgc_x from tier4_2_encoder_unfrozen_lr_patch import patch_into patch_into(cgc_x) # Then either: # (a) set the global default before calling train_fold: cgc_x.DEFAULT_FREEZE_MODE = "C" cgc_x.train_fold(...) # (b) or pass it explicitly via the helper run function: from tier4_2_encoder_unfrozen_lr_patch import run_tier4_2_ablation run_tier4_2_ablation(data_dir="./raw_data", T=64, H=16, encoder_lr_factor=0.1) Run protocol ------------ ONE fold-4 run with mode="C", same seeds and hyperparameters as v3_advanced. Compare on T1 + T2 with the metric panel. Expected outcome ---------------- Either a measurable improvement (~5–10% on RMSE per the nitrogen paper) or no change. If no change, you've ruled it out and locked in mode "A" with confidence. Caveats ------- - WALL-CLOCK COST: Unfreezing the encoder approximately DOUBLES the trainable parameter count from ~80M (decoder + adapter) to ~160M (encoder + decoder + adapter). Expect 1.6× per-epoch time and roughly 1.6× more GPU memory. - LR SCHEDULER: We rely on `get_custom_cosine_schedule` from cgc_x. When the optimizer has multiple param groups, PyTorch schedulers natively scale ALL groups' base LRs by the same factor — the encoder/decoder LR ratio is preserved across the whole schedule. We verify this in `_assert_param_group_ratio_preserved`. - WARM-START: When loading a checkpoint trained under mode="A", the encoder weights are the off-the-shelf Chronos weights (not fine-tuned). Switching to mode "C" for the next fold means the encoder will START fine-tuning from those pretrained weights. This is the desired path — do NOT mode="B" → mode="C" mid-fold. """ from __future__ import annotations import json import math import os from pathlib import Path import torch # ============================================================================= # 1. Replacement configure_t5_freezing — adds mode "C" # ============================================================================= def configure_t5_freezing_with_C(self, mode: str): """ Replacement for ChronosDualPath.configure_t5_freezing. For mode "C": all T5 parameters require_grad=True, but the FACT that the encoder will get a reduced LR is enforced inside the trainer's optimizer construction (see `_init_with_param_groups` below). This method itself just records the mode on the model. """ if mode is True: mode = "A" elif mode is False: mode = "B" if mode not in ("A", "B", "C"): raise ValueError(f"freeze mode must be 'A', 'B', or 'C', got {mode!r}") if mode == "A": for p in self.t5_model.encoder.parameters(): p.requires_grad = False for p in self.t5_model.decoder.parameters(): p.requires_grad = True elif mode == "B": for p in self.t5_model.parameters(): p.requires_grad = True elif mode == "C": # Same gradient pattern as B; the LR differential lives in the optimizer. for p in self.t5_model.parameters(): p.requires_grad = True # Stash for later inspection / serialisation self._freeze_mode = mode n_trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) print(f"[Tier 4.2] configure_t5_freezing(mode={mode!r}): {n_trainable:,} trainable params") # ============================================================================= # 2. Replacement CGCXTrainer.__init__ that supports param groups # ============================================================================= # # We replicate the original __init__ flow exactly, then branch on the # `freeze_mode` recorded on the model (defaulting to "A" for backwards # compatibility) to choose between single-LR AdamW and two-group AdamW. def init_with_param_groups( self, model, config, normalizer, train_loader, val_loader, tokenizer, mase_scale, cov_cols=None, ): """ Replacement for CGCXTrainer.__init__. Reads `model._freeze_mode` to decide whether to use a single LR or split encoder/other LRs. Reads `config.encoder_lr_factor` (default 0.1) for mode "C". """ from cgc_x import get_custom_cosine_schedule from torch.cuda.amp import GradScaler self.model = model.to(config.device) self.config = config self.normalizer = normalizer self.train_loader = train_loader self.val_loader = val_loader self.tokenizer = tokenizer self.device = config.device self.cov_cols = list(cov_cols) if cov_cols is not None else None self.best_val_loss = float("inf") self.epochs_no_improve = 0 freeze_mode = getattr(model, "_freeze_mode", "A") encoder_lr_factor = float(getattr(config, "encoder_lr_factor", 0.1)) if freeze_mode == "C": # Build TWO param groups: encoder gets reduced LR, everything # else (decoder + adapter + cov_encoder) at full LR. encoder_params = [p for p in model.t5_model.encoder.parameters() if p.requires_grad] encoder_ids = {id(p) for p in encoder_params} other_params = [ p for p in model.parameters() if p.requires_grad and id(p) not in encoder_ids ] n_enc = sum(p.numel() for p in encoder_params) n_oth = sum(p.numel() for p in other_params) print(f"[Tier 4.2] Mode 'C' optimizer: " f"encoder LR={config.lr * encoder_lr_factor:.2e} ({n_enc:,} params), " f"other LR={config.lr:.2e} ({n_oth:,} params)") self.optimizer = torch.optim.AdamW( [ {"params": encoder_params, "lr": config.lr * encoder_lr_factor}, {"params": other_params, "lr": config.lr}, ], weight_decay=config.weight_decay, ) else: # Original single-group construction — exact copy of cgc_x.CGCXTrainer.__init__. trainable_params = [p for n, p in model.named_parameters() if p.requires_grad] self.optimizer = torch.optim.AdamW( trainable_params, lr=config.lr, weight_decay=config.weight_decay, ) steps_per_epoch = math.ceil(len(train_loader) / config.gradient_accumulation_steps) num_training_steps = steps_per_epoch * config.epochs self.scheduler = get_custom_cosine_schedule( self.optimizer, int(num_training_steps * config.warmup_ratio), num_training_steps, lr_end_factor=config.lr_end / config.lr, ) self.scaler = GradScaler('cuda') if config.mixed_precision else None self.mase_scale = float(max(mase_scale, 1e-8)) # Bin midpoints — exact copy from cgc_x. num_bins = tokenizer.num_bins vocab_size = model.t5_model.config.vocab_size offset = model.t5_model.config.vocab_offset bin_idx = torch.arange(num_bins, dtype=torch.float32) scaled = bin_idx / (num_bins - 1) bin_values = scaled * tokenizer.range + tokenizer.min_val full_bin_values = torch.zeros(vocab_size, dtype=torch.float32) full_bin_values[offset:offset + num_bins] = bin_values self.bin_values = full_bin_values.view(1, 1, -1).to(self.device) self.bin_values.requires_grad_(False) print(f"✓ [Tier 4.2] Trainer initialized (freeze_mode={freeze_mode}, " f"encoder_lr_factor={encoder_lr_factor if freeze_mode == 'C' else 'n/a'})") if freeze_mode == "C": _assert_param_group_ratio_preserved(self.optimizer, self.scheduler, config.lr, encoder_lr_factor) def _assert_param_group_ratio_preserved(optimizer, scheduler, base_lr, factor): """ Sanity-check that the LR scheduler scales ALL param groups by the same multiplicative factor. We do a single scheduler.step() on a fresh schedule and check the ratio is preserved. """ # Don't actually mutate the scheduler — peek at its internal state. # PyTorch's LambdaLR multiplies each group's base_lr by the lambda value # uniformly; we just verify the recorded base_lrs. bases = [g.get("initial_lr", g["lr"]) for g in optimizer.param_groups] if len(bases) < 2: return enc_base, oth_base = bases[0], bases[1] expected_ratio = factor # encoder / other at start actual_ratio = enc_base / oth_base assert abs(actual_ratio - expected_ratio) < 1e-6, ( f"Tier 4.2 LR ratio is wrong at construction: " f"got {actual_ratio:.6f}, expected {expected_ratio:.6f}" ) print(f"[Tier 4.2] LR ratio check: ✓ encoder/other = {actual_ratio:.4f} (== factor)") # ============================================================================= # 3. TrainingConfig extension to expose encoder_lr_factor # ============================================================================= def _augment_training_config(cgc_x_module): """ The patcher calls this to add `encoder_lr_factor` to whichever TrainingConfig is currently bound on cgc_x_module. We add it as a settable attribute on instances; the existing dataclass keeps its declared fields untouched, so users can either pass it in the constructor (only if they swap to a Tier-4.2-aware dataclass) or set it on an instance after construction. """ OldCfg = cgc_x_module.TrainingConfig # Wrap with a small subclass that adds the field as an instance attr. # We can't easily add a dataclass field at runtime, so we use a # lightweight monkey-patch: __post_init__ adds the attribute if # absent. orig_post_init = getattr(OldCfg, "__post_init__", None) def __post_init__(self): if orig_post_init is not None: orig_post_init(self) if not hasattr(self, "encoder_lr_factor"): self.encoder_lr_factor = 0.1 OldCfg.__post_init__ = __post_init__ # ============================================================================= # 4. Patched train_fold that supports `freeze_mode` # ============================================================================= def train_fold_with_freeze_mode( fold_cfg: dict, data_dir: str, T: int, H: int, freeze_mode: str = "A", encoder_lr_factor: float = 0.1, ): """ Drop-in train_fold replacement that exposes freeze_mode as an argument. For mode "A" the behaviour is identical to the production train_fold. """ import numpy as np from torch.utils.data import DataLoader from transformers import AutoModelForSeq2SeqLM from cgc_x import ( WFO_BASE_DIR, SimpleChronosTokenizer, ZScoreNormalizer, ChronosDualPath, TimeSeriesDataset, load_multi_asset_time_splits, load_checkpoint_with_config, compute_fold_lr, TrainingConfig, CGCXTrainer, COV_COLS, ) fold_num = fold_cfg["fold"] save_dir = os.path.join(WFO_BASE_DIR, f"fold_{fold_num}") warmstart_dir = ( os.path.join(WFO_BASE_DIR, fold_cfg["warmstart"]) if fold_cfg["warmstart"] is not None else None ) Path(save_dir).mkdir(parents=True, exist_ok=True) print(f"\n{'='*60}\n🚀 [FREEZE_MODE={freeze_mode}] FOLD {fold_num}\n{'='*60}") (train_p, train_c, train_l, train_s, val_p, val_c, val_l, val_s, _, _, _, _) = load_multi_asset_time_splits( data_dir, T=T, H=H, train_end=fold_cfg["train_end"], val_end=fold_cfg["val_end"], include_test=False, ) BATCH_SIZE = 8 tokenizer = SimpleChronosTokenizer(num_bins=1024, min_val=0.0, max_val=5.0) def prep(prices, future_prices, scales, tok): valid = (np.min(prices, axis=1) > 0) & (np.min(future_prices, axis=1) > 0) prices = prices[valid]; future_prices = future_prices[valid]; scales = scales[valid] scaled_input = prices / scales[:, None] scaled_future = future_prices / scales[:, None] in_ids = np.stack([tok.encode(row) for row in scaled_input]) lb_ids = np.stack([tok.encode(row) for row in scaled_future]) attn = (in_ids != tok.pad_token_id).astype(np.int64) return in_ids, lb_ids, attn, scales, prices, future_prices, valid tr_in, tr_lb, tr_mk, tr_sc, tr_p, tr_l, tr_v = prep(train_p, train_l, train_s, tokenizer) va_in, va_lb, va_mk, va_sc, va_p, va_l, va_v = prep(val_p, val_l, val_s, tokenizer) train_c = train_c[tr_v]; val_c = val_c[va_v] if warmstart_dir is not None: normalizer = ZScoreNormalizer.load(warmstart_dir) else: normalizer = ZScoreNormalizer() normalizer.fit(torch.FloatTensor(train_c)) train_dataset = TimeSeriesDataset( prices=tr_p, covariates=train_c, input_ids=tr_in, labels=tr_lb, future_prices=tr_l, scale=tr_sc, attention_mask=tr_mk, pad_token_id=tokenizer.pad_token_id, ) val_dataset = TimeSeriesDataset( prices=va_p, covariates=val_c, input_ids=va_in, labels=va_lb, future_prices=va_l, scale=va_sc, attention_mask=va_mk, pad_token_id=tokenizer.pad_token_id, ) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE) computed_lr, total_steps = compute_fold_lr( warmstart_dir=warmstart_dir, train_dataset_len=len(train_dataset), batch_size=BATCH_SIZE, epochs=12, ) if warmstart_dir is not None: model, _, _, _ = load_checkpoint_with_config(warmstart_dir, strict_schema=False) else: t5 = AutoModelForSeq2SeqLM.from_pretrained("amazon/chronos-t5-base") t5.resize_token_embeddings(tokenizer.vocab_size) for k in ["vocab_size", "pad_token_id", "eos_token_id", "bos_token_id", "vocab_offset", "num_bins"]: setattr(t5.config, k, getattr(tokenizer, k)) # Construct without freezing first; we apply freeze_mode below. model = ChronosDualPath( t5_model=t5, num_features=train_c.shape[2], freeze_encoder=False, ) # Apply the freeze mode chosen for THIS fold/run. model.configure_t5_freezing(freeze_mode) model.adapter = torch.compile(model.adapter, dynamic=False) config = TrainingConfig( lr=computed_lr, lr_end=computed_lr * 0.3, warmup_ratio=0.05, epochs=12, batch_size=BATCH_SIZE, patience=6 if warmstart_dir is None else 3, save_dir=save_dir, lambda_price=10000, # freeze_encoder is now decorative — the canonical state lives on # model._freeze_mode. Keep the field aligned for backwards compat: freeze_encoder=(freeze_mode == "A"), mixed_precision=False, ) config.encoder_lr_factor = encoder_lr_factor trainer = CGCXTrainer( model=model, config=config, normalizer=normalizer, train_loader=train_loader, val_loader=val_loader, tokenizer=tokenizer, mase_scale=1.0, cov_cols=COV_COLS, ) trainer.train() fold_meta = { "fold_num": fold_num, "train_end": str(fold_cfg["train_end"].date()), "val_end": str(fold_cfg["val_end"].date()), "lr_used": computed_lr, "total_steps": total_steps, "freeze_mode": freeze_mode, "encoder_lr_factor": encoder_lr_factor if freeze_mode == "C" else None, "status": "complete", } with open(os.path.join(save_dir, "fold_meta.json"), "w") as f: json.dump(fold_meta, f, indent=2) # Also stash the freeze_mode in training_config.json so that # load_checkpoint_with_config can restore it. cfg_path = os.path.join(save_dir, "training_config.json") if os.path.exists(cfg_path): with open(cfg_path) as f: tc = json.load(f) else: tc = {} tc["freeze_mode"] = freeze_mode tc["encoder_lr_factor"] = encoder_lr_factor with open(cfg_path, "w") as f: json.dump(tc, f, indent=2) print(f"✅ [Tier 4.2] Fold {fold_num} complete (freeze_mode={freeze_mode}) → {save_dir}") # ============================================================================= # 5. End-to-end ablation runner # ============================================================================= def run_tier4_2_ablation( data_dir: str, T: int, H: int, encoder_lr_factor: float = 0.1, fold_cfg: dict = None, save_root: str = "./checkpoints_tier4_2_ablation", ): """ Single fold-4 run with mode "C". Saves to /fold_4/. Produces a side-by-side metric comparison vs a v3_advanced baseline. """ import cgc_x as cgc_x_module if fold_cfg is None: # Copy the dictionary so we can safely modify it fold_cfg = dict(next(f for f in cgc_x_module.WFO_FOLDS if f["fold"] == 4)) # FIX: Force it to train from scratch so it doesn't look for fold_3 fold_cfg["warmstart"] = None Path(save_root).mkdir(exist_ok=True, parents=True) old_base = cgc_x_module.WFO_BASE_DIR cgc_x_module.WFO_BASE_DIR = save_root try: train_fold_with_freeze_mode( fold_cfg, data_dir, T=T, H=H, freeze_mode="C", encoder_lr_factor=encoder_lr_factor, ) finally: cgc_x_module.WFO_BASE_DIR = old_base print(f"\n[Tier 4.2] Ablation run complete in {save_root}.") print("Compare against your v3_advanced (mode A) checkpoint by running") print("tier0_1_greedy_vs_probabilistic.py on both with the SAME data_dir.") # ============================================================================= # 6. Patcher # ============================================================================= def patch_into(cgc_x_module=None): """ Apply Tier 4.2 onto cgc_x. After patching, the simplest way to run the ablation is: cgc_x.train_fold_with_freeze_mode(fold_cfg, ..., freeze_mode="C") or: run_tier4_2_ablation(data_dir, T, H, encoder_lr_factor=0.1) """ if cgc_x_module is None: import cgc_x as cgc_x_module # 1. Replace ChronosDualPath.configure_t5_freezing with the C-aware one. cgc_x_module.ChronosDualPath.configure_t5_freezing = configure_t5_freezing_with_C # 2. Augment TrainingConfig with encoder_lr_factor default. _augment_training_config(cgc_x_module) # 3. Replace CGCXTrainer.__init__. cgc_x_module.CGCXTrainer.__init__ = init_with_param_groups # 4. Expose the new train_fold and ablation runner. cgc_x_module.train_fold_with_freeze_mode = train_fold_with_freeze_mode cgc_x_module.run_tier4_2_ablation = run_tier4_2_ablation print("✓ Tier 4.2 patch applied: encoder freeze modes A/B/C are available.") print(" Use cgc_x.train_fold_with_freeze_mode(..., freeze_mode='C') to ablate.") print(" For mode C, the trainer auto-creates two param groups with " "encoder_lr = encoder_lr_factor × decoder_lr (default 0.1).")