Development

​

Table of Contents

• Test Suite
  • Layout
  • Running Tests Locally
  • CI Workflows
  • Markers
• Pre-Commit Hooks
• Adding a New Model
  • Implement the Modeling Code
  • Register a Mini Preset
  • Run the Smoke Test
• Adding a Custom VLM Implementation

​

Test Suite

​

Layout

• tests/unit/ — fast-running, hermetic tests for isolated logic: config parsing and validation, advantage / loss / scheduler / packer math, individual dataset paths, model-conversion roundtrips, etc. Tests that need a GPU are tagged with the gpu marker.
• tests/integration/ — full-stack RL/SFT runs on a tiny model end-to-end through inference + orchestrator + trainer.
• tests/nightly/ — runs the configs in examples/ every night to catch regressions in the shipped examples.

​

Running Tests Locally

uv run pytest -v                                           # everything
uv run pytest tests/unit -v                                # unit only
uv run pytest tests/integration -v                         # integration only
uv run pytest -v -m "not gpu"                              # CPU-only subset (mirrors CPU CI)
uv run pytest -v -m gpu                                    # GPU-only subset
uv run pytest tests/integration/test_reverse_text.py -vvs  # one specific scenario

​

CI Workflows

​

Markers

• gpu — gate a test that needs CUDA. CPU CI uses -m "not gpu"; the GPU unit job uses -m gpu.
• slow — gate a test that’s expensive enough you’d usually skip it locally. Deselect with -m "not slow".

​

Pre-Commit Hooks

uv run pre-commit install

​

Adding a New Model

​

Implement the Modeling Code

​

Register a Mini Preset

ARCH_PRESETS = {
    "glm4_moe": {
        "config_class": Glm4MoeConfig,
        "config_kwargs": dict(hidden_size=1024, num_hidden_layers=24, n_routed_experts=8, ...),
        "hf_model_class": HFGlm4MoeForCausalLM,
        "prime_model_class": PrimeRLGlm4MoeForCausalLM,
        "tokenizer_source": "THUDM/GLM-4-9B-0414",
    },
    # add your arch here
}

​

Run the Smoke Test

uv run python scripts/mini_moe.py --arch glm4_moe --output-dir ./mini-glm-moe

uv run python scripts/mini_moe.py --arch glm4_moe --output-dir ./mini-glm-moe --verify-only

uv run sft \
  --model.name ./mini-glm-moe \
  --data.name PrimeIntellect/Reverse-Text-SFT \
  --max_steps 200 \
  --optim.lr 1e-4 \
  --ckpt.weights

uv run rl @ configs/ci/integration/reverse-text-moe/start.toml \
  --model.name samsja/mini-glm-moe \
  --trainer.model.impl custom \
  --inference.gpu-memory-utilization 0.7 \
  --inference.model.max-model-len 2048

• No crashes. Validates the full inference + orchestrator + trainer pipeline end-to-end.
• Finite, non-zero KL. Confirms the reference distribution is meaningful.
• Loss reasonable. Not NaN, not stuck.

​

Requirements for merging a new model

• The model is correctly registered and defines and all the required methods - such as convert_hf_layer_to_tt and convert_tt_layer_to_hf.
• The small smoke test passes.

​

Adding a Custom VLM Implementation

1. Custom text model first. The VLM body wraps a custom *ForCausalLM (see Adding a New Model), so the text side — including its state-dict conversion and KL-mismatch table — comes first.
2. Composite VLM body. A *VLMModel that holds the HF vision encoder and the custom text model, with a prepare_inputs_embeds_and_position_ids step: embed tokens, run the vision encoder, scatter image embeddings over placeholder tokens, and build MRoPE 3D positions from mm_token_type_ids (the renderer owns the token→modality mapping). The unified *ForCausalLM dispatches on the config: composite config → VLM path, text config → text path.
3. Always run the vision encoder. Text-only micro-batches must feed the encoder dummy pixels and graft the result into the graph with zero contribution (inputs_embeds + image_embeds.sum() * 0.0) so FSDP/EP collectives stay symmetric across ranks when the encoder is trainable.
4. Packed-boundary consumption. Samples pack into shared rows with per-document boundaries in seq_lens; every custom model’s forward() declares the typed seq_lens/seq_lens_are_pre_shard parameters (the trainer passes them unconditionally) and must honor the boundaries — varlen flash cu_seqlens, linear-attention state resets per document, and a loud rejection on attention paths that can’t (see the packed-batch guard in any modeling file). Set supports_packed_multimodal_training on the VLM model once packed rows are handled — RL fails loudly at startup for VLM models without it.
5. Registration. Register the composite model_type in _CUSTOM_VLM_MAPPING (models/__init__.py) so get_model dispatches to the custom class, and describe the family in VLM_REGISTRY (utils/vlm.py).
6. Context parallelism (optional). CP-capable VLMs implement set_context_parallel_attributes and shard embeds/positions inside the model after the vision merge; the trainers defer sharding to the model for MRoPE batches under ulysses.
7. Validation. Same bar as text models: the KL-mismatch table for the text path, plus an SFT run and an RL run on a real multimodal dataset (the color-codeword environment is the reference task).