"""Library entry point for a KempnerForge training job.
``scripts/train.py`` is a thin CLI wrapper over :func:`run_training`: it
parses argv, loads the config, and calls in. Each phase below is a
separate function so an experiment can reuse the scaffold — or swap one
phase — without editing the shared script.
"""
from __future__ import annotations
from pathlib import Path
from typing import cast
import torch
import torch.distributed as dist
from kempnerforge.checkpoint.manager import CheckpointManager
from kempnerforge.config.job import JobConfig
from kempnerforge.config.vlm import MoTConfig
from kempnerforge.distributed.parallel import (
apply_ac,
apply_float8,
apply_fsdp2,
build_parallel_model,
default_mp_policy,
)
from kempnerforge.distributed.setup import destroy_distributed
from kempnerforge.distributed.tensor_parallel import apply_tensor_parallel
from kempnerforge.metrics.logger import get_logger
from kempnerforge.metrics.tracker import MetricsTracker
from kempnerforge.model.mot import mot_warm_start_from_text_stack
from kempnerforge.model.transformer import Transformer
from kempnerforge.model.vlm import inner_transformer
from kempnerforge.profiling.profiler import build_profiler
from kempnerforge.resilience.elastic import resolve_resume_path
from kempnerforge.resilience.health import NaNDetector
from kempnerforge.resilience.signal_handler import ShutdownHandler
from kempnerforge.training.data_pipeline import (
build_data_pipeline,
build_eval_dataloader,
build_phase_state,
)
from kempnerforge.training.freeze import (
apply_freeze_specs,
effective_freeze,
freeze_meta_at_step,
)
from kempnerforge.training.hooks import HookRunner
from kempnerforge.training.loop import (
LossFn,
StepFn,
TrainingSession,
run_training_loop,
select_step_fn,
)
from kempnerforge.training.loss import build_loss_fn
from kempnerforge.training.optimizer import build_optimizer
from kempnerforge.training.runtime import PipelineBundle, RuntimeContext, setup_distributed
from kempnerforge.training.scheduler import build_scheduler
logger = get_logger(__name__)
[docs]
def build_model(
config: JobConfig, runtime: RuntimeContext, loss_fn: LossFn
) -> tuple[torch.nn.Module, PipelineBundle | None]:
"""Build the model with the full parallelism stack applied.
With ``distributed.pp > 1`` this builds this rank's pipeline stage plus
its schedule (which is why ``loss_fn`` is needed here); otherwise it
delegates to ``build_parallel_model`` and returns no pipeline.
"""
tc = config.train
device = runtime.device
device_mesh = runtime.device_mesh
mp_policy = default_mp_policy(tc.param_dtype)
if config.distributed.pp <= 1:
model = build_parallel_model(
config.model,
device,
device_mesh,
vision_config=config.vision_encoder,
adapter_config=config.adapter,
vlm_config=config.vlm,
frames_per_clip=(config.video.max_frames if config.video is not None else 1),
ac_mode=tc.activation_checkpointing,
mp_policy=mp_policy,
param_dtype=tc.param_dtype,
compile_model=tc.compile_model,
fp8=tc.is_fp8,
)
return model, None
from kempnerforge.distributed.pipeline_parallel import (
build_pipeline_schedule,
build_pipeline_stage,
build_stage_module,
get_pp_rank,
get_pp_size,
)
if device_mesh is None:
# Unreachable: JobConfig.validate rejects pp > 1 at world_size 1, the
# only case init_distributed returns no mesh. Explicit so it survives -O.
raise RuntimeError("pipeline parallelism requires a device mesh")
pp_rank = get_pp_rank(device_mesh)
pp_size = get_pp_size(device_mesh)
tp_enabled_pp = "tp" in device_mesh.mesh_dim_names # type: ignore[reportOperatorIssue]
# apply_{float8,ac,fsdp2} are annotated for Transformer; a stage module
# exposes the same block structure, hence the cast.
if tp_enabled_pp:
# Meta-device init: same pattern as non-PP TP path.
# Avoids OOM for large PP stages that don't fit on one GPU before TP shards them.
with torch.device("meta"):
stage_mod = cast("Transformer", build_stage_module(config.model, pp_rank, pp_size))
apply_tensor_parallel(stage_mod, device_mesh)
if tc.is_fp8:
apply_float8(stage_mod)
apply_ac(stage_mod, tc.activation_checkpointing)
apply_fsdp2(stage_mod, device_mesh, mp_policy=mp_policy)
stage_mod.to_empty(device=device)
stage_mod.init_weights_and_freqs()
stage_mod.to(dtype=tc.param_dtype)
else:
stage_mod = cast(
"Transformer",
build_stage_module(config.model, pp_rank, pp_size).to(
device=device, dtype=tc.param_dtype
),
)
if tc.is_fp8:
apply_float8(stage_mod)
apply_ac(stage_mod, tc.activation_checkpointing)
apply_fsdp2(stage_mod, device_mesh, mp_policy=mp_policy)
model: torch.nn.Module = stage_mod
if tc.compile_model:
logger.info("Compiling model with torch.compile...")
model = torch.compile(model) # type: ignore[assignment]
n_params = sum(p.numel() for p in model.parameters())
logger.info(f"Model (PP stage {pp_rank}/{pp_size}): {n_params:,} parameters")
pp_stage = build_pipeline_stage(
model, # type: ignore[arg-type]
device_mesh,
device,
batch_size=tc.batch_size,
seq_len=tc.seq_len,
param_dtype=tc.param_dtype,
)
pp_schedule = build_pipeline_schedule(
stage=pp_stage,
n_microbatches=tc.grad_accum_steps,
loss_fn=loss_fn,
schedule=config.distributed.pp_schedule.value,
)
return model, PipelineBundle(rank=pp_rank, size=pp_size, schedule=pp_schedule)
[docs]
def build_checkpoint_manager(
config: JobConfig,
runtime: RuntimeContext,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
pipeline: PipelineBundle | None,
) -> CheckpointManager:
"""Create the checkpoint manager, scoping DCP to the PP stage when needed.
With PP, each stage has different parameters — DCP needs a group scoped
to ranks within the same PP stage (all non-PP mesh dimensions), and each
stage saves DCP shards to its own subdirectory to avoid file collisions.
"""
ckpt_pg = None
ckpt_pp_rank = None
device_mesh = runtime.device_mesh
if pipeline is not None and device_mesh is not None:
ckpt_pp_rank = pipeline.rank
non_pp_dims = [d for d in device_mesh.mesh_dim_names if d != "pp"] # type: ignore[reportOptionalIterable]
if len(non_pp_dims) == 1:
ckpt_pg = device_mesh[non_pp_dims[0]].get_group()
elif len(non_pp_dims) > 1:
ckpt_pg = device_mesh[tuple(non_pp_dims)].get_group()
return CheckpointManager(
config.checkpoint,
model,
optimizer,
process_group=ckpt_pg,
pp_rank=ckpt_pp_rank,
)
[docs]
def restore_checkpoint(
config: JobConfig,
model: torch.nn.Module,
scheduler: object,
ckpt_mgr: CheckpointManager,
) -> tuple[int, int]:
"""Auto-resume from ``latest`` (or ``checkpoint.load_path``).
Returns ``(step, tokens_seen)`` — ``(0, 0)`` when there is nothing to
resume from. Also restores the metrics run ids, runs the MoT warm-start
hook, and re-applies the effective freeze at the resumed step.
"""
# Every rank must resume from the same checkpoint. resolve_resume_path
# stats `.metadata` to test durability, and on NFS/Lustre independent
# stats can disagree briefly across ranks; ranks that disagree resume
# from different steps, or hang in the collective load below when some
# skip it. Use a rank-0-authoritative answer broadcast to all ranks,
# as CheckpointManager.load does for train_state.pt.
if dist.is_initialized():
chosen = [
str(resolve_resume_path(config.checkpoint.dir) or "") if dist.get_rank() == 0 else ""
]
dist.broadcast_object_list(chosen, src=0)
resume_path = Path(chosen[0]) if chosen[0] else None
else:
resume_path = resolve_resume_path(config.checkpoint.dir)
if not resume_path and not config.checkpoint.load_path:
return 0, 0
vlm_cfg = config.vlm
# On resume the expected freeze metadata reflects the post-transition
# state at the saved step (effective_freeze handles step-boundary
# transitions). Peek at the saved step via metadata.json before invoking
# load() so the comparison uses the same step the checkpoint was written at.
vlm_freeze_expected = None
if config.is_vlm:
assert vlm_cfg is not None # narrowed by is_vlm
probe_step = ckpt_mgr.peek_saved_step(str(resume_path) if resume_path else None) or 0
vlm_freeze_expected = freeze_meta_at_step(probe_step, vlm_cfg)
# A resume always restores full state; the exclusion applies only to a load_path warm start.
warm_start_exclude = (
(config.checkpoint.exclude_from_loading or None) if resume_path is None else None
)
step, tokens_seen, ckpt_extra_loaded = ckpt_mgr.load(
path=str(resume_path) if resume_path else None,
scheduler=scheduler,
exclude_keys=warm_start_exclude,
vlm_freeze_expected=vlm_freeze_expected,
)
if ckpt_extra_loaded.get("wandb_run_id"):
config.metrics.wandb_run_id = ckpt_extra_loaded["wandb_run_id"]
if ckpt_extra_loaded.get("mlflow_run_id"):
config.metrics.mlflow_run_id = ckpt_extra_loaded["mlflow_run_id"]
# MoT warm-start: translate dense TransformerBlock weights from a
# JD/text-only checkpoint into per-modality copies inside every
# MoTBlock. Runs once at the start of training (resume_path is
# None or step == 0); a real resume of an in-flight MoT run
# already has the MoT-shaped state in the checkpoint and skips
# this hook.
if isinstance(vlm_cfg, MoTConfig) and vlm_cfg.mot_warm_start_from_text and step == 0:
source = torch.load(vlm_cfg.mot_warm_start_path, map_location="cpu", weights_only=True)
if isinstance(source, dict) and "model" in source:
source = source["model"]
mot_warm_start_from_text_stack(inner_transformer(model), source)
logger.info(
f"MoT warm-start: copied dense block weights from {vlm_cfg.mot_warm_start_path}"
)
# Apply effective freeze at the resumed step so requires_grad reflects
# the post-transition state of any stages with start_step <= loaded_step.
# Build-time apply only handles the base freeze list.
if vlm_cfg is not None and vlm_cfg.freeze_schedule:
valid_modules = set(vlm_cfg.module_patterns.keys())
specs = effective_freeze(step, vlm_cfg.freeze, vlm_cfg.freeze_schedule, valid_modules)
apply_freeze_specs(model, specs, vlm_cfg.module_patterns)
logger.info(f"Resumed at step={step}; applied effective freeze ({len(specs)} specs)")
return step, tokens_seen
[docs]
def run_training(
config: JobConfig,
*,
step_fn: StepFn | None = None,
hooks: HookRunner | None = None,
) -> None:
"""Run a full training job: build every phase, run the loop, tear down.
``step_fn`` and ``hooks`` let an experiment own the step body or register
hooks without copying the build phases. Both default to what
``scripts/train.py`` uses.
"""
runtime = setup_distributed(config)
# Bound before the try so the teardown can tell "never built" from "built".
tracker: MetricsTracker | None = None
finished = False
try:
shutdown_handler = ShutdownHandler(timeout_sec=config.train.shutdown_timeout_sec)
shutdown_handler.register()
nan_detector = NaNDetector(action="warn", max_consecutive=10)
loss_fn = build_loss_fn(config.train)
model, pipeline = build_model(config, runtime, loss_fn)
optimizer = build_optimizer(model, config.optimizer)
scheduler = build_scheduler(optimizer, config.scheduler, max_steps=config.train.max_steps)
ckpt_mgr = build_checkpoint_manager(config, runtime, model, optimizer, pipeline)
step, tokens_seen = restore_checkpoint(config, model, scheduler, ckpt_mgr)
tracker = MetricsTracker(config, num_gpus=runtime.world_size)
tracker.init_backends(config)
prof = build_profiler(config.profiling, rank=runtime.rank)
data = build_data_pipeline(config, runtime)
# Apply any dataloader state stashed during load(). Runs after dataloader
# construction because the loader's identity depends on phase scheduling
# that load() restores. No-op when resuming without a prior dataloader
# state or when the loader is not stateful (plain TorchDataLoader).
if data.dataloader is not None:
ckpt_mgr.apply_dataloader_state(data.dataloader)
eval_dataloader = build_eval_dataloader(config, runtime)
phases = build_phase_state(config, data, step)
session = TrainingSession(
config=config,
runtime=runtime,
model=model,
optimizer=optimizer,
scheduler=scheduler,
loss_fn=loss_fn,
step_fn=step_fn or select_step_fn(config),
data=data,
phases=phases,
checkpointer=ckpt_mgr,
tracker=tracker,
hooks=hooks or HookRunner(),
nan_detector=nan_detector,
shutdown_handler=shutdown_handler,
pipeline=pipeline,
eval_dataloader=eval_dataloader,
profiler=prof,
)
run_training_loop(session, step=step, tokens_seen=tokens_seen)
finished = True
finally:
# As a library call this can raise and the caller can keep going, so
# the metrics run must not outlive it. close() is rank-local: backends
# are built on rank 0 only, so it is a no-op elsewhere.
if tracker is not None:
tracker.close()
# destroy_process_group() is the cooperative teardown -- it blocks in
# _wait_for_pending_works() and shuts backends down in a deliberate
# order because ncclCommAbort has been collective in some NCCL
# versions; torch ships _abort_process_group() as the separate
# error-path API. Peers are still mid-collective while we unwind, so
# on the failure path leave the group to the launcher, as main did.
if finished:
destroy_distributed()