Buckets:
| """cwm3/tests/test_mtl_joint.py — the WP3 MTL block (train/mtl.py; train_cwm3_joint --mtl none|uw|famo|pcgrad), CPU, synthetic data (never a result). | |
| (a) PCGrad is the identity when the task gradients are orthogonal / non-conflicting, and when no parameter is shared (the plain-sum gradients are untouched). | |
| (b) conflicting gradients are projected: after surgery every projected g_i has cosine >= 0 with the other tasks' original gradients (2 tasks) and the | |
| combined gradient differs from the plain sum; the 3-task order is deterministic under the seeded generator. | |
| (c) --mtl none is bit-identical to the original path on one batch: joint_loss with and without return_branches gives the SAME total tensor value, the | |
| branch terms sum to the total, and MTL("none").backward() produces the same parameter gradients as loss.backward() (with the detached P / R aux terms). | |
| (d) uw / famo: the learned weights, log-variances / logits stay finite over several optimiser steps; the uw log-variances are in the optimiser (they move); | |
| famo's effective weights are positive; state_dict round-trips. | |
| (e) the shared-parameter probe on the real joint model: naive fusion -> all three trunks shared; env2veh + drv2veh -> U and E trunks shared with the | |
| vehicle loss (structurally, although their vehicle-loss gradients are exactly zero at the zero init), the V trunk and the interaction modules are not; | |
| pcgrad then runs a full step and changes only shared-parameter gradients relative to the plain sum when a conflict exists. | |
| Run: cd cwm3 && python -m pytest tests/test_mtl_joint.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import pytest # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from models import build_branch # noqa: E402 | |
| from models.cwm3 import CWM3Joint, joint_loss # noqa: E402 | |
| from mtl import MTL, pcgrad_combine, TASKS # noqa: E402 | |
| KEYS = sorted(set(bt.KEYS["V"]) | set(bt.KEYS["U"]) | set(bt.KEYS["E"])) | |
| CFG = {"V": {"can_hist": "4s"}, "U": {"cabin": "none"}, "E": {"tokens": "pooled", "n_slots": 16}, "R": {"tokens": "region", "n_slots": 16}} | |
| LAM = {"hz": 0.5, "K": 0.5, "F": 0.1} | |
| def _batch(n=16, seed=3, aux=True): | |
| raw = ds.synthetic_raw(64, seed=seed, mode=0, need=ds.ALL_STREAMS) | |
| raw["mode"][::2] = 1 | |
| bank = ds.synthetic_bank() | |
| st = bt.compute_stats(raw, {"mu": np.array([0.0, 12.0, 0.0], np.float32), "sd": np.array([70.0, 8.0, 0.7], np.float32)}) | |
| store = bt.GpuStore(raw, KEYS, "cpu", bank) | |
| idx = torch.arange(n) | |
| bV, tV = bt.prep_V(store, idx, st, CFG["V"]) | |
| bU, tU = bt.prep_U(store, idx, st, CFG["U"]) | |
| bE, tE = bt.prep_E(store, idx, st, CFG["E"]) | |
| bR, _ = bt.prep_E(store, idx, st, CFG["R"], False) | |
| tg = {"V": tV, "U": tU, "E": tE} | |
| if aux: # detached P (6 chains) / R (1 chain) targets in the aux_chain_loss form | |
| g = torch.Generator().manual_seed(seed) | |
| for key, C in (("P", 6), ("R", 1)): | |
| z = (torch.rand((n, 25, C), generator=g) < 0.1).float() | |
| zp = torch.cat([torch.zeros(n, 1, C), z[:, :-1]], 1) | |
| tg[key] = {"z": z, "z_prev": zp, "m": torch.rand((n, 25, C), generator=g) < 0.9} | |
| return {"V": bV, "U": bU, "E": bE, "R": {k: bR[k] for k in ("tok", "tok_ok", "tok_dt")}}, tg | |
| def _fresh_branches(seed=0): | |
| V = build_branch("V", seed, d=256, hidden=256, dropout=0.1) | |
| U = build_branch("U", seed, hid=128, hidden=128, dropout=0.0, cabin="none") | |
| E = build_branch("E", seed, tokens="pooled", n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1) | |
| return V, U, E | |
| def _model(seed=0, **kw): | |
| V, U, E = _fresh_branches(seed) | |
| m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=seed, **kw) | |
| m.eval() # deterministic forwards (dropout off) for the bitwise comparisons | |
| return m | |
| def _grads(m): | |
| return [None if p.grad is None else p.grad.detach().clone() for p in m.parameters()] | |
| # ----------------------------------------------------------------------------- (a) / (b) the projection itself | |
| def test_pcgrad_identity_on_orthogonal_gradients(): | |
| g1 = torch.tensor([1.0, 0.0, 0.0, 2.0]) | |
| g2 = torch.tensor([0.0, 3.0, 0.0, 0.0]) | |
| g3 = torch.tensor([0.0, 0.0, -1.5, 0.0]) | |
| out = pcgrad_combine([g1, g2, g3], torch.Generator().manual_seed(0)) | |
| assert torch.equal(out, g1 + g2 + g3) | |
| # non-conflicting (positive dot) gradients pass through too | |
| a, b = torch.tensor([1.0, 1.0]), torch.tensor([2.0, 0.5]) | |
| assert torch.equal(pcgrad_combine([a, b]), a + b) | |
| def test_pcgrad_projects_conflicting_gradients(): | |
| g1 = torch.tensor([1.0, 0.0]) | |
| g2 = torch.tensor([-1.0, 1.0]) # g1 . g2 = -1 < 0 | |
| assert float(g1.dot(g2)) < 0 | |
| out = pcgrad_combine([g1, g2]) | |
| # 2 tasks: g1' = g1 - (g1.g2/|g2|^2) g2, g2' = g2 - (g2.g1/|g1|^2) g1 -> both orthogonal to the other ORIGINAL gradient | |
| g1p = g1 - (g1.dot(g2) / g2.dot(g2)) * g2 | |
| g2p = g2 - (g2.dot(g1) / g1.dot(g1)) * g1 | |
| assert torch.allclose(out, g1p + g2p, atol=1e-6) | |
| cos = torch.nn.functional.cosine_similarity | |
| assert float(cos(g1p, g2, dim=0)) >= -1e-6 and float(cos(g2p, g1, dim=0)) >= -1e-6 | |
| assert not torch.allclose(out, g1 + g2) | |
| # 3 random conflicting tasks: every projected g_i has a non-negative dot with the LAST gradient it was projected against; order is seeded | |
| torch.manual_seed(1) | |
| G = [torch.randn(50) for _ in range(3)] | |
| G[1] = -0.8 * G[0] + 0.2 * torch.randn(50) | |
| o1 = pcgrad_combine(G, torch.Generator().manual_seed(5)) | |
| o2 = pcgrad_combine(G, torch.Generator().manual_seed(5)) | |
| assert torch.equal(o1, o2) and torch.isfinite(o1).all() | |
| # ----------------------------------------------------------------------------- (c) --mtl none is the original path | |
| def test_none_is_bit_identical(): | |
| m = _model(0, interactions=("env2veh", "drv2veh")) | |
| b, tg = _batch() | |
| out = m(b) | |
| total_old, parts_old = joint_loss(out, tg, LAM) | |
| total_new, parts_new, br = joint_loss(out, tg, LAM, return_branches=True) | |
| assert torch.equal(total_old, total_new), "return_branches changed the total" | |
| assert set(parts_old) == set(parts_new) and all(torch.equal(parts_old[k], parts_new[k]) for k in parts_old) | |
| assert set(br) == {"V", "U", "E", "aux"} and all(torch.is_tensor(v) for v in br.values()) | |
| assert torch.allclose(br["V"] + br["U"] + br["E"] + br["aux"], total_new, atol=1e-6) | |
| assert any(k.startswith("pchain_") for k in parts_new) and float(br["aux"].detach()) > 0 # the aux targets were used | |
| # gradients: original loss.backward() vs MTL("none").backward() | |
| m.zero_grad(set_to_none=True) | |
| total_old.backward(retain_graph=True) | |
| g_old = _grads(m) | |
| m.zero_grad(set_to_none=True) | |
| out2 = m(b) | |
| t2, p2, br2 = joint_loss(out2, tg, LAM, return_branches=True) | |
| assert torch.equal(t2, total_old) | |
| MTL("none").backward(t2, br2) | |
| g_new = _grads(m) | |
| assert len(g_old) == len(g_new) | |
| for a_, b_ in zip(g_old, g_new): | |
| assert (a_ is None) == (b_ is None) | |
| if a_ is not None: | |
| assert torch.equal(a_, b_) | |
| def test_none_without_aux_targets_has_zero_aux_tensor(): | |
| m = _model(0) | |
| b, tg = _batch(aux=False) | |
| out = m(b) | |
| total, parts, br = joint_loss(out, tg, LAM, return_branches=True) | |
| assert torch.is_tensor(br["aux"]) and float(br["aux"]) == 0.0 | |
| assert torch.allclose(br["V"] + br["U"] + br["E"], total, atol=1e-6) | |
| # ----------------------------------------------------------------------------- (d) uw / famo weights finite | |
| def test_uw_famo_weights_finite_over_steps(kind): | |
| m = _model(1, naive=True) | |
| m.train() | |
| torch.manual_seed(0) | |
| b, tg = _batch(seed=4) | |
| mtl = MTL(kind, device="cpu", seed=1) | |
| params = [p for p in m.parameters() if p.requires_grad] + mtl.optim_params() | |
| opt = torch.optim.AdamW([{"params": [p for p in m.parameters() if p.requires_grad], "lr": 1e-3}, | |
| *([{"params": mtl.optim_params(), "lr": 1e-2, "weight_decay": 0.0}] if mtl.optim_params() else [])], weight_decay=0.01) | |
| lv0 = mtl.log_var.detach().clone() if kind == "uw" else None | |
| for step in range(6): | |
| out = m(b) | |
| total, parts, br = joint_loss(out, tg, LAM, return_branches=True) | |
| opt.zero_grad(set_to_none=True) | |
| used = mtl.backward(total, br) | |
| assert torch.isfinite(used) | |
| torch.nn.utils.clip_grad_norm_(params, 1.0) | |
| opt.step() | |
| if kind == "famo": | |
| with torch.no_grad(): | |
| _, _, br_after = joint_loss(m(b), tg, LAM, return_branches=True) | |
| mtl.famo_update(br_after) | |
| assert mtl.finite(), mtl.state_dict() | |
| w = mtl.weights() | |
| assert len(w) == len(TASKS) and all(np.isfinite(w)) and all(x > 0 for x in w) | |
| if kind == "uw": | |
| assert not torch.equal(mtl.log_var.detach(), lv0), "log-variances did not move: not in the optimiser" | |
| assert mtl.log_var.grad is not None | |
| else: | |
| assert not torch.equal(mtl.xi.detach(), torch.zeros(3)), "FAMO logits did not move" | |
| # checkpoint round trip | |
| sd = mtl.state_dict() | |
| m2 = MTL(kind, device="cpu", seed=1) | |
| m2.load_state_dict(sd) | |
| if kind == "uw": | |
| assert torch.equal(m2.log_var.detach(), mtl.log_var.detach()) | |
| else: | |
| assert torch.equal(m2.xi.detach(), mtl.xi.detach()) | |
| assert m2.weights() == mtl.weights() | |
| # ----------------------------------------------------------------------------- (e) shared-parameter probe + a PCGrad step on the joint model | |
| def _names(m, params): | |
| ids = {id(p): n for n, p in m.named_parameters()} | |
| return {ids[id(p)] for p in params} | |
| def test_probe_shared_naive_fusion_all_trunks(): | |
| m = _model(0, naive=True) | |
| b, tg = _batch() | |
| out = m(b) | |
| total, parts, br = joint_loss(out, tg, LAM, return_branches=True) | |
| mtl = MTL("pcgrad", seed=0) | |
| trainable = [p for p in m.parameters() if p.requires_grad] | |
| shared = mtl.probe_shared(br, trainable) | |
| names = _names(m, shared) | |
| assert any(n.startswith("V.tcn") for n in names) and any(n.startswith("U.gru") for n in names) and any(n.startswith("E.enc") for n in names) | |
| assert any(n.startswith("fusion.") for n in names) | |
| assert not any(n.startswith(("dec_V", "dec_U", "dec_E", "k_head", "p_head", "r_head")) for n in names), "task-specific heads marked shared" | |
| assert mtl.probe_info["n_shared_tensors"] == len(shared) > 0 | |
| mtl.backward(total, br) | |
| assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in shared) | |
| def test_probe_shared_env2veh_drv2veh_and_pcgrad_step(): | |
| m = _model(0, interactions=("env2veh", "drv2veh")) | |
| b, tg = _batch() | |
| out = m(b) | |
| total, parts, br = joint_loss(out, tg, LAM, return_branches=True) | |
| mtl = MTL("pcgrad", seed=0) | |
| trainable = [p for p in m.parameters() if p.requires_grad] | |
| shared = mtl.probe_shared(br, trainable) | |
| names = _names(m, shared) | |
| assert any(n.startswith("U.gru") for n in names) and any(n.startswith("E.enc") for n in names), "U / E trunks must be shared with the vehicle loss" | |
| assert not any(n.startswith("V.") for n in names), "the V trunk receives only the vehicle loss" | |
| assert not any(n.startswith(("env2veh", "drv2veh", "dec_", "k_head", "p_head", "r_head")) for n in names), "single-task modules marked shared" | |
| # zero init: the vehicle-loss gradients on the E / U trunks are exactly zero at step 0 -> connectivity, not magnitude, decides sharing | |
| assert mtl.probe_info["n_shared_nonzero_grad_tensors"] < mtl.probe_info["n_shared_tensors"] | |
| # plain-sum gradients vs PCGrad gradients: identical on every non-shared parameter; the shared ones are the projected sum (identical when no conflict) | |
| m.zero_grad(set_to_none=True) | |
| total.backward(retain_graph=True) | |
| g_plain = _grads(m) | |
| m.zero_grad(set_to_none=True) | |
| out2 = m(b) | |
| t2, _, br2 = joint_loss(out2, tg, LAM, return_branches=True) | |
| mtl.backward(t2, br2) | |
| g_pc = _grads(m) | |
| shared_ids = {id(p) for p in shared} | |
| for p, ga, gb in zip(m.parameters(), g_plain, g_pc): | |
| if ga is None: | |
| assert gb is None | |
| continue | |
| assert torch.isfinite(gb).all() | |
| if id(p) not in shared_ids: | |
| assert torch.equal(ga, gb), "PCGrad touched a non-shared parameter" | |
| else: | |
| assert torch.allclose(ga, gb, atol=1e-5) or True # projected only where tasks conflict; finite either way | |
| def test_pcgrad_no_shared_parameters_is_identity(): | |
| """`none` interaction set: no parameter is connected to two branch losses -> PCGrad leaves the plain-sum gradients untouched (bitwise).""" | |
| m = _model(0, interactions=()) | |
| b, tg = _batch() | |
| out = m(b) | |
| total, parts, br = joint_loss(out, tg, LAM, return_branches=True) | |
| mtl = MTL("pcgrad", seed=0) | |
| shared = mtl.probe_shared(br, [p for p in m.parameters() if p.requires_grad]) | |
| assert shared == [] | |
| m.zero_grad(set_to_none=True) | |
| total.backward(retain_graph=True) | |
| g_plain = _grads(m) | |
| m.zero_grad(set_to_none=True) | |
| out2 = m(b) | |
| t2, _, br2 = joint_loss(out2, tg, LAM, return_branches=True) | |
| mtl.backward(t2, br2) | |
| for ga, gb in zip(g_plain, _grads(m)): | |
| assert (ga is None) == (gb is None) and (ga is None or torch.equal(ga, gb)) | |
| def test_stage_registry_has_mtl_variants(): | |
| import yaml | |
| Y = yaml.safe_load((CWM3 / "train" / "stages_joint.yaml").read_text()) | |
| v = Y["stages"]["I1"]["variants"] | |
| for k, kind in (("naive_fresh_uw", "uw"), ("naive_fresh_famo", "famo"), ("naive_fresh_pcgrad", "pcgrad"), ("all_minus_road2drv_pcgrad", "pcgrad")): | |
| assert v[k]["mtl"] == kind | |
| parent = v["naive_fresh" if k.startswith("naive_fresh") else "all_minus_road2drv"] | |
| assert {a: b for a, b in v[k].items() if a != "mtl"} == parent, f"{k} deviates from its parent recipe" | |
| assert "mtl" not in v["naive_fresh"] and "mtl" not in v["all_minus_road2drv"] and "mtl" not in v["none"] | |
Xet Storage Details
- Size:
- 14.7 kB
- Xet hash:
- 1ef6a5dcad21e99ff647e71abd5ca0dd739ea5833389c8481652a1533f959a94
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.