ezpz.examples.fsdp_tp⚓︎

@ezpz.timeitlogit(rank=ezpz.get_rank())
def main(args: argparse.Namespace) -> int:
    """Entrypoint to set up distributed context and dispatch training."""
    ezpz.silence_noisy_loggers()
    t0 = time.perf_counter()
    rank = ezpz.distributed.setup_torch(tensor_parallel_size=args.tp, seed=args.seed)
    t_setup = time.perf_counter()
    base_dir = args.outdir if args.outdir else None
    # Collective (broadcasts the shared timestamp) — every rank must call it.
    outdir = get_example_outdir(WBPROJ_NAME, base_dir=base_dir)
    # Create the W&B run HERE, so it exists before the slow startup path.
    # History (inside train()) is constructed only after tokenization,
    # model build, FLOP counting, FSDP wrapping and torch.compile — on a
    # 4-node agpt-2b run that is ~60s, and at 20b it spans the OOM-prone
    # build/compile phase. Creating the run at History time meant a job
    # that died before the first step uploaded nothing at all.
    #
    # This does NOT move ownership away from History: History still builds
    # the tracker, and its WandbBackend still calls setup_wandb. It simply
    # *adopts* this run rather than creating one — setup_wandb passes
    # reinit=None and wandb.init() returns the existing run object when one
    # is live (verified on wandb 0.24.0 and 0.28.1), so config updates from
    # both sites merge onto a single run.
    #
    # Gated on the same backend resolution History uses, so
    # EZPZ_TRACKER_BACKENDS=none (or a csv/mlflow-only selection) still
    # keeps W&B entirely out of the run. WANDB_MODE=disabled/offline
    # remains the other opt-out, and setup_wandb additionally no-ops when
    # verify_wandb() fails or rank != 0 — it performs no collectives, so
    # calling it early cannot deadlock or diverge ranks.
    if rank == 0 and "wandb" in ezpz.tracker.resolve_backend_names():
        ezpz.setup_wandb(project_name=WBPROJ_NAME, dir=outdir)
        # Dumped *after* wandb.init so the resolved config lands in the
        # run's captured console log too (wandb does not capture stdout
        # retroactively).
        jstr = json.dumps(vars(args), indent=2, sort_keys=True, default=str)
        logger.info(f"config:\n{jstr}")
    logger.info("Outputs will be saved to %s", outdir)
    # W&B run created above; History (inside train()) adopts it and adds
    # the CSV/JSONL backends.
    train_start = time.perf_counter()
    # nullcontext (prof=None) unless --profile / --pyinstrument-profiler set.
    with profiling_context_from_args(args, outdir) as prof:
        # Pass main()'s t0 (captured before setup_torch) so
        # train/restart_seconds covers the full cold path incl. distributed
        # init — the dominant cost of a real --auto-retry failover.
        history = train(
            args=args, outdir=outdir, profiler=prof, process_start=t0
        )
    train_end = time.perf_counter()
    timings = {
        "main/setup_torch": t_setup - t0,
        "main/train": train_end - train_start,
        "main/total": train_end - t0,
        "timings/training_start": train_start - t0,
        "timings/train_duration": train_end - train_start,
        "timings/end-to-end": train_end - t0,
    }
    logger.info("Timings: %s", timings)
    history.tracker.log(
        {
            (f"timings/{k}" if not k.startswith("timings/") else k): v
            for k, v in timings.items()
        }
    )
    if ezpz.get_rank() == 0:
        dataset = history.finalize(
            outdir=outdir,
            run_name=WBPROJ_NAME,
            dataset_fname="train",
        )
        del dataset  # logged by finalize()
    return 0

def parallelize(
    model: nn.Module,
    device_mesh: DeviceMesh,
    mixed_precision: Optional[MixedPrecisionPolicy],
    reshard_after_forward: str = "always",
    activation_checkpoint: str = "none",
    loss_parallel: bool = False,
    meta_init: bool = False,
    device: Optional["torch.device"] = None,
) -> nn.Module:
    """Apply tensor parallelism + FSDP2 (``fully_shard``) to the model.

    FSDP2 shards each module group independently (embedding, every
    TransformerBlock, then [norm, output], then the root). This per-module
    sharding keeps the backward-pass gradient/activation memory bounded —
    in particular for the large 256K-vocab embedding and output projection
    — where FSDP1's single flat-parameter wrap would OOM at long sequence
    length. Activation checkpointing (when requested) is applied to each
    block BEFORE ``fully_shard`` so the checkpoint envelope sits inside the
    sharded unit (torchtitan's ordering).

    ``meta_init``: when True the model was built on the ``meta`` device (no
    storage). Pre-shard ``init_weights`` is skipped; after ``fully_shard`` the
    sharded params are materialized on ``device`` via ``to_empty`` and then
    ``init_weights(buffer_device=device)`` fills them — so the full dense model
    is never placed on one device (avoids the large-model build OOM). Requires
    ``device``.
    """
    tp_mesh = device_mesh["tp"]

    # Choose the mesh fully_shard shards over:
    #   dp_replicate > 1 -> 2D (dp_replicate, dp_shard) submesh; FSDP2 reads a
    #                       2D DP mesh as HSDP (replicate across the outer dim,
    #                       shard within the inner) automatically.
    #   dp_replicate == 1 -> 1D dp_shard mesh; plain FSDP sharding, identical
    #                        to the pre-HSDP behavior (avoids a needless
    #                        size-1 replicate wrap).
    if device_mesh["dp_replicate"].size() > 1:
        fsdp_dp_mesh = device_mesh["dp_replicate", "dp_shard"]
    else:
        fsdp_dp_mesh = device_mesh["dp_shard"]

    reshard = _reshard_arg(reshard_after_forward)

    # Dense path: init the real params now. Meta path: skip — the params are on
    # `meta` (no storage), so init happens after fully_shard via to_empty +
    # init_weights(buffer_device=device) below.
    if not meta_init:
        model.init_weights()  # type: ignore

    # Only apply tensor/sequence parallelism when the tp mesh dim is > 1.
    # At tp=1 (FSDP-only) the TP plan is pure overhead: SequenceParallel
    # still wraps norms as DTensors sharded over a size-1 tp dim, which
    # produces a `_NormPartial` placement that must be all-reduced — and
    # combined with FSDP2's dp sharding triggers the "2 sequential
    # all_reduce ... suboptimal" warning every step, for zero benefit
    # (there's nothing to shard across a 1-rank tp group). torchtitan
    # guards the same way (`if parallel_dims.tp_enabled`).
    if tp_mesh.size() > 1:
        model = parallelize_module(
            model,
            tp_mesh,
            {
                "tok_embeddings": RowwiseParallel(
                    input_layouts=Replicate(),
                    output_layouts=Shard(1),
                ),
                "norm": SequenceParallel(),
                # With loss_parallel, keep logits vocab-sharded (Shard(-1))
                # and return the LOCAL [N, vocab/tp] tensor so the loss can run
                # vocab-parallel CE (no full-vocab all-gather). Otherwise gather
                # to Replicate() so the loss sees full-vocab logits (default).
                "output": ColwiseParallel(
                    input_layouts=Shard(1),
                    output_layouts=Shard(-1) if loss_parallel else Replicate(),
                    use_local_output=bool(loss_parallel),
                ),
            },
        )

        assert isinstance(model.layers, Iterable)
        for _, transformer_block in enumerate(model.layers):
            layer_tp_plan = {
                "attention_norm": SequenceParallel(),
                "attention": PrepareModuleInput(
                    input_layouts=(Shard(1), None),  # type:ignore
                    desired_input_layouts=(Replicate(), None),  # type:ignore
                ),
                "attention.wq": ColwiseParallel(),
                "attention.wk": ColwiseParallel(),
                "attention.wv": ColwiseParallel(),
                "attention.wo": RowwiseParallel(output_layouts=Shard(1)),
                "ffn_norm": SequenceParallel(),
                "feed_forward": PrepareModuleInput(
                    input_layouts=(Shard(1),),
                    desired_input_layouts=(Replicate(),),
                ),
                "feed_forward.w1": ColwiseParallel(),
                "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
                "feed_forward.w3": ColwiseParallel(),
            }

            attn_layer = transformer_block.attention  # type: ignore
            attn_layer.n_heads = attn_layer.n_heads // tp_mesh.size()
            attn_layer.n_kv_heads = attn_layer.n_kv_heads // tp_mesh.size()
            parallelize_module(
                module=transformer_block,  # type: ignore
                device_mesh=tp_mesh,
                parallelize_plan=layer_tp_plan,
            )

    # Activation checkpointing must wrap each block BEFORE fully_shard so the
    # checkpoint envelope lives inside the FSDP2 unit (torchtitan ordering;
    # the reverse — AC after sharding — is the FSDP1 order and is wrong for
    # FSDP2). _apply_activation_checkpointing replaces each block in
    # `model.layers` in-place with a compile-aware CheckpointWrapper.
    if activation_checkpoint != "none":
        _apply_activation_checkpointing(model, activation_checkpoint)

    # FSDP2: shard each module group on the dp sub-mesh. Per-module sharding
    # (vs FSDP1's one flat param) is what keeps backward memory bounded.
    fsdp_kwargs = {"mesh": fsdp_dp_mesh, "reshard_after_forward": reshard}
    if mixed_precision is not None:
        fsdp_kwargs["mp_policy"] = mixed_precision

    # Embedding first (largest single param: vocab*dim).
    if getattr(model, "tok_embeddings", None) is not None:
        fully_shard(model.tok_embeddings, **fsdp_kwargs)
    # Each transformer block (or its CheckpointWrapper) as its own unit.
    assert isinstance(model.layers, Iterable)
    for block in model.layers:
        fully_shard(block, **fsdp_kwargs)
    # norm + output together (output is the other vocab*dim-sized param).
    if (
        getattr(model, "norm", None) is not None
        and getattr(model, "output", None) is not None
    ):
        fully_shard([model.norm, model.output], **fsdp_kwargs)
    # Root last.
    fully_shard(model, **fsdp_kwargs)

    # Meta path: params are now sharded DTensors still on `meta`. Materialize
    # ONLY this rank's shard on the real device (to_empty — no full-model copy),
    # then init_weights fills the sharded params and recomputes the freqs_cis
    # buffer on `device` (to_empty leaves buffer data uninitialized). On resume,
    # dcp.load later overwrites the params + persistent buffers.
    if meta_init:
        assert device is not None, "meta_init=True requires a device"
        model.to_empty(device=device)
        model.init_weights(buffer_device=device)  # type: ignore

    _configure_fsdp_gradient_division(model)

    logger.info(f"Model after parallelization (FSDP2):\n{model=}\n")
    return model

def parse_args(argv: Optional[list[str]] = None):
    """CLI parser for 2D parallel (TP/SP + FSDP) training."""
    if argv is None:
        argv = sys.argv[1:]
    parser = argparse.ArgumentParser(
        description="2D Parallel Training",
        formatter_class=DefaultsFormatter,
    )
    parser.add_argument(
        "--dim",
        type=int,
        default=256,
        help=(
            "Model hidden / embedding dimension (a.k.a. d_model). Overridden "
            "when --model selects a preset."
        ),
    )
    parser.add_argument(
        "--n-layers",
        type=int,
        default=32,
        help=(
            "Number of TransformerBlocks stacked in the model. Overridden "
            "when --model selects a preset."
        ),
    )
    parser.add_argument(
        "--n-heads",
        type=int,
        default=32,
        help=(
            "Number of attention heads per layer. Must divide --dim. "
            "Overridden when --model selects a preset."
        ),
    )
    parser.add_argument(
        "--n-kv-heads",
        type=int,
        default=4,
        help=(
            "Number of key/value heads for grouped-query attention (GQA). "
            "Must divide --n-heads. Set equal to --n-heads for standard MHA. "
            "Overridden when --model selects a preset."
        ),
    )
    parser.add_argument(
        "--multiple-of",
        type=int,
        default=360,
        help=(
            "Round the SwiGLU FFN hidden dim up to a multiple of this value "
            "(for hardware-friendly shapes). Ignored when --hidden-dim is "
            "set explicitly."
        ),
    )
    parser.add_argument(
        "--ffn-dim-multiplier",
        type=float,
        default=None,
        help=(
            "Scale factor applied to the SwiGLU FFN hidden dim before the "
            "--multiple-of rounding step. None (default) means no extra "
            "scaling; Llama2-style models use 1.3. Ignored when "
            "--hidden-dim is set explicitly."
        ),
    )
    parser.add_argument(
        "--hidden-dim",
        type=int,
        default=None,
        help=(
            "Override SwiGLU FFN hidden dim. When None (default), TransformerBlock "
            "derives it as `4 * dim` and FeedForward applies the 2/3 + "
            "ffn_dim_multiplier + multiple_of pipeline. Set this to a concrete "
            "value (e.g. 11008 for agpt-2b, 14336 for agpt-20b) to bypass the "
            "formula and hit a published architecture exactly."
        ),
    )
    parser.add_argument(
        "--rope-theta",
        type=float,
        default=10000.0,
        help=(
            "Base frequency for RoPE positional embeddings. Llama1/2 used "
            "10000 (the default); Llama3 uses 500000; agpt-2b uses 50000."
        ),
    )
    parser.add_argument(
        "--norm-eps",
        type=float,
        default=1e-5,
        help="Epsilon added to RMSNorm denominators for numerical stability.",
    )
    parser.add_argument(
        "--vocab-size",
        type=int,
        default=32_000,
        help=(
            "Tokenizer vocabulary size. Sets the embedding table and output "
            "projection sizes; must match the tokenizer used for the dataset."
        ),
    )
    parser.add_argument(
        "--lr",
        type=float,
        default=3e-3,
        help="Peak learning rate for the AdamW optimizer.",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=5,
        help="Number of passes over the training dataset.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=1,
        help=(
            "Per-DP-rank training batch size (a.k.a. micro-batch). "
            "Global batch = --batch-size * (world_size / --tp)."
        ),
    )
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        # No `choices=` — accepts both preset names (validated in
        # apply_model_preset) AND free-form HF repo IDs like
        # `meta-llama/Llama-3.2-1B`. Disambiguation is by the `/` character:
        # presence of `/` => HF repo ID; absence => preset/alias lookup.
        help=(
            "Model size preset (overrides dim/layer defaults). "
            "Presets: debug/small/medium/large/xl/xxl/xxxl/agpt-2b/agpt-20b. "
            "xl/xxl/xxxl accept long-form aliases (`xlarge`/`extra-large`, etc). "
            "agpt presets accept `agpt2b`/`agpt_2b` etc. "
            "Pass a HuggingFace repo id with a `/` (e.g. "
            "`meta-llama/Llama-3.2-1B`) to load HF weights instead — that "
            "path forces --tp 1 (FSDP-only)."
        ),
    )
    parser.add_argument(
        "--test-batch-size",
        type=int,
        default=1000,
        help=(
            "Per-DP-rank batch size for the eval/test loader. Only "
            "consumed by the MNIST data path; ignored for random and HF "
            "datasets."
        ),
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=0,
        help=(
            "Subprocess workers for the DataLoader. 0 (default) loads "
            "in-process — fine for tokenized HF datasets; bump for "
            "image pipelines or heavy on-the-fly preprocessing."
        ),
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help=(
            "Seed for torch/numpy/python RNGs (forwarded to "
            "ezpz.setup_torch). None (default) leaves the RNGs unseeded "
            "for non-deterministic runs."
        ),
    )
    parser.add_argument(
        "--tp",
        type=int,
        default=2,
        help=(
            "Tensor-parallel degree (a.k.a. TP / Megatron-style sharding). "
            "Must divide WORLD_SIZE. The remaining dimension "
            "(WORLD_SIZE / --tp) is used for FSDP data parallelism. "
            "Set to 1 for FSDP-only. Forced to 1 when --model is a HF "
            "repo id."
        ),
    )
    parser.add_argument(
        "--dp-replicate",
        type=int,
        default=1,
        help=(
            "Data-parallel REPLICATE degree (HSDP outer dim). Weights are "
            "replicated across this many groups; within each group they are "
            "sharded across --dp-shard ranks. Default 1 = no replication "
            "(pure FSDP sharding, i.e. today's behavior). Set >1 for HSDP "
            "(e.g. shard within a node, replicate across nodes). Mirrors "
            "torchtitan's data_parallel_replicate_degree. Constraint: "
            "dp_replicate * dp_shard * tp == WORLD_SIZE."
        ),
    )
    parser.add_argument(
        "--dp-shard",
        type=int,
        default=-1,
        help=(
            "Data-parallel SHARD degree (FSDP inner dim). Weights are "
            "sharded across this many ranks within each replicate group. "
            "Default -1 = 'use all remaining ranks' = "
            "WORLD_SIZE / (dp_replicate * tp), which reproduces today's "
            "flat data-parallel behavior. Mirrors torchtitan's "
            "data_parallel_shard_degree."
        ),
    )
    parser.add_argument(
        "--reshard-after-forward",
        dest="reshard_after_forward",
        nargs="?",
        const="always",
        default="always",
        choices=list(RESHARD_POLICIES),
        help=(
            "FSDP2 reshard_after_forward policy (memory vs. comm tradeoff). "
            "`always` (default, ZeRO-3): reshard params after forward — "
            "lowest memory, re-all-gathers params in backward. `never` "
            "(ZeRO-2): keep params gathered after forward — more memory, "
            "skips the backward all-gather. Bare `--reshard-after-forward` "
            "== `always`; `--no-reshard-after-forward` == `never`. For HSDP "
            "(replicate + shard) use --dp-replicate / --dp-shard."
        ),
    )
    parser.add_argument(
        "--no-reshard-after-forward",
        dest="reshard_after_forward",
        action="store_const",
        const="never",
        help="Alias for --reshard-after-forward never (ZeRO-2).",
    )
    # Deprecated legacy alias, hidden from --help. Resolved post-parse by
    # _resolve_reshard_after_forward: full_shard->always, shard_grad_op/
    # no_shard->never (with a deprecation warning), hybrid_shard*->hard error.
    parser.add_argument(
        "--sharding-strategy",
        dest="sharding_strategy",
        type=str,
        default=None,
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--activation-checkpoint",
        "--ac",
        type=str,
        default="none",
        # `full` is an alias for `block` for compatibility with
        # torchtitan's CLI surface (their `activation_checkpoint_mode`
        # uses `full` for what we call `block` — every transformer
        # block wrapped). Resolved in _apply_activation_checkpointing.
        choices=["none", "block", "full", "selective"],
        help=(
            "Activation checkpointing strategy. "
            "`none` (default) keeps all forward activations in memory. "
            "`block` (alias: `full`) wraps each TransformerBlock — typical "
            "30-40 pct activation memory reduction, ~20 pct throughput hit "
            "(matches torchtitan's default for agpt-2b/agpt-20b). "
            "`selective` checkpoints only the attention computation inside "
            "each block — ~15-20 pct memory reduction, ~10 pct throughput "
            "hit. Trade activation memory for recomputation cost — useful "
            "when OOM-ing during training (NOT during init; for init-time "
            "OOM consider increasing --tp or reducing --seq-len). "
            "NOTE: cannot be combined with --compile (upstream AOTAutograd "
            "DeviceMesh-in-saved-tensors bug — see the --compile warning). "
            "With FSDP2 you usually don't need --ac anyway; it was a "
            "workaround for the FSDP1 backward-memory OOM that FSDP2 fixes."
        ),
    )
    parser.add_argument(
        "--meta-init",
        type=str,
        default="auto",
        choices=["auto", "on", "off"],
        help=(
            "Build the native Transformer on the `meta` device, then "
            "materialize only each rank's shard after FSDP2 sharding "
            "(torchtitan pattern). Avoids the OOM from moving the full dense "
            "model onto one device before sharding, which otherwise caps model "
            "size at what fits whole on a single GPU (~2-8B) regardless of "
            "node count. `auto` (default) enables it for large native models "
            "(>= ~6B params) and keeps small models on the exact dense init "
            "path (bit-for-bit reproducible); `on` forces it for any native "
            "model; `off` forces the legacy dense path. Ignored for HF "
            "`from_pretrained` models (they load real pretrained weights). "
            "Override the auto threshold with EZPZ_META_INIT_MIN_PARAMS."
        ),
    )
    parser.add_argument(
        "--max-grad-norm",
        type=float,
        default=1.0,
        help=(
            "Clip gradients to this L2 norm before the optimizer step. "
            "Set to 0 (or negative) to disable gradient clipping."
        ),
    )
    parser.add_argument(
        "--outdir",
        type=str,
        default=None,
        help=(
            "Base directory for metrics logs + the History report. None "
            "(default) writes under the current working directory. (Model "
            "checkpoints go to --ckpt-dir, not here.)"
        ),
    )
    parser.add_argument(
        "--ckpt-dir",
        "--ckpt_dir",
        type=str,
        default=None,
        dest="ckpt_dir",
        help=(
            "Directory for DCP (sharded) checkpoints. When set, enables "
            "checkpoint save (see --save-interval) AND auto-resume: on "
            "startup the latest complete checkpoint here is loaded and "
            "training continues from it (unless --no-resume). This is what "
            "makes `ezpz launch --auto-retry` resume across attempts."
        ),
    )
    parser.add_argument(
        "--save-interval",
        "--save_interval",
        type=int,
        default=0,
        dest="save_interval",
        help=(
            "Save a checkpoint every N optimizer steps (requires "
            "--ckpt-dir). 0 (default) disables saving."
        ),
    )
    parser.add_argument(
        "--train-iters",
        "--train_iters",
        type=int,
        default=0,
        dest="train_iters",
        help=(
            "Stop after N optimizer steps, regardless of --epochs. 0 "
            "(default) runs the full --epochs pass. Step-based cap for "
            "fixed-length runs / restart-time experiments."
        ),
    )
    parser.add_argument(
        "--no-resume",
        action="store_true",
        help=(
            "Ignore any existing checkpoint in --ckpt-dir and start fresh "
            "(step 0). Default behavior auto-resumes from the latest."
        ),
    )
    parser.add_argument(
        "--async-ckpt",
        "--async_ckpt",
        action="store_true",
        dest="async_ckpt",
        help=(
            "Save checkpoints asynchronously: stage to fast node-local "
            "--ckpt-stage-dir (background thread, overlaps training), then "
            "fan out to the durable --ckpt-dir on shared FS. Requires "
            "--ckpt-dir. Resume is unchanged (always from --ckpt-dir)."
        ),
    )
    parser.add_argument(
        "--ckpt-stage-dir",
        "--ckpt_stage_dir",
        type=str,
        default=None,
        dest="ckpt_stage_dir",
        help=(
            "Node-local staging dir for --async-ckpt (default "
            "/tmp/ezpz-ckpt-<jobid>). Transient — NOT resumable on its own; "
            "only the fanned-out --ckpt-dir copy is durable."
        ),
    )
    # parser.add_argument('--dataset', type=str, default='random')
    parser.add_argument(
        "--dataset",
        type=str,
        default="eliplutchok/fineweb-small-sample",
        help=(
            "Training dataset. Special values: `mnist` (image debug "
            "dataset) and `random` (synthetic tokens, no IO). Anything "
            "else is treated as a HuggingFace dataset repo id."
        ),
    )
    parser.add_argument(
        "--tokenizer_name",
        type=str,
        default="meta-llama/llama-2-7b-hf",
        help=(
            "HuggingFace tokenizer repo id used to tokenize the HF "
            "dataset. Auto-overridden to --model when --model is a HF "
            "repo id and --tokenizer_name wasn't passed explicitly."
        ),
    )
    parser.add_argument(
        "--hf-split",
        "--hf_split",
        type=str,
        default="train",
        help="Dataset split to load.",
    )
    parser.add_argument(
        "--hf-text-column",
        "--hf_text_column",
        type=str,
        default="text",
        help="Column containing raw text in the dataset.",
    )
    parser.add_argument(
        "--hf-limit",
        "--hf_limit",
        type=int,
        default=0,
        help=(
            "Maximum number of rows to sample from the HF dataset. "
            "0 (default) = no limit (use the full dataset). Pass a "
            "positive value (e.g. `--hf-limit 512`) to subsample for "
            "smoke tests. Subsampling is deterministic given "
            "$EZPZ_HF_SAMPLE_SEED."
        ),
    )
    # parser.add_argument('--max_batch_size', type=int, default=None)
    parser.add_argument(
        "--seq-len",
        type=int,
        default=int(os.environ.get("SEQ_LEN", 1024)),
        help=(
            "Training sequence length (tokens per sample). Defaults to "
            "$SEQ_LEN if set, otherwise 1024. Must be <= --max-seq-len."
        ),
    )
    parser.add_argument(
        "--max-seq-len",
        type=int,
        default=32768,
        help=(
            "Maximum sequence length the model is built to support — "
            "sets the RoPE frequency table size and the attention "
            "scratch budget. Increase if you raise --seq-len."
        ),
    )
    parser.add_argument(
        "--fp32",
        action="store_true",
        help="Disable mixed precision (use fp32) for debugging NaNs.",
    )
    parser.add_argument(
        "--compile",
        action="store_true",
        help=(
            "Compile each TransformerBlock with torch.compile after "
            "FSDP/TP wrap (matches torchtitan's apply_compile pattern). "
            "Per-block compile dodges the Dynamo + DTensor _MaskPartial "
            "graph break that whole-model compile hits on TP-wrapped "
            "tok_embeddings, and amortizes compile cost across N layers."
        ),
    )
    parser.add_argument(
        "--compile-mode",
        type=str,
        default="default",
        choices=["default", "reduce-overhead", "max-autotune"],
        help=(
            "torch.compile mode (only used when --compile is set). "
            "`default` is safest. `reduce-overhead` enables cudagraphs "
            "for small models / large batches. `max-autotune` does "
            "extensive kernel search — slow startup, fastest steady state."
        ),
    )
    parser.add_argument(
        "--act-mem-budget",
        type=float,
        default=1.0,
        help=(
            "Activation-memory budget for the inductor min-cut partitioner "
            "(sets torch._functorch.config.activation_memory_budget). Only "
            "takes effect with --compile. 1.0 (default) saves ALL "
            "activations (no recompute); lower values let the compiler "
            "recompute activations in backward to cut peak memory — e.g. "
            "0.5 keeps ~half. This is how torchtitan fits larger batches "
            "for the same model (its MemoryBudgetAC sets 0.5). Try 0.5 if "
            "you OOM in backward at a batch size that should fit."
        ),
    )
    parser.add_argument(
        "--loss-impl",
        type=str,
        default="eager",
        choices=[
            "eager",
            "chunked",
            "chunked-backward",
            "compiled",
            "loss-parallel",
            "fused-linear",
        ],
        help=(
            "Cross-entropy implementation. The large-vocab output path is the "
            "memory bottleneck: a full (B*T, vocab) fp32 logits tensor + its "
            "grad (agpt-2b 256K vocab, seq=8192, bs=2: ~16.8 GiB EACH) can OOM "
            "a GPU tile (UR_RESULT_ERROR_OUT_OF_RESOURCES) even when the model "
            "fits. Pick by what you need (numbers = measured agpt-2b tp=1):\n"
            "  • `eager` (default): plain F.cross_entropy on full logits. "
            "Simplest; OOMs at agpt-2b bs2/seq8192. Use for small vocab/seq.\n"
            "  • `chunked`: chunks only the FORWARD (--loss-chunk-size). Does "
            "NOT bound backward; still OOMs at large vocab. Rarely useful.\n"
            "  • `chunked-backward`: custom autograd Function that also bounds "
            "the backward graph (recomputes each chunk's grad), saving ~one "
            "full logits buffer vs eager. General + model-agnostic (works for "
            "HF models, no torch.compile needed) — good at MODERATE vocab/seq "
            "or when compile is unavailable. Still holds two logit-sized "
            "buffers, so it does NOT fix the very-large-vocab OOM (use "
            "fused-linear/compiled there).\n"
            "  • `compiled`: torch.compile fuses log_softmax+NLL+backward so "
            "the full transient is never materialized (torchtitan's approach). "
            "Fits (~45 GB) and is the FASTEST that fits (~28%% MFU). Needs "
            "working torch.compile. Best default when it fits.\n"
            "  • `fused-linear` (Liger/Cut-CE): runs the output projection "
            "per row-chunk so the full (B*T,vocab) logits/grad are NEVER built "
            "— bounds BOTH row and vocab dims. LOWEST memory (~32 GB, below "
            "compiled) at ~24%% MFU; trades a little speed for headroom (bigger "
            "batch/seq). ezpz Transformer + tp=1 only (HF / tp>1 fall back to "
            "compiled).\n"
            "  • `loss-parallel`: vocab-parallel CE sharding the vocab across "
            "TP ranks (each holds vocab/tp) via TP all-reduces. Bounds the "
            "VOCAB dim; only helps at tp>1 (at tp=1 falls back to eager). At "
            "tp>1 it is also the only correct path (plain CE hits a "
            "Tensor/DTensor mismatch on Replicate logits). ~23 GB/rank, "
            "~34%% MFU at tp=2.\n"
            "NOTE: `--compile` only compiles the transformer blocks, NOT the "
            "loss, so it does NOT by itself fix the loss transient — use "
            "--loss-impl for that."
        ),
    )
    parser.add_argument(
        "--loss-chunk-size",
        type=int,
        default=1024,
        help=(
            "Row-chunk size (number of (B*T) token rows per cross-entropy "
            "chunk) for --loss-impl=chunked, chunked-backward, and "
            "fused-linear. Smaller = lower peak memory, more kernel launches. "
            "Ignored for eager/compiled/loss-parallel."
        ),
    )
    # max_batch_size: int = 32
    # max_seq_len: int = 32768
    # Shared profiler flags (--profile / --pyinstrument-profiler / etc.),
    # consumed by profiling_context_from_args around the training loop.
    add_profiling_args(parser)
    args = parser.parse_args(argv)
    apply_model_preset(args, argv)
    # Fold the deprecated --sharding-strategy alias into reshard_after_forward
    # (and hard-error the removed hybrid_shard* values).
    _resolve_reshard_after_forward(args)
    _maybe_enable_cpu_backend_for_async_ckpt(args)
    return args

@ezpz.timeitlogit(rank=ezpz.get_rank())
def train(
    args: argparse.Namespace,
    outdir: Path | str | os.PathLike,
    profiler: Optional[Any] = None,
    process_start: Optional[float] = None,
) -> int:
    """Run TP/SP + FSDP training and optionally log metrics.

    Args:
        args: Parsed CLI namespace.
        outdir: Output directory for metrics / reports.
        profiler: Optional active profiler (``torch.profiler.profile`` or
            ``None``) from :func:`ezpz.profile.profiling_context_from_args`.
            When non-None, ``profiler.step()`` is called once per training
            step so the schedule (wait/warmup/active/repeat) advances.
    """
    # Timestamp for the restart-time metric. Prefer the caller's
    # process-start (main() captures it BEFORE setup_torch, so
    # train/restart_seconds includes distributed init — the dominant cost of a
    # real cold failover). Fall back to now() when train() is called directly.
    _train_t0 = process_start if process_start is not None else perf_counter()
    _restart_logged = False
    _optim_dtype_logged = False  # log AdamW state dtype once, after step 1
    _precision_in_jsonl = False  # fold precision summary into JSONL once
    world_size = ezpz.distributed.get_world_size()
    assert world_size % args.tp == 0, "WORLD_SIZE must be divisible by TP"
    # Resolve the data-parallel topology: dp_replicate (HSDP outer) x
    # dp_shard (FSDP inner). Defaults (replicate=1, shard=-1) reproduce the
    # flat FSDP dp dim exactly.
    dp_replicate, dp_shard = _resolve_dp_degrees(
        world_size=world_size,
        tp=args.tp,
        dp_replicate=args.dp_replicate,
        dp_shard=args.dp_shard,
    )
    dpsize = dp_replicate * dp_shard  # == flattened "dp" size
    # Global batch size = per-DP-rank micro-batch x data-parallel degree.
    # Data parallelism is dp_replicate * dp_shard (== dpsize): BOTH the HSDP
    # replicate dim and the FSDP shard dim get a distinct data slice (the
    # DistributedSampler shards over num_replicas=dpsize). TP is NOT data
    # parallel — tp ranks process the SAME batch (sharded across the hidden
    # dim), so tp does not enter here. Backfilled onto args so it lands in the
    # wandb run config (History logs vars(args); we also push it to
    # wandb.config below). Computed here, not in parse_args, because the
    # default --dp-shard -1 ("use all remaining ranks") only resolves to a
    # concrete degree once WORLD_SIZE is known.
    #
    # This is the general Megatron/NeMo formula
    #   gbs = micro_batch * world_size * grad_accum
    #         / (tensor_parallel * pipeline_parallel)
    # specialized to this example: there is no pipeline parallelism
    # (pipeline_parallel = 1) and no gradient accumulation (grad_accum = 1 —
    # zero_grad/backward/step run every iteration), and world_size / tp ==
    # dp_replicate * dp_shard == dpsize. If either is ever added, this must
    # gain the corresponding factor (* grad_accum, / pipeline_parallel).
    args.global_batch_size = args.batch_size * dpsize
    if ezpz.get_rank() == 0:
        logger.info(
            "global_batch_size=%d (batch_size=%d x dp_replicate=%d x "
            "dp_shard=%d; tp=%d does not scale the batch)",
            args.global_batch_size,
            args.batch_size,
            dp_replicate,
            dp_shard,
            args.tp,
        )
    # fused-linear (Liger/Cut-CE): needs the model's hidden states + output
    # weight, so only the ezpz Transformer (not HF models) and — for now —
    # only tp=1 (tp>1 needs vocab-shard composition, gated to a follow-up).
    # Resolved against the actual model after it's built (see below).
    want_fused_linear = getattr(args, "loss_impl", "eager") == "fused-linear"
    # loss-parallel (vocab-sharded CE) only does anything at tp>1; at tp=1 the
    # local vocab IS the full vocab. NOTE: `use_loss_parallel` is resolved
    # LATER (after the HF branch may force args.tp=1), so an HF model launched
    # with --tp>1 --loss-impl=loss-parallel doesn't enter vocab-parallel CE
    # with a stale TP group. See the resolution just before parallelize().
    # 3D mesh: (dp_replicate, dp_shard, tp). We then flatten the two DP dims
    # into a single named "dp" dim so all existing dp consumers
    # (DistributedSampler num_replicas/rank, loss dp_group) keep working with
    # correct rank arithmetic. fully_shard picks the DP submesh explicitly
    # (2D -> HSDP when dp_replicate > 1, else 1D dp_shard == today's FSDP).
    device_mesh = ezpz.init_device_mesh_safe(
        str(ezpz.get_torch_device()),
        (dp_replicate, dp_shard, args.tp),
        mesh_dim_names=("dp_replicate", "dp_shard", "tp"),
    )
    # Flatten the two DP dims into a single "dp" dim. Must go through
    # flatten_device_mesh_safe: `_flatten` builds the flattened PG via
    # split_group when the default PG is device-bound, which xccl (XPU)
    # doesn't support — the same workaround init_device_mesh_safe applies.
    ezpz.distributed.flatten_device_mesh_safe(
        device_mesh[("dp_replicate", "dp_shard")], "dp"
    )
    logger.info(f"Device mesh created:\n{device_mesh=}")

    hf_dataset = None
    hf_tokenizer = None
    if args.dataset.lower() not in {"mnist", "random"}:
        from ezpz.data.hf import get_hf_text_dataset

        seed = int(os.environ.get("EZPZ_HF_SAMPLE_SEED", "1337"))
        hf_dataset, hf_tokenizer = get_hf_text_dataset(
            dataset_name=args.dataset,
            split=args.hf_split,
            text_column=args.hf_text_column,
            tokenizer_name=args.tokenizer_name,
            seq_len=args.seq_len,
            limit=args.hf_limit,
            seed=seed,
        )
        if hf_tokenizer.vocab_size != args.vocab_size:
            logger.warning(
                "Overriding vocab_size from %s to tokenizer vocab_size=%s",
                args.vocab_size,
                hf_tokenizer.vocab_size,
            )
            args.vocab_size = hf_tokenizer.vocab_size

    # HF repo IDs are forced to the HF code path; the ezpz Transformer
    # construction below would silently produce a randomly-initialized model
    # with the wrong architecture for the requested repo.
    is_hf_model = bool(args.model and "/" in args.model)

    config = ModelArgs(
        dim=args.dim,
        n_layers=args.n_layers,
        n_heads=args.n_heads,
        n_kv_heads=args.n_kv_heads,
        batch_size=args.batch_size,
        vocab_size=args.vocab_size,
        multiple_of=args.multiple_of,
        hidden_dim=args.hidden_dim,
        rope_theta=args.rope_theta,
        ffn_dim_multiplier=args.ffn_dim_multiplier,
        norm_eps=args.norm_eps,
        max_seq_len=args.max_seq_len,
    )
    logger.info(f"config:\n{config}")
    metrics_every = int(os.environ.get("EZPZ_METRICS_EVERY", "1"))
    track_logits = os.environ.get("EZPZ_TRACK_LOGITS", "0") == "1"
    track_hist = os.environ.get("EZPZ_TRACK_HIST", "0") == "1"
    track_act_hist = os.environ.get("EZPZ_TRACK_ACT_HIST", "1") == "1"
    hist_bins = int(os.environ.get("EZPZ_HIST_BINS", "64"))
    hist_samples = int(os.environ.get("EZPZ_HIST_SAMPLES", "20000"))
    dataset_tag = args.dataset.lower().replace("/", "_")
    # Update wandb config with model args (run already initialised in main).
    # NOTE: `config` (ModelArgs) is not mutated later, so it's safe to push
    # here. The resolved CLI `args`, however, ARE still mutated below (the HF
    # branch forces args.tp=1; args.loss_impl can be normalized to
    # "compiled"), so vars(args) is pushed just before the training loop
    # instead (see below) — logging it here would record the requested
    # settings rather than the effective ones actually used.
    if (
        ezpz.get_rank() == 0
        and wandb is not None
        and getattr(wandb, "run", None) is not None
    ):
        from dataclasses import asdict

        wandb.config.update(asdict(config))  # type:ignore

    device_type = ezpz.distributed.get_torch_device_type()
    device = (
        torch.device("cpu")
        if device_type == "cpu"
        else torch.device(f"{device_type}:{ezpz.get_local_rank()}")
    )
    # Decide meta-device init (native large models only) before building, so the
    # native build below can go on `meta` and never place the full dense model
    # on one device. HF models always resolve to False (they load real weights).
    meta_init = _resolve_meta_init(args, config, is_hf_model)
    if is_hf_model:
        # HF path: pull arch + weights from the hub. The ezpz Transformer
        # above is skipped entirely. Note we still built `config` above so
        # downstream logging / wandb.config.update(asdict(config)) doesn't
        # crash, but it does NOT reflect the real HF architecture — that's
        # in `model.config` after the load below.
        from transformers import AutoModelForCausalLM

        if args.tp > 1:
            logger.warning(
                "HF model %s requested with --tp=%d; ezpz's TP plan is "
                "hardcoded to its own Transformer module names and won't "
                "match HF's LlamaDecoderLayer / GemmaDecoderLayer / ... "
                "Forcing --tp 1 (FSDP-only).",
                args.model,
                args.tp,
            )
            args.tp = 1
        hf_dtype = torch.float32 if args.fp32 else torch.bfloat16
        hf_token = os.environ.get("HF_TOKEN") or os.environ.get(
            "HUGGING_FACE_HUB_TOKEN"
        )
        logger.info(
            "Loading HF model %s (dtype=%s)%s",
            args.model,
            hf_dtype,
            " with HF_TOKEN" if hf_token else "",
        )
        model = AutoModelForCausalLM.from_pretrained(
            args.model,
            torch_dtype=hf_dtype,
            token=hf_token,
        )
    else:
        if meta_init:
            # Build on `meta`: no storage allocated, so even a 20B model costs
            # nothing here. parallelize() shards it, then to_empty materializes
            # only this rank's shard on the real device.
            with torch.device("meta"):
                model = Transformer.from_model_args(config)
        else:
            model = Transformer.from_model_args(config)
    mstr = summarize_model(
        model,
        verbose=False,
        depth=2,
    )
    logger.info(f"\n{mstr}")
    # Meta models are materialized (sharded) inside parallelize() via to_empty;
    # moving a meta model with .to(device) would NOT allocate real storage, so
    # skip it here. Dense + HF paths place the real model on-device now.
    if not meta_init:
        model.to(device)

    # FLOPs estimation: try the exact fake-tensor path first, fall back
    # to the linear-scaling probe if it fails.
    #
    # FAKE-TENSOR PATH (preferred): runs the forward+backward at the
    # real (batch, seq) shape under FakeTensorMode (shape-only tensors,
    # no allocations → no OOM) with sdpa_kernel(MATH) forced so SDPA
    # decomposes into bmms that FlopCounterMode can see. Exact count,
    # attention included.
    #
    # LINEAR-SCALING PROBE (fallback): runs at (1, 128) with real
    # tensors and scales by token ratio. Exact for O(seq·dim) MLP/proj
    # ops, but UNDER-COUNTS attention because the O(seq²·dim) Q·Kᵀ and
    # attn·V matmuls don't scale linearly. Worse, on CPU and on fused
    # SDPA backends (flash / efficient / cuDNN), FlopCounterMode often
    # reports zero for the SDPA op entirely — so both probe and actual
    # silently drop attention from the count. Reported MFU is then a
    # lower bound: real utilization is at least the printed number,
    # often significantly higher on long-seq runs.
    _model_flops = try_estimate_fake(
        model, (args.batch_size, args.seq_len)
    )
    if _model_flops > 0:
        logger.info(
            "FLOPs counted exactly via FakeTensorMode at shape "
            "(batch=%d, seq=%d): %.3e (includes attention).",
            args.batch_size,
            args.seq_len,
            _model_flops,
        )
    elif meta_init:
        # The real-tensor probe below would run an actual forward, which a
        # meta model (no storage) cannot do. Skip it — MFU/TFLOPS just stay 0
        # for meta-init runs when the fake-tensor count is unavailable.
        if ezpz.get_rank() == 0:
            logger.warning(
                "Fake-tensor FLOP estimate returned 0 and model is on `meta` "
                "(--meta-init); skipping the real-tensor probe. train/tflops "
                "and train/mfu will be 0 for this run."
            )
    else:
        _flops_probe_batch = 1
        _flops_probe_seq = min(128, args.seq_len)
        _flops_probe = try_estimate(
            model, (_flops_probe_batch, _flops_probe_seq)
        )
        _actual_tokens = args.batch_size * args.seq_len
        _probe_tokens = _flops_probe_batch * _flops_probe_seq
        _model_flops = int(
            _flops_probe * _actual_tokens / max(_probe_tokens, 1)
        )
        if args.seq_len > _flops_probe_seq:
            logger.warning(
                "Fake-tensor FLOP estimate failed; falling back to "
                "linear-scaling probe (probe seq=%d -> actual seq=%d). "
                "This under-counts O(seq^2) attention by ~%dx; reported "
                "MFU is a lower bound (real utilization is at least this "
                "high).",
                _flops_probe_seq,
                args.seq_len,
                args.seq_len // _flops_probe_seq,
            )

    # FSDP2 mixed-precision policy (param in bf16, reduce in fp32). None when
    # --fp32 is set (pure fp32 for NaN debugging).
    mp_config: Optional[MixedPrecisionPolicy] = None
    _reduce_dtype = torch.float32  # default (also the effective dtype under --fp32)
    if not args.fp32:
        # reduce_dtype: fp32 gradient reduce-scatter is more accurate, but for
        # a large-vocab output projection (e.g. agpt's 256K) the single
        # reduce-scatter tensor can exceed CCL's ~2GB-per-message MPI limit
        # (256K*2048*4B = 2.1GB) → `atl_mpi !req.is_completed`. Set
        # EZPZ_REDUCE_DTYPE=bf16 to halve the collective size (1.05GB) and
        # stay under the limit. Validated against an explicit set: a typo
        # silently falling back to fp32 would re-trigger the very CCL
        # failure this escape hatch exists to avoid, so raise instead.
        _reduce_dtype_env = os.environ.get("EZPZ_REDUCE_DTYPE", "fp32")
        _reduce_dtype_key = _reduce_dtype_env.lower()
        if _reduce_dtype_key == "fp32":
            _reduce_dtype = torch.float32
        elif _reduce_dtype_key in ("bf16", "bfloat16"):
            _reduce_dtype = torch.bfloat16
        else:
            raise ValueError(
                f"Invalid EZPZ_REDUCE_DTYPE={_reduce_dtype_env!r}. Expected "
                "one of: 'fp32', 'bf16', 'bfloat16' (case-insensitive)."
            )
        mp_config = MixedPrecisionPolicy(
            param_dtype=torch.bfloat16,
            reduce_dtype=_reduce_dtype,
        )
    # Resolve loss-parallel NOW that args.tp is final (the HF branch above may
    # have forced tp=1). loss-parallel only does anything at tp>1; at tp=1 the
    # local vocab IS the full vocab, so normalize to compiled CE rather than
    # leaving loss_impl='loss-parallel' to hit an unhandled impl at the call
    # site. This also covers the HF + --tp>1 --loss-impl=loss-parallel case:
    # tp is now 1, so we fall back instead of entering vocab-parallel CE with a
    # stale TP group.
    use_loss_parallel = (
        getattr(args, "loss_impl", "eager") == "loss-parallel" and args.tp > 1
    )
    if getattr(args, "loss_impl", "eager") == "loss-parallel" and not use_loss_parallel:
        logger.warning(
            "--loss-impl=loss-parallel requires tp>1 (got tp=%d); falling back "
            "to compiled CE.",
            args.tp,
        )
        args.loss_impl = "compiled"
    if is_hf_model:
        # HF path: FSDP2-only wrap (no TP — the TP plan is ezpz-specific).
        # Apply activation checkpointing first (HF models use their own
        # gradient_checkpointing_enable inside _apply_activation_checkpointing),
        # then fully_shard each decoder block + the root.
        if args.activation_checkpoint != "none":
            _apply_activation_checkpointing(model, args.activation_checkpoint)

        # Find the decoder block stack: the SINGLE deepest non-empty
        # ModuleList (e.g. `model.model.layers`). Collecting every ModuleList
        # would over-shard MoE/multimodal models; the deepest one is reliably
        # the decoder stack (matches _find_block_list / HF's _no_split_modules).
        deepest_modlist: Optional[torch.nn.ModuleList] = None
        deepest_depth = -1
        deepest_len = -1
        for name, module in model.named_modules():
            if (
                isinstance(module, torch.nn.ModuleList)
                and len(module) > 0
            ):
                depth = name.count(".")
                if depth > deepest_depth or (
                    depth == deepest_depth and len(module) > deepest_len
                ):
                    deepest_depth = depth
                    deepest_len = len(module)
                    deepest_modlist = module

        # Same HSDP mesh selection as parallelize(): 2D DP submesh when
        # replicating, else the 1D shard mesh (identical to pre-HSDP).
        if device_mesh["dp_replicate"].size() > 1:
            hf_dp_mesh = device_mesh["dp_replicate", "dp_shard"]
        else:
            hf_dp_mesh = device_mesh["dp_shard"]
        hf_fsdp_kwargs = {
            "mesh": hf_dp_mesh,
            "reshard_after_forward": _reshard_arg(args.reshard_after_forward),
        }
        if mp_config is not None:
            hf_fsdp_kwargs["mp_policy"] = mp_config
        if deepest_modlist is not None:
            for block in deepest_modlist:
                fully_shard(block, **hf_fsdp_kwargs)
        else:
            logger.warning(
                "HF model: no decoder ModuleList found; sharding only the "
                "root module (per-layer memory savings will be reduced)."
            )
        fully_shard(model, **hf_fsdp_kwargs)
        _configure_fsdp_gradient_division(model)
    else:
        # TP + FSDP2. parallelize() applies activation checkpointing per
        # block BEFORE fully_shard (correct FSDP2 ordering), so we do NOT
        # re-apply it afterwards.
        # loss-parallel needs vocab-sharded (local) logits out of the output
        # projection; only meaningful at tp>1 (at tp=1 there's nothing to
        # shard, so it falls back to eager — see the loss call site).
        model = parallelize(
            model,
            device_mesh,
            mp_config,
            reshard_after_forward=args.reshard_after_forward,
            activation_checkpoint=args.activation_checkpoint,
            loss_parallel=use_loss_parallel,
            meta_init=meta_init,
            device=device,
        )
    if args.compile:
        # Activation-memory budget for the inductor min-cut partitioner.
        # Default 1.0 = save every activation (no recompute); < 1.0 lets the
        # compiler recompute a fraction of activations in backward to cut
        # peak memory. This is the knob torchtitan's MemoryBudgetAC sets
        # (0.5) — it's why TT fits a larger batch than this example for the
        # same model. Global config, applies to every compiled block below.
        if args.act_mem_budget != 1.0:
            import torch._functorch.config as _functorch_config

            _functorch_config.activation_memory_budget = args.act_mem_budget
            logger.info(
                "Set activation_memory_budget=%.3f (inductor will recompute "
                "activations in backward to cut peak memory).",
                args.act_mem_budget,
            )
        if args.activation_checkpoint != "none":
            # --ac + --compile together trip an upstream AOTAutograd bug:
            #   AssertionError: expected all tensors_saved_with_vc_check to
            #   be Tensors, got [... DeviceMesh]
            # The non-reentrant checkpoint_wrapper saves a DeviceMesh into
            # the autograd graph, which the compiled-backward saved-tensors
            # check rejects. Under FSDP2 every sharded module carries a
            # DeviceMesh, so this fires even at --tp 1 (with FSDP1 it
            # required --tp > 1). Repro + triage:
            # torchtitan/.../docs/upstream-issues/repro_devicemesh_in_saved_tensors.py
            # Not fixable here — drop one of --ac / --compile. (With FSDP2
            # you typically no longer need --ac for memory; it was a
            # workaround for the FSDP1 OOM that FSDP2 already resolves.)
            logger.warning(
                "--compile + --activation-checkpoint=%s will likely crash "
                "with an AOTAutograd 'tensors_saved_with_vc_check ... "
                "DeviceMesh' assertion (upstream bug; fires under FSDP2 even "
                "at --tp 1). Drop one of --ac / --compile. Note FSDP2 usually "
                "removes the need for --ac (it fixed the FSDP1 OOM).",
                args.activation_checkpoint,
            )
        # Compile each TransformerBlock individually rather than the whole
        # model. This is what torchtitan does (apply_compile in
        # torchtitan/models/.../infra/parallelize.py) and it dodges the
        # Dynamo + DTensor _MaskPartial graph break that whole-model
        # compile hits on TP-wrapped tok_embeddings:
        #
        #   RuntimeError when making fake tensor call: call_method
        #   redistribute(...) on DTensor(_MaskPartial(...))
        #
        # The embedding's RowwiseParallel output_fn does a redistribute
        # from _MaskPartial → Shard(1), which Dynamo can't trace under
        # fake tensors. Excluding the embedding from compile (and only
        # compiling the blocks) keeps the speedup where it matters
        # (attention + MLP, the repeated structure) without exposing
        # Dynamo to the TP output transform. Bonus: compile cost is paid
        # once for one block and reused across N layers, not N times.
        #
        # Find the block list: ezpz Transformer has `.layers`, HF
        # decoder-only models nest it as `.model.layers`.
        block_container = None
        if hasattr(model, "layers"):
            block_container = model.layers
        elif hasattr(model, "model") and hasattr(model.model, "layers"):
            block_container = model.model.layers
        if block_container is None:
            logger.warning(
                "Could not find a TransformerBlock list (model.layers or "
                "model.model.layers) — falling back to whole-model "
                "torch.compile, which may hit DTensor graph breaks."
            )
            model = torch.compile(model, mode=args.compile_mode)
        else:
            logger.info(
                "Compiling each TransformerBlock with torch.compile"
                "(mode=%s, fullgraph=True) — %d blocks.",
                args.compile_mode,
                len(block_container),
            )
            for layer_id, block in block_container.named_children():
                compiled = torch.compile(
                    block, mode=args.compile_mode, fullgraph=True
                )
                block_container.register_module(layer_id, compiled)
    base_model = model
    if not hasattr(base_model, "layers"):
        base_model = getattr(model, "_fsdp_wrapped_module", model)
    # Resolve fused-linear eligibility now that the model exists. Requires the
    # ezpz Transformer (has `.output` weight + return_hidden) and tp=1 for now.
    use_fused_linear = False
    if want_fused_linear:
        has_output = hasattr(base_model, "output") and hasattr(
            base_model.output, "weight"
        )
        if is_hf_model or not has_output:
            logger.warning(
                "--loss-impl=fused-linear needs the ezpz Transformer "
                "(hidden states + output weight); falling back to compiled "
                "for this model."
            )
            # Normalize so the `_compute_loss` call site actually runs compiled
            # CE — leaving loss_impl='fused-linear' would fall through to an
            # unhandled impl (now an error; previously silent eager + OOM).
            args.loss_impl = "compiled"
        elif args.tp > 1:
            logger.warning(
                "--loss-impl=fused-linear with tp>1 is not yet supported "
                "(needs vocab-shard composition); falling back to compiled."
            )
            args.loss_impl = "compiled"
        else:
            use_fused_linear = True
    act_activations: dict[str, torch.Tensor] = {}
    act_handles: list[torch.utils.hooks.RemovableHandle] = []
    if track_hist and track_act_hist and ezpz.get_rank() == 0 and not is_hf_model:
        # `_register_activation_hooks` indexes into `model.layers[i]`, which
        # is ezpz-Transformer specific. HF models nest blocks under
        # `model.model.layers` (Llama/Mistral) or `model.gpt_neox.layers`
        # (GPT-NeoX) etc., so the hook registration would key-error. Skip
        # the hooks for HF runs; the rest of the metrics still work.
        hist_layers_spec = os.environ.get(
            "EZPZ_HIST_LAYERS", f"0,{config.n_layers - 1}"
        )
        layer_ids = _parse_hist_layers(hist_layers_spec, config.n_layers)
        act_activations, act_handles = _register_activation_hooks(
            base_model, layer_ids
        )
    logger.info(f"Creating optimizer=AdamW with lr={args.lr}")

    # Prefer the fused AdamW kernel (single kernel for the whole param
    # update) — it's what torchtitan uses and it's measurably faster than
    # `foreach` on XPU. Fall back to foreach if fused isn't supported for
    # this build/device (older torch, CPU, etc.).
    try:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=args.lr,
            betas=(0.9, 0.95),
            eps=1e-8,
            weight_decay=0.1,
            fused=True,
        )
    except (RuntimeError, ValueError) as exc:
        logger.warning(
            "Fused AdamW unavailable (%s); falling back to foreach=True.", exc
        )
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=args.lr,
            betas=(0.9, 0.95),
            eps=1e-8,
            weight_decay=0.1,
            foreach=True,
        )

    # Log the ACTUAL precision of each component (introspects the live model +
    # MP policy, not the intended config), push it to the W&B run config, and
    # keep the dict to (a) fold into the first metrics row and (b) amend with
    # the optimizer-state dtype after the first step (on a FRESH run the state
    # dict is empty until then; a RESUMED run has it populated by
    # load_checkpoint, so reading it post-step-1 is correct either way).
    _precision_summary = _log_precision_summary(model, mp_config, _reduce_dtype, args)

    # --- Resume from checkpoint (auto-detect latest) --------------------------
    # Both model AND optimizer exist and are fully sharded here, so this is the
    # correct point to restore sharded DCP state into them. Auto-resume (no
    # flag) is what lets `ezpz launch --auto-retry` — which relaunches the
    # IDENTICAL command each attempt — pick up where a failed attempt left off.
    # Async checkpointing needs a durable target to fan out to; a node-local
    # stage dir alone is not resumable after a failure (see _checkpoint.py).
    if args.async_ckpt and not args.ckpt_dir:
        raise ValueError("--async-ckpt requires --ckpt-dir (the durable target)")
    ckpt_stage_dir = args.ckpt_stage_dir
    if args.async_ckpt and not ckpt_stage_dir:
        jobid = (
            os.environ.get("PBS_JOBID")
            or os.environ.get("SLURM_JOB_ID")
            or str(os.getpid())
        ).split(".")[0]
        ckpt_stage_dir = f"/tmp/ezpz-ckpt-{jobid}"
    # In-flight async checkpoint handle (drained before the next save + at exit).
    pending_ckpt = None

    resume_meta: "dict[str, object] | None" = None
    if args.ckpt_dir and not args.no_resume:
        from ezpz.examples._checkpoint import load_checkpoint

        resume_meta = load_checkpoint(args.ckpt_dir, model, optimizer)
        if resume_meta is not None and ezpz.get_rank() == 0:
            logger.info(
                "RESUMED from step=%s (epoch=%s, batch_offset=%s, "
                "tokens_seen=%s)",
                resume_meta.get("step"),
                resume_meta.get("epoch"),
                resume_meta.get("batch_offset"),
                resume_meta.get("tokens_seen"),
            )

    # reuse device for input placement

    tp_group = device_mesh.get_group("tp")
    if args.dataset.lower() == "mnist":
        data_prefix = Path(os.getcwd()).joinpath(
            ".cache", "ezpz", "data", f"{args.dataset.lower()}"
        )
        from ezpz.data.vision import get_mnist
        from ezpz.data.distributed import TPBroadcastDataLoader

        data = get_mnist(
            outdir=Path(data_prefix),
            train_batch_size=args.batch_size,
            test_batch_size=args.test_batch_size,
            num_replicas=dpsize,
            rank=device_mesh.get_local_rank("dp"),
            pin_memory=True,
            num_workers=args.num_workers,
        )
        dataset = data["dataset"]
        sampler = data["sampler"]
        dataloader = data["dataloader"]
        if args.tp > 1:
            dataloader = TPBroadcastDataLoader(dataloader, tp_group)
    elif args.dataset.lower() == "random":
        from ezpz.data.distributed import get_random_dataset_fsdp_tp

        data = get_random_dataset_fsdp_tp(
            batch_size=args.batch_size,
            vocab_size=args.vocab_size,
            seq_length=args.seq_len,
            dp_group=device_mesh.get_group("dp"),
            tp_group=tp_group,
            broadcast_within_tp=True,
            drop_last=True,
        )
        dataset = data["dataset"]
        sampler = data["sampler"]
        dataloader = data["dataloader"]
    # if args.dataset.lower() != "random":
    else:
        from ezpz.data.distributed import TPBroadcastDataLoader

        assert hf_dataset is not None
        dataset = hf_dataset
        # drop_last=True (both sampler + loader) so every batch has a static
        # batch dim — a ragged tail triggers a torch.compile recompile at the
        # epoch boundary that OOMs the compiled CE. See _build_hf_dataloader.
        sampler, dataloader = _build_hf_dataloader(
            dataset,
            batch_size=args.batch_size,
            dpsize=dpsize,
            dp_rank=device_mesh.get_local_rank("dp"),
            world_size=ezpz.get_world_size(),
        )
        if args.tp > 1:
            dataloader = TPBroadcastDataLoader(dataloader, tp_group)

    # ezpz.breakpoint(0)

    logger.info("Starting 2D training...")
    model.train()

    # outdir = Path(args.outdir).joinpath(ezpz.utils.get_timestamp())
    metrics_path = Path(outdir).joinpath(
        f"metrics-{ezpz.distributed.get_rank()}.jsonl"
    )
    Path(outdir).mkdir(parents=True, exist_ok=True)
    history = ezpz.history.History(
        project_name=WBPROJ_NAME,
        config={"args": vars(args), **ezpz.get_dist_info()},
        outdir=outdir,
        report_dir=outdir,
        report_enabled=True,
        jsonl_path=metrics_path,
        jsonl_overwrite=True,
        # Disable cross-rank history aggregation while profiling (either
        # profiler) — the all-gather of per-rank metrics perturbs the very
        # step times the profiler is measuring.
        distributed_history=(
            1 < world_size <= 384
            and not getattr(args, "pytorch_profiler", False)
            and not getattr(args, "pyinstrument_profiler", False)
        ),
    )

    # Re-push the precision summary now that History has created the W&B run.
    # The initial push (at optimizer-build time, above) runs BEFORE this and so
    # no-ops on `wandb.run is None`; without this second push a run that dies
    # before its first optimizer step would leave W&B with no precision
    # diagnostics at all. Cheap + idempotent (config update, allow_val_change).
    _push_precision_to_wandb(_precision_summary)

    # For TP, input needs to be the same across all TP ranks.
    # while for SP, input can be different across all ranks
    # We will use dp_rank for setting the random seed
    # to mimic the behavior of the dataloader
    # x = torch.tensor((args.batch_size, args.seq_len))
    x = torch.tensor(0)
    global_step = 0
    # Cumulative count of training tokens consumed across the whole run
    # (summed global tokens/step). Logged as train/tokens_seen — the standard
    # x-axis for loss-vs-tokens curves. See the metrics block: global
    # tokens/step = batch * full_seq_len * dpsize (full pre-shard seq length,
    # not the SP-local shard, so it's exact and rank-invariant).
    tokens_seen = 0
    # Checkpoint timings from the PREVIOUS step, folded into the next step's
    # metrics dict so they reach JSONL + W&B:
    #   - train/ckpt_save_seconds  : sync save's blocking write (all 23 GB to
    #     the durable dir on the training thread).
    #   - train/ckpt_stage_seconds : async save's CPU-stage stall only (the
    #     cheap part — copy state to host, kick off the background write).
    #   - train/ckpt_drain_seconds : async fan-out (/tmp -> shared FS) blocking
    #     time at the START of the next step. This is the EXPENSIVE half of an
    #     async save and was previously untimed — it lands between steps, so it
    #     is captured by neither ckpt_stage_seconds nor train/dt. The honest
    #     per-save stall for async is stage + drain, NOT stage alone.
    pending_stage_seconds: "float | None" = None
    pending_save_seconds: "float | None" = None
    pending_drain_seconds: "float | None" = None
    # Resume bookkeeping. When resuming, seed the counters from the checkpoint
    # and reconstruct (start_epoch, resume_offset) so the loop skips
    # already-consumed batches. drop_last=True on both sampler and loader makes
    # batches-per-epoch deterministic, so global_step -> (epoch, offset) is
    # exact. batches_per_epoch may be 0 for iterable/unsized loaders — guard it.
    start_epoch = 0
    resume_offset = 0
    if resume_meta is not None:
        global_step = int(resume_meta.get("step", 0) or 0)
        tokens_seen = int(resume_meta.get("tokens_seen", 0) or 0)
        try:
            batches_per_epoch = len(dataloader)
        except TypeError:
            batches_per_epoch = 0
        if batches_per_epoch > 0:
            start_epoch = global_step // batches_per_epoch
            resume_offset = global_step % batches_per_epoch
        else:
            # Unsized loader: fall back to the saved epoch/offset if present.
            start_epoch = int(resume_meta.get("epoch", 0) or 0)
            resume_offset = int(resume_meta.get("batch_offset", 0) or 0)
        # Time-to-first-post-resume-step: measured against the training entry
        # timestamp captured at the top of train() (see _train_t0), logged once
        # on the first completed step below as train/restart_seconds.
        _resumed = True
    else:
        _resumed = False
    # Push the RESOLVED CLI args to the wandb run config now — after every
    # args mutation (HF tp-force, loss_impl normalization) and after the
    # global_batch_size backfill — so the logged config reflects the settings
    # actually used for the rest of training, not the requested ones.
    if (
        ezpz.get_rank() == 0
        and wandb is not None
        and getattr(wandb, "run", None) is not None
    ):
        # allow_val_change since some keys may already be present from main().
        wandb.config.update(vars(args), allow_val_change=True)  # type:ignore
    for epoch in range(start_epoch, args.epochs):
        if sampler is not None:
            sampler.set_epoch(epoch)
        for idx, batch in enumerate(dataloader):
            # Skip already-consumed batches in the resumed epoch only.
            if _resumed and epoch == start_epoch and idx < resume_offset:
                continue
            # Step-based stop (independent of --epochs).
            if args.train_iters and global_step >= args.train_iters:
                break
            # Finalize the backgrounded async fan-out AS SOON AS every rank's
            # copy is done — not deferred to the next save boundary. Each step
            # this cheaply votes (a torch all-reduce on the main thread, in
            # lockstep across ranks — the same footing as the barriers already
            # here, no cross-thread hazard) whether all ranks finished copying;
            # when they have, it stamps the durable .complete marker. This keeps
            # the saved-but-not-yet-durable window to ~copy-duration (~seconds)
            # instead of a full save interval, so a crash falls back at most ~1
            # interval like a sync save. The guard is rank-uniform (pending_ckpt
            # is set/cleared identically on every rank), so all ranks enter the
            # collective together. None-safe / no-op until ready.
            if args.async_ckpt and pending_ckpt is not None:
                from ezpz.examples._checkpoint import try_finalize_if_ready

                _fin_t0 = perf_counter()
                if try_finalize_if_ready(pending_ckpt) is not None:
                    pending_drain_seconds = perf_counter() - _fin_t0
                    pending_ckpt = None
            ezpz.distributed.synchronize()
            t0 = perf_counter()
            attn_mask = None
            if isinstance(batch, dict) and "input_ids" in batch:
                x = batch["input_ids"]
                attn_mask = batch.get("attention_mask")
            else:
                x = batch
            assert isinstance(x, torch.Tensor)
            x = x.to(device)
            x = x.to(torch.long)
            if args.dataset == "random":
                inp = x[:, :-1]
                labels = x[:, 1:]
            else:
                inp = x[:, :-1]
                labels = x[:, 1:]
            inp = inp.to(device)
            labels = labels.to(device)
            if attn_mask is not None:
                attn_mask = attn_mask.to(device)
            # fused-linear gets hidden states (B,T,dim) instead of logits, so
            # the loss can form logits in chunks and never materialize the full
            # (B,T,vocab) tensor. Otherwise the model returns logits as usual.
            if use_fused_linear:
                pred = model(inp, return_hidden=True)
            else:
                pred = model(inp)
            # HF causal-LM models return a CausalLMOutput dataclass with a
            # `.logits` tensor; ezpz's Transformer returns logits directly.
            if hasattr(pred, "logits"):
                pred = pred.logits
            # pred is (B,T,vocab) logits, or (B,T,dim) hidden under
            # fused-linear; either way dim-1 is the (SP-local) seq length.
            local_seq_len = pred.shape[1]
            if labels.shape[1] != local_seq_len:
                labels = _slice_for_sequence_parallel(labels, local_seq_len)
            if attn_mask is not None:
                if attn_mask.shape[1] > 1:
                    attn_labels = attn_mask[:, 1:]
                else:
                    attn_labels = attn_mask
                if attn_labels.shape[1] != local_seq_len:
                    attn_labels = _slice_for_sequence_parallel(
                        attn_labels, local_seq_len
                    )
            # Build a single ignore-mask (attention-pad OR tokenizer-pad) and
            # apply it with ONE masked_fill, instead of two .clone()s + two
            # boolean index-assigns. Each .clone() + labels[mask]=-100 was a
            # separate aten::copy_ per step; this collapses them to one copy.
            # -100 can't collide with a valid pad_id (>=0), so mask order is
            # irrelevant. (Profiling: agpt-2b aten::copy_ was ~8.8% of step.)
            pad_id = getattr(dataset, "pad_id", None)
            ignore_mask = None
            if attn_mask is not None:
                ignore_mask = attn_labels == 0
            if pad_id is not None:
                pad_mask = labels == int(pad_id)
                ignore_mask = (
                    pad_mask if ignore_mask is None else (ignore_mask | pad_mask)
                )
            if ignore_mask is not None:
                labels = labels.masked_fill(ignore_mask, -100)
            ezpz.distributed.synchronize()
            t1 = perf_counter()
            tp_mod = getattr(ezpz, "tp", None)
            tp_rank = (
                getattr(tp_mod, "get_tensor_parallel_rank", lambda: 0)()
                if tp_mod is not None
                else 0
            )
            # First-step finite/max debug stats. Gated behind
            # EZPZ_TRACK_LOGITS because `torch.isfinite(pred)` allocates a
            # full `(B, T, vocab)`-shaped bool tensor on the un-reduced
            # logits — at agpt's 256K vocab and long seq that's multiple GB
            # materialized *before* the loss, which can OOM a run that would
            # otherwise fit (the loss itself may be chunked/compiled to stay
            # bounded, but this debug probe is not). Off by default.
            # Skip the logits probe under fused-linear: `pred` is hidden
            # states there, not logits, so finite/max-abs of it is meaningless
            # (and the full logits are never materialized by design).
            if track_logits and not use_fused_linear and epoch == 0 and idx == 0:
                pred_finite = torch.isfinite(pred)
                pred_nonfinite = int((~pred_finite).sum().item())
                pred_max = float(pred.abs().max().item())
                logger.info(
                    "pred_stats rank=%s tp=%s shape=%s nonfinite=%s max_abs=%s",
                    ezpz.get_rank(),
                    tp_rank,
                    tuple(pred.shape),
                    pred_nonfinite,
                    f"{pred_max:.6f}",
                )
            if use_fused_linear:
                # `pred` is hidden states (B,T,dim); the fused loss runs the
                # output projection MODULE per row-chunk (so FSDP unshards the
                # weight + routes its grad) and never materializes the full
                # (B,T,vocab) logits or its grad.
                loss = _cross_entropy_fused_linear(
                    pred,
                    base_model.output,
                    labels,
                    ignore_index=-100,
                    chunk_size=args.loss_chunk_size,
                )
            elif use_loss_parallel:
                # `pred` is this rank's local [B, T, vocab/tp] shard
                # (output projection has use_local_output=True under
                # loss-parallel). Vocab-parallel CE reduces across the TP
                # group; global vocab is needed to compute shard bounds.
                loss = _cross_entropy_vocab_parallel(
                    pred,
                    labels,
                    ignore_index=-100,
                    global_vocab_size=args.vocab_size,
                    tp_group=tp_group,
                )
            else:
                # tp>1 non-loss-parallel: `pred` is a REPLICATED DTensor
                # (output ColwiseParallel output_layouts=Replicate,
                # use_local_output=False) but `labels` is plain, so plain CE
                # would raise "mixed torch.Tensor and DTensor". Localize to a
                # plain tensor first (no-op at tp=1/HF). See
                # _localize_logits_for_loss for the full rationale + guard.
                loss = _compute_loss(
                    _localize_logits_for_loss(pred),
                    labels,
                    impl=args.loss_impl,
                    ignore_index=-100,
                    chunk_size=args.loss_chunk_size,
                )
            if epoch == 0 and idx == 0:
                valid_labels = int((labels != -100).sum().item())
                logger.info(
                    "loss_inputs rank=%s tp=%s local_seq_len=%s labels=%s valid_labels=%s",
                    ezpz.get_rank(),
                    tp_rank,
                    local_seq_len,
                    tuple(labels.shape),
                    valid_labels,
                )
                # loss = F.cross_entropy(
                #     pred.flatten(0, 1),
                #     labels.flatten(0, 1),
                # )
                # loss = output.loss
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            grad_norm_preclip = None
            if args.max_grad_norm > 0:
                grad_norm_preclip = torch.nn.utils.clip_grad_norm_(
                    model.parameters(), args.max_grad_norm
                )
            optimizer.step()
            # AdamW state (exp_avg/exp_avg_sq) is allocated by the first step on
            # a fresh run (a resumed run already has it from load_checkpoint);
            # log its dtype once, now that it exists (also amends the summary +
            # W&B config with the optimizer-state dtype).
            if not _optim_dtype_logged:
                _log_optimizer_state_dtype(optimizer, _precision_summary)
                _optim_dtype_logged = True
            ezpz.distributed.synchronize()
            t2 = perf_counter()
            global_step += 1
            # Advance the torch.profiler schedule once per optimizer step.
            # No-op when not profiling (profiler is None).
            if profiler is not None:
                profiler.step()
            metrics: dict[str, object] = {
                "train/iter": global_step,
                "train/epoch": epoch,
                "train/bidx": idx,
                "train/loss": loss.item(),
                "train/dt": t2 - t0,
                "train/dtf": t1 - t0,
                "train/dtb": t2 - t1,
            }
            if grad_norm_preclip is not None:
                metrics["grad/norm_preclip"] = float(grad_norm_preclip)
            if global_step % max(metrics_every, 1) == 0:
                metrics.update(_collect_param_grad_stats(model, device))
                metrics["opt/iter"] = (global_step,)
                metrics["opt/lr"] = float(optimizer.param_groups[0]["lr"])
                metrics["input/iter"] = (global_step,)
                metrics["input/max"] = float(x.max().item())
                metrics["input/min"] = float(x.min().item())
                metrics["labels/valid"] = float((labels != -100).sum().item())
                if track_logits:
                    pred_finite = torch.isfinite(pred)
                    metrics["logits/nonfinite"] = float(
                        (~pred_finite).sum().item()
                    )
                    metrics["logits/max_abs"] = float(pred.abs().max().item())
                if track_hist and ezpz.get_rank() == 0:
                    logits_sample = _sample_tensor_values(pred, hist_samples)
                    if logits_sample is not None:
                        logits_hist = _histogram_dict(logits_sample, hist_bins)
                        if logits_hist is not None:
                            metrics[f"hist/{dataset_tag}/logits"] = logits_hist
                    layer_grad_norms = _collect_layer_grad_norms(base_model)
                    if layer_grad_norms:
                        layer_grad_hist = _histogram_dict(
                            torch.tensor(layer_grad_norms), hist_bins
                        )
                        if layer_grad_hist is not None:
                            metrics[
                                f"hist/{dataset_tag}/grad_norm_per_layer"
                            ] = layer_grad_hist
                    if track_act_hist and act_activations:
                        for act_key, act_tensor in act_activations.items():
                            act_sample = _sample_tensor_values(
                                act_tensor, hist_samples
                            )
                            act_hist = _histogram_dict(act_sample, hist_bins)
                            if act_hist is not None:
                                metrics[
                                    f"hist/{dataset_tag}/activations/{act_key}"
                                ] = act_hist
                    if history.tracker.get_backend("wandb") is not None:
                        _wandb_log_histograms(
                            metrics, step=global_step, enabled=track_hist
                        )
            # Reuse the train/dt we already computed above so the MFU
            # denominator can never silently drift from the reported
            # step time.
            dt_step = float(metrics["train/dt"])  # type: ignore[arg-type]
            if _model_flops > 0 and dt_step > 0:
                # Per-DEVICE TFLOPS / MFU. `_model_flops` is counted on the
                # FULL, un-sharded model (estimated before parallelize() applies
                # TP), so it is the work of ONE data-parallel group — done
                # COLLECTIVELY by the group's `tp` GPUs. Divide by args.tp so
                # each metric reflects a single GPU's share; otherwise per-GPU
                # TFLOPS/MFU over-count by exactly `tp` (2x at tp=2, 4x at tp=4;
                # at tp=1 this is a no-op). Divide by tp, NOT world_size: unlike
                # the GLOBAL token count (tps_per_gpu ÷ world_size), _model_flops
                # is already per-DP-group, so only the tp ranks sharing that one
                # model must be divided out — the dpsize groups each do this
                # full-model work independently.
                flops_per_gpu = _model_flops / args.tp
                metrics["train/tflops"] = flops_per_gpu / dt_step / 1e12
                metrics["train/mfu"] = compute_mfu(flops_per_gpu, dt_step)
            # Throughput.
            #   - train/tps         : global tokens/sec across all GPUs
            #   - train/tps_per_gpu : per-GPU tokens/sec (torchtitan's `tgs`)
            #
            # Global tokens/step = batch * FULL pre-shard seq len (inp.shape[1])
            # * dpsize. `inp` is Replicate() across the tp group (the TP plan
            # shards only the embedding OUTPUT, never the input), so inp.shape[1]
            # is the full sequence, identical on every rank — exact and
            # rank-invariant even though only rank 0 logs. tp does NOT enter the
            # GLOBAL count: the tp ranks hold the SAME sequence (full-length
            # logits under the default Replicate() output, or Shard(1) slices
            # summing to inp.shape[1] under fused-linear / loss-parallel), never
            # distinct sequences. On the HF path (tp forced to 1, no SP) the
            # tp-dim ranks see DUPLICATE samples, so multiplying by dpsize (not
            # world_size) counts each distinct token once.
            # Global tokens processed THIS step across all distinct-data ranks.
            tokens_this_step = args.batch_size * inp.shape[1] * dpsize
            tokens_seen += tokens_this_step
            # Cumulative consumed training tokens — the standard x-axis for
            # loss curves. Accumulated every step (this block runs each step),
            # so it's exact regardless of the metrics-logging interval.
            metrics["train/tokens"] = tokens_this_step
            metrics["train/tokens_seen"] = tokens_seen
            if dt_step > 0:
                # Per-GPU throughput = global tokens / (actual GPU count) / dt.
                # Divide by world_size, NOT dpsize: under TP the `tp` GPUs in a
                # data-parallel group process that group's tokens TOGETHER, so
                # the per-GPU rate is tp× lower. (tokens_per_rank / dt would be
                # per-DP-group, over-counting per-GPU by `tp`; at tp=1
                # world_size==dpsize so this is unchanged.)
                metrics["train/tps_per_gpu"] = tokens_this_step / world_size / dt_step
                metrics["train/tps"] = tokens_this_step / dt_step
            # Async-ckpt staging time from the previous step's save, carried
            # here so it lands in the JSONL + W&B alongside the other train/*
            # metrics (set once, then cleared).
            if pending_stage_seconds is not None:
                metrics["train/ckpt_stage_seconds"] = pending_stage_seconds
                pending_stage_seconds = None
            if pending_save_seconds is not None:
                metrics["train/ckpt_save_seconds"] = pending_save_seconds
                pending_save_seconds = None
            if pending_drain_seconds is not None:
                metrics["train/ckpt_drain_seconds"] = pending_drain_seconds
                pending_drain_seconds = None
            # Restart time: on the FIRST completed step after a resume, log how
            # long from train() entry (process init + dist setup + model build
            # + dcp.load + this step) to a productive step. This is the full
            # cold path a real --auto-retry failover pays. Logged once.
            if _resumed and not _restart_logged:
                metrics["train/restart_seconds"] = perf_counter() - _train_t0
                _restart_logged = True
            # Device memory: empty on CPU/MPS, 4 keys on CUDA/XPU.
            metrics |= ezpz.get_memory_metrics(prefix="train/")
            history.update(metrics, summarize=False)
            # Write the precision summary as its OWN JSONL record, once.
            # Deliberately NOT merged into `metrics`: History is a NUMERIC
            # store (it coerces values with float() and builds an xarray
            # Dataset), so string dtype names there break get_dataset() with
            # "ValueError: too many dimensions 'str'" — taking the end-of-run
            # report/plots down with it. _write_jsonl_entry is string-safe and
            # independent of the numeric store. Emitted after step 1 so
            # optimizer_states is populated.
            if not _precision_in_jsonl:
                try:
                    history._write_jsonl_entry(
                        {f"precision/{k}": v for k, v in _precision_summary.items()}
                    )
                except Exception:  # noqa: BLE001 — logging must not break training
                    logger.debug("could not write precision JSONL entry", exc_info=True)
                _precision_in_jsonl = True
            history.log_metrics(
                metrics,
                logger=logger,
                debug_prefixes=("hist/", "grad/", "input/", "labels/", "param/"),
                include_summary=True,
                rank0_only_summary=True,
            )
            if epoch == 0 and idx == 0:
                logger.info(f"{x.shape}")
            # Save a checkpoint every --save-interval optimizer steps.
            if (
                args.ckpt_dir
                and args.save_interval
                and global_step % args.save_interval == 0
            ):
                _ckpt_meta = {
                    "tokens_seen": tokens_seen,
                    "epoch": epoch,
                    # Next batch to run on resume within `epoch`.
                    "batch_offset": idx + 1,
                }
                if args.async_ckpt:
                    from ezpz.examples._checkpoint import (
                        finalize_fanout,
                        save_checkpoint_async,
                        start_fanout,
                    )

                    # Only one async save may be in flight (DCP constraint), so
                    # ensure the previous one is finalized before starting a new
                    # one. Normally the per-step try_finalize_if_ready already
                    # finalized it (pending_ckpt is None here); this is the
                    # fallback for when the save interval is SHORTER than the
                    # copy — then finalize_fanout blocks on the still-running
                    # copy. Only record the drain time when we actually finalized
                    # here (else keep the per-step poll's measurement). None-safe.
                    if pending_ckpt is not None:
                        _drain_t0 = perf_counter()
                        finalize_fanout(pending_ckpt)
                        pending_drain_seconds = perf_counter() - _drain_t0
                        pending_ckpt = None
                    _t_stage = perf_counter()
                    pending_ckpt = save_checkpoint_async(
                        args.ckpt_dir,
                        ckpt_stage_dir,
                        global_step,
                        model,
                        optimizer,
                        meta=_ckpt_meta,
                    )
                    # Caller-thread stall = staging time only. Stash it so the
                    # NEXT step's metrics dict carries train/ckpt_stage_seconds
                    # through history.update -> JSONL + W&B (this save block runs
                    # after the current step's metrics were already written).
                    _stage_s = perf_counter() - _t_stage
                    pending_stage_seconds = _stage_s
                    # Kick the /tmp -> shared-FS fan-out onto a background
                    # thread; it overlaps the next save interval of training and
                    # is finalized at the next save boundary (above).
                    start_fanout(pending_ckpt)
                    if ezpz.get_rank() == 0:
                        logger.info(
                            "train/ckpt_stage_seconds=%.4f (async stage @ "
                            "step %d)",
                            _stage_s,
                            global_step,
                        )
                else:
                    from ezpz.examples._checkpoint import save_checkpoint

                    # Synchronous save BLOCKS the training loop for the full
                    # write. Time it (train/ckpt_save_seconds) so the sync-vs-
                    # async trade-off is measurable — this is the stall async
                    # removes. Folded into the next step's metrics dict.
                    _t_save = perf_counter()
                    save_checkpoint(
                        args.ckpt_dir,
                        global_step,
                        model,
                        optimizer,
                        meta=_ckpt_meta,
                    )
                    pending_stage_seconds = None  # (async-only; keep clear)
                    _save_s = perf_counter() - _t_save
                    pending_save_seconds = _save_s
                    if ezpz.get_rank() == 0:
                        logger.info(
                            "train/ckpt_save_seconds=%.4f (sync save @ step %d)",
                            _save_s,
                            global_step,
                        )
        # Step-based stop breaks the inner loop; also break the epoch loop.
        if args.train_iters and global_step >= args.train_iters:
            break
    # Finish any in-flight async checkpoint so a run that ends right after a
    # save doesn't lose it (the background fan-out may still be running), then
    # tear down the fan-out worker. drain() joins the background copy if one is
    # in flight, else runs it inline; either way it does the barrier + marker.
    if args.async_ckpt:
        from ezpz.examples._checkpoint import drain, shutdown_fanout_pool

        drain(pending_ckpt)
        shutdown_fanout_pool()
    if act_handles:
        for handle in act_handles:
            handle.remove()
    ezpz.distributed.barrier()
    logger.info("Finished 2D training")
    return history