"""Top-level job configuration aggregating all sub-configs."""
from __future__ import annotations
from dataclasses import dataclass, field
from kempnerforge.config.adapter import AdapterConfig
from kempnerforge.config.checkpoint import CheckpointConfig
from kempnerforge.config.data import DataConfig
from kempnerforge.config.distributed import DistributedConfig
from kempnerforge.config.eval import EvalConfig
from kempnerforge.config.metrics import MetricsConfig
from kempnerforge.config.model import FLEX_BLOCK_SIZE, ModelConfig
from kempnerforge.config.optimizer import OptimizerConfig
from kempnerforge.config.profiling import ProfilingConfig
from kempnerforge.config.scheduler import SchedulerConfig
from kempnerforge.config.training import TrainConfig
from kempnerforge.config.video import VideoConfig
from kempnerforge.config.vision import VisionEncoderConfig
from kempnerforge.config.vlm import MoMaConfig, MoTConfig, VLMConfig
_HF_VISION_ENCODER_TYPES = frozenset({"siglip2", "clip"})
def _arch_of(vlm: VLMConfig) -> str:
"""The arch of ``vlm``'s class, which is what the parallelism checks act on.
The instance's ``arch`` field can name another arch, for example after a
``--vlm.arch`` override, which keeps the class the TOML section built.
"""
return type(vlm).arch
[docs]
@dataclass
class JobConfig:
"""Top-level configuration aggregating all sub-configs.
The VLM stack lives as three sibling top-level sections,
``[vision_encoder]`` / ``[adapter]`` / ``[vlm]``. A pure text run
leaves all three at ``None`` and writes only ``[model]`` (plus the
other operational sections). A VLM run sets ``[vlm]`` and is
required to provide ``[vision_encoder]``; ``[adapter]`` defaults to
the registered 2-layer MLP and may be omitted.
"""
model: ModelConfig = field(default_factory=ModelConfig)
train: TrainConfig = field(default_factory=TrainConfig)
optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
data: DataConfig = field(default_factory=DataConfig)
eval: EvalConfig = field(default_factory=EvalConfig)
distributed: DistributedConfig = field(default_factory=DistributedConfig)
checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
metrics: MetricsConfig = field(default_factory=MetricsConfig)
profiling: ProfilingConfig = field(default_factory=ProfilingConfig)
vision_encoder: VisionEncoderConfig | None = None
adapter: AdapterConfig | None = None
vlm: VLMConfig | None = None
video: VideoConfig | None = None
plugins: list[str] = field(default_factory=list)
def __post_init__(self) -> None:
"""Cross-section invariants that fire at construction time.
``validate(world_size)`` covers checks that depend on the
distributed world size; everything that is decidable from the
config alone runs here so loader-time mistakes surface before
any distributed init.
"""
if not isinstance(self.plugins, list) or not all(isinstance(p, str) for p in self.plugins):
raise ValueError(f"plugins must be a list of module paths (got {self.plugins!r})")
if self.video is not None:
if self.vlm is None:
raise ValueError(
"[video] is set but [vlm] is missing; video training runs through "
"the VLM wrapper, so a [vlm] section (and [vision_encoder]) is required."
)
if self.vlm.tile_grid:
raise ValueError(
"vlm.tile_grid tiles the image path, but [video] is set; a video run "
"samples frames instead, so leave vlm.tile_grid at 0"
)
if self.vlm is not None:
if self.vision_encoder is None:
raise ValueError(
"[vlm] is set but [vision_encoder] is missing; a VLM run requires a "
"vision encoder section."
)
if self.adapter is None:
self.adapter = AdapterConfig()
if self.vision_encoder.type in _HF_VISION_ENCODER_TYPES and (
self.vision_encoder.feature_dim > 0 or self.vision_encoder.num_tokens > 0
):
import logging
overrides = []
if self.vision_encoder.feature_dim > 0:
overrides.append(f"feature_dim={self.vision_encoder.feature_dim}")
if self.vision_encoder.num_tokens > 0:
overrides.append(f"num_tokens={self.vision_encoder.num_tokens}")
logging.getLogger(__name__).warning(
"vision_encoder.type=%r is HF-backed and probes its own dims; "
"setting %s overrides the probe and may desync from the loaded "
"model. Leave both at 0 to use the probed values.",
self.vision_encoder.type,
" and ".join(overrides),
)
if self.vision_encoder.num_tokens > 0:
assert self.adapter is not None
per_frame = self.adapter.output_num_tokens(self.vision_encoder.num_tokens)
frames = self.views_per_sample
visual_tokens = frames * per_frame
residual_image_tokens = self.vlm.residual_stream_image_tokens(visual_tokens)
required = residual_image_tokens + self.vlm.max_text_len
if self.model.max_seq_len < required:
raise ValueError(
f"model.max_seq_len ({self.model.max_seq_len}) insufficient for VLM: "
f"residual_image_tokens ({residual_image_tokens}) + "
f"vlm.max_text_len ({self.vlm.max_text_len}) = {required}"
)
if self.data.pack_sequences and self.vlm is not None:
if self.video is not None:
raise ValueError("data.pack_sequences is not supported with [video]")
if self.vlm.arch != "joint_decoder":
raise ValueError(
"data.pack_sequences on a VLM run requires vlm.arch='joint_decoder' "
f"(got {self.vlm.arch!r}); the other arches place image tokens as a "
"prefix only"
)
if self.train.is_fp8:
raise ValueError(
"data.pack_sequences on a VLM run does not support "
"train.mixed_precision='fp8': the packed head runs over the batch's text "
"tokens, whose count is the inner dimension of the head's FP8 "
"weight-gradient GEMM and must be a multiple of 16"
)
if self.optimizer.lr_multipliers:
if self.vlm is None:
raise ValueError(
"optimizer.lr_multipliers names modules of vlm.module_patterns, so it needs "
"a [vlm] section"
)
unknown = sorted(set(self.optimizer.lr_multipliers) - set(self.vlm.module_patterns))
if unknown:
raise ValueError(
f"optimizer.lr_multipliers names unknown modules {unknown}; the modules of "
f"vlm.module_patterns are {sorted(self.vlm.module_patterns)}"
)
if self.model.attention_backend == "flex" and self.train.seq_len < FLEX_BLOCK_SIZE:
raise ValueError(
f"attention_backend='flex' requires train.seq_len >= {FLEX_BLOCK_SIZE}, "
f"got {self.train.seq_len}. Use attention_backend='sdpa' for shorter sequences."
)
@property
def is_vlm(self) -> bool:
"""Whether this job builds a ``VLMWrapper`` around the text backbone."""
return self.vlm is not None
@property
def is_video(self) -> bool:
"""Whether this job trains on a video dataset (a VLM sub-mode)."""
return self.video is not None
@property
def views_per_sample(self) -> int:
"""Images the encoder sees per sample: a clip's frames, a tiled image's views, or one."""
if self.video is not None:
return self.video.max_frames
if self.vlm is not None:
return self.vlm.views_per_image
return 1
[docs]
def validate(self, world_size: int = 1) -> None:
"""Run cross-config validations that depend on the world size."""
self.distributed.validate_world_size(world_size)
if self.train.seq_len > self.model.max_seq_len:
raise ValueError(
f"train.seq_len ({self.train.seq_len}) exceeds "
f"model.max_seq_len ({self.model.max_seq_len})"
)
if self.model.tie_embeddings and self.distributed.pp > 1:
raise ValueError(
"Tied embeddings are not supported with pipeline parallelism "
"(embedding and output head must be on different stages)"
)
if self.eval.enabled and not self.eval.dataset_path and not self.eval.hf_dataset_name:
raise ValueError(
"eval.enabled is True but no eval data source is configured. "
"Set eval.dataset_path or eval.hf_dataset_name."
)
if (
self.eval.enabled
and self.distributed.pp > 1
and self.eval.steps % self.train.grad_accum_steps
):
raise ValueError(
f"eval.steps ({self.eval.steps}) must be a multiple of train.grad_accum_steps "
f"({self.train.grad_accum_steps}) under pipeline parallelism: eval runs its "
f"batches through the training schedule, grad_accum_steps at a time."
)
if self.train.fp32_master_weights and self.model.is_moe:
resolved = self.distributed.resolve(world_size)
if max(resolved.dp_replicate, resolved.dp_shard, resolved.tp, resolved.ep) == 1:
raise ValueError(
"train.fp32_master_weights with an MoE model needs a mesh FSDP2 wraps "
"(more than one GPU): without one the forward runs under torch.autocast, "
"the expert inputs stay fp32 and the experts lose grouped GEMM. Turn "
"fp32_master_weights off or train on more than one GPU."
)
if self.model.is_moe and self.train.compile_model:
import logging
logging.getLogger(__name__).warning(
"torch.compile is not yet optimized for MoE dispatch (data-dependent shapes "
"cause graph breaks). Set compile_model=false for MoE models."
)
if self.distributed.tp > 1:
if isinstance(self.vlm, MoTConfig):
raise ValueError(
f"vlm.arch={_arch_of(self.vlm)!r} does not support tensor parallelism "
f"(distributed.tp={self.distributed.tp}): the tensor-parallel plan is keyed "
"on TransformerBlock's module names, and MoTBlock renames the attention "
"modules (attn_norm, attn) and holds mlp_norm and mlp as per-modality "
"ModuleDicts, so the build fails or no parameter is sharded. Set "
"distributed.tp = 1."
)
if (
self.vlm is not None
and self.distributed.pp == 1
and self.model.tp_sequence_parallel
):
ways = (
"tied embeddings give (MoMa keeps model.num_experts = 0)"
if isinstance(self.vlm, MoMaConfig)
else "an MoE model or tied embeddings give"
)
raise ValueError(
f"vlm.arch={_arch_of(self.vlm)!r} does not support tensor parallelism with "
f"sequence parallelism on (distributed.tp={self.distributed.tp}): tensor "
"parallelism shards the activations along the sequence for a dense model "
"with untied embeddings, and the first forward then fails on mixed "
"torch.Tensor and DTensor operands. Set distributed.tp = 1. This arch "
"runs under tensor parallelism only with sequence parallelism off, which "
f"{ways}."
)
if self.model.n_heads % self.distributed.tp != 0:
raise ValueError(
f"n_heads ({self.model.n_heads}) must be divisible by "
f"tp ({self.distributed.tp})"
)
if self.model.n_kv_heads and self.model.n_kv_heads % self.distributed.tp != 0:
raise ValueError(
f"n_kv_heads ({self.model.n_kv_heads}) must be divisible by "
f"tp ({self.distributed.tp})"
)
if self.train.is_fp8 and self.distributed.tp > 1:
raise ValueError(
"FP8 + Tensor Parallelism is not yet supported (torchao Float8Linear "
"does not compose with DTensor sharding). Use FP8 with FSDP only, "
"or TP without FP8."
)
if (
self.model.is_moe
and self.model.moe_router == "sigmoid_topk"
and self.train.activation_checkpointing == "full"
):
raise ValueError(
"model.moe_router = 'sigmoid_topk' does not support "
"train.activation_checkpointing = 'full': the sigmoid router updates its expert "
"bias in place in every training forward, so the backward's replay of a "
"checkpointed block routes with a moved bias, can pick other experts than the "
"forward did, and then raises CheckpointError. Use activation_checkpointing = "
"'selective', which checkpoints attention only and leaves the MoE layers out "
"of the replay."
)
if self.model.is_moe and self.distributed.pp > 1:
alternatives = (
"Use FSDP instead; MoT runs neither TP nor EP."
if isinstance(self.vlm, MoTConfig)
else "Use FSDP, TP, or EP instead."
)
raise ValueError(
"MoE + Pipeline Parallelism is not supported. MoE layers use "
"data-dependent routing that is incompatible with pipeline stage "
f"splitting. {alternatives}"
)
if self.distributed.ep > 1:
if isinstance(self.vlm, MoTConfig):
raise ValueError(
f"vlm.arch={_arch_of(self.vlm)!r} does not support expert parallelism "
f"(distributed.ep={self.distributed.ep}): the FSDP2 wrap for expert-parallel "
"blocks expects layer.attention and layer.mlp, and MoTBlock names its "
"attention attn and keeps its FFNs per modality, so a MoT model with experts "
"fails to build. Set distributed.ep = 1."
)
if isinstance(self.vlm, MoMaConfig):
raise ValueError(
f"vlm.arch={_arch_of(self.vlm)!r} does not support expert parallelism "
f"(distributed.ep={self.distributed.ep}): apply_expert_parallel partitions "
"MoEMLP experts only, and MoMa's per-modality expert groups need an EP-aware "
"dispatch that is not wired yet, so every rank would keep every expert. Set "
"distributed.ep = 1."
)
if not self.model.is_moe:
raise ValueError("ep > 1 requires an MoE model (num_experts > 0)")
if self.model.num_experts % self.distributed.ep != 0:
raise ValueError(
f"num_experts ({self.model.num_experts}) must be divisible by "
f"ep ({self.distributed.ep})"
)
if self.is_vlm:
assert self.vlm is not None
assert self.vision_encoder is not None
if self.vision_encoder.num_tokens > 0:
assert self.adapter is not None
per_frame = self.adapter.output_num_tokens(self.vision_encoder.num_tokens)
frames = self.views_per_sample
visual_tokens = frames * per_frame
residual_image_tokens = self.vlm.residual_stream_image_tokens(visual_tokens)
required = residual_image_tokens + self.vlm.max_text_len
if self.train.seq_len < required:
raise ValueError(
f"train.seq_len ({self.train.seq_len}) insufficient for VLM: "
f"residual_image_tokens ({residual_image_tokens}) + "
f"vlm.max_text_len ({self.vlm.max_text_len}) = {required}"
)
if self.distributed.pp > 1:
raise ValueError(
"VLM + Pipeline Parallelism is not supported on this branch. "
"The vlm-cross-attention branch will address PP integration."
)
if isinstance(self.vlm, MoMaConfig):
if self.model.num_experts > 0:
raise ValueError(
"MoMa + model.num_experts > 0 is rejected. MoMa derives expert "
"counts per modality from vlm.moma_experts_per_modality; set "
"model.num_experts=0 (the dense-LLM default) when using arch='moma'."
)
if self.train.compile_model:
import logging
logging.getLogger(__name__).warning(
"torch.compile is not yet validated for MoMa dispatch "
"(modality_ids-based scatter/gather + expert-choice top-k cause "
"graph breaks). Set compile_model=false for MoMa models."
)