Tier_4/tier4_2_encoder_unfrozen_lr_patch.py

"""
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 <save_root>/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).")