diff --git a/giant/model/wgan.py b/giant/model/wgan.py index 161b416..ac6a852 100644 --- a/giant/model/wgan.py +++ b/giant/model/wgan.py @@ -17,7 +17,11 @@ def gradient_penalty( is for Stage 2's variable-length slot vector: both the interpolate and the critic's gradient are zeroed on padded dims first, so the norm target of 1 is only ever asked of genuine content, not the padding convention shared - by both `real` and `fake`. + by both `real` and `fake`. Rows fully masked out (e.g. `n_sec == 0`, so + every slot is padding) have no real content to constrain the gradient + norm to 1 — `x_hat`/`grad` are forced to all-zero for such a row, which + would otherwise contribute a constant `(||0|| - 1)^2 == 1` bias to the + mean regardless of critic behavior — so they're excluded from the mean. """ eps = torch.rand(real.size(0), 1, device=real.device) x_hat = eps * real + (1 - eps) * fake @@ -28,7 +32,11 @@ def gradient_penalty( grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0] if mask is not None: grad = grad * mask - return ((grad.norm(2, dim=1) - 1) ** 2).mean() + penalty = (grad.norm(2, dim=1) - 1) ** 2 + if mask is not None: + valid = (mask.sum(dim=1) > 0).float() + return (penalty * valid).sum() / valid.sum().clamp_min(1.0) + return penalty.mean() def critic_loss( diff --git a/giant/pipeline.py b/giant/pipeline.py index ed12e55..21d6f6d 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -1,3 +1,4 @@ +import os from dataclasses import dataclass from pathlib import Path @@ -256,6 +257,20 @@ def run_train_job( out_dir = Path(out_dir) + # Soft warning (never blocks) — CLAUDE.md's Compute environment section + # asks that shared portal machines (portal1/deepthought{,2}/bms{1..3}) + # stay within ~1/4 of CPU/RAM so as not to disturb other users' jobs; + # DataLoader's num_workers has no awareness of that on its own. + cpu_count = os.cpu_count() or 1 + quota = max(1, cpu_count // 4) + if num_workers > quota: + echo( + f"warning: --num-workers={num_workers} exceeds ~1/4 of this " + f"machine's {cpu_count} CPU(s) ({quota}) — portal machines are " + "shared with other users (see CLAUDE.md's Compute environment " + "section)" + ) + router_cfg = m["router"] if t["mode"] == "wgan" and router_cfg.get("enabled"): raise ValueError( diff --git a/giant/train.py b/giant/train.py index 0977a86..a6324bc 100644 --- a/giant/train.py +++ b/giant/train.py @@ -564,7 +564,9 @@ def train( # would never finish and cosine decay would barely move. steps_per_epoch = max(total_train_batches, 1) if mode == "wgan": - steps_per_epoch = max(total_train_batches // (n_critic + 1), 1) + # Generator steps fire every n_critic-th batch (did_g_step = + # step_count % n_critic == 0 in _wgan_train_step), not n_critic + 1. + steps_per_epoch = max(total_train_batches // n_critic, 1) warmup_steps = warmup_epochs * steps_per_epoch total_steps = max(epochs * steps_per_epoch, 1) @@ -579,6 +581,41 @@ def train( ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None + def _build_checkpoint(epoch: int, global_step: int, best_val_loss: float) -> dict: + ckpt: dict = { + "model": stage1_model.state_dict(), + "sec_decoder": sec_decoder.state_dict(), + "optimizer": optimizer.state_dict(), + "lr_sched": lr_sched.state_dict(), + "epoch": epoch, + "best_val_loss": best_val_loss, + "global_step": global_step, + } + if mode == "wgan": + assert ( + critic is not None + and sec_critic is not None + and optimizer_d is not None + ) + ckpt["critic"] = critic.state_dict() + ckpt["sec_critic"] = sec_critic.state_dict() + ckpt["optimizer_d"] = optimizer_d.state_dict() + if ema_decay > 0: + assert ema_stage1_model is not None and ema_sec_decoder is not None + ckpt["model_ema"] = ema_stage1_model.state_dict() + ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict() + if normalizer_dict is not None: + ckpt["normalizer"] = normalizer_dict + if pdg_map is not None: + ckpt["pdg_map"] = pdg_map + if mat_map is not None: + ckpt["mat_map"] = mat_map + if proc_map is not None: + ckpt["proc_map"] = proc_map + if model_config is not None: + ckpt["model_config"] = model_config + return ckpt + start_epoch = 1 best_val_loss = float("inf") resumed_global_step = 0 @@ -601,6 +638,15 @@ def train( critic.load_state_dict(ckpt["critic"]) sec_critic.load_state_dict(ckpt["sec_critic"]) optimizer_d.load_state_dict(ckpt["optimizer_d"]) + # Mirrors the `lr` fixup below for the generator optimizer: + # optimizer_d.load_state_dict() above restores the checkpoint's + # own critic LR, which would otherwise silently override an + # explicit `--critic-lr` passed on this resume. optimizer_d has + # no LR scheduler (unlike `optimizer`/`lr_sched`), so this is a + # flat set rather than a schedule-relative one. + resumed_critic_lr = critic_lr if critic_lr is not None else lr + for group in optimizer_d.param_groups: + group["lr"] = resumed_critic_lr optimizer.load_state_dict(ckpt["optimizer"]) lr_sched.load_state_dict(ckpt["lr_sched"]) start_epoch = ckpt.get("epoch", 0) + 1 @@ -835,6 +881,20 @@ def train( bar.close() if shutdown.requested: + # Epoch was interrupted mid-loop, so there's no val_loss to + # weigh a "best" checkpoint against — save the in-progress + # weights as last.pt only, under the last *fully completed* + # epoch number so --resume restarts this epoch from scratch + # rather than skipping it (weights/optimizer state are still + # kept, so those partial-epoch batches aren't wasted work). + ckpt = _build_checkpoint(epoch - 1, global_step, best_val_loss) + torch.save(ckpt, out_dir / "last.pt") + last_completed_epoch = epoch - 1 + print( + f"saved in-progress weights from partway through epoch " + f"{epoch} to {out_dir / 'last.pt'} " + f"(resume will restart epoch {epoch})" + ) break train_loss = train_loss_sum / max(train_n, 1) @@ -1071,38 +1131,7 @@ def train( # must never decrease. wandb_run.log(metrics_row, step=global_step) - ckpt: dict = { - "model": stage1_model.state_dict(), - "sec_decoder": sec_decoder.state_dict(), - "optimizer": optimizer.state_dict(), - "lr_sched": lr_sched.state_dict(), - "epoch": epoch, - "best_val_loss": best_val_loss, - "global_step": global_step, - } - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - ckpt["critic"] = critic.state_dict() - ckpt["sec_critic"] = sec_critic.state_dict() - ckpt["optimizer_d"] = optimizer_d.state_dict() - if ema_decay > 0: - assert ema_stage1_model is not None and ema_sec_decoder is not None - ckpt["model_ema"] = ema_stage1_model.state_dict() - ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict() - if normalizer_dict is not None: - ckpt["normalizer"] = normalizer_dict - if pdg_map is not None: - ckpt["pdg_map"] = pdg_map - if mat_map is not None: - ckpt["mat_map"] = mat_map - if proc_map is not None: - ckpt["proc_map"] = proc_map - if model_config is not None: - ckpt["model_config"] = model_config + ckpt = _build_checkpoint(epoch, global_step, best_val_loss) if val_loss < best_val_loss: best_val_loss = val_loss diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index f0711ee..4c84c45 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -108,13 +108,13 @@ def _tiny_cfg(**train_overrides): def _run(data, out_dir, cfg=None, **kwargs): echoed: list[str] = [] + kwargs.setdefault("num_workers", 0) run_train_job( data=data, cfg=cfg or _tiny_cfg(), out_dir=out_dir, device=torch.device("cpu"), shuffle_buffer=64, - num_workers=0, echo=echoed.append, **kwargs, ) @@ -143,6 +143,22 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch): assert "normalizer: cache hit" in joined +def test_run_train_job_warns_when_num_workers_exceeds_shared_quota( + tmp_path, data, monkeypatch +): + monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2 + echo = _run(data, tmp_path / "out", num_workers=3) + assert any("num-workers=3" in m and "exceeds" in m for m in echo) + + +def test_run_train_job_no_warning_when_num_workers_within_shared_quota( + tmp_path, data, monkeypatch +): + monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2 + echo = _run(data, tmp_path / "out", num_workers=2) + assert not any("exceeds" in m for m in echo) + + def test_run_train_job_no_cache_setup_never_writes_sidecar(tmp_path, data): _run(data, tmp_path / "out", cache_setup=False) assert not setup_cache.sidecar_path(data).exists()