DOCSEARCHbyAlgolia
Submit Docs

torchtitan

repository·main·Indexed 26 days ago

https://github.com/pytorch/torchtitan

A PyTorch native platform for rapid experimentation and large-scale training of generative AI models. It provides a clean-room implementation of PyTorch scaling techniques optimized for models like Llama 3.1. Features include the ForgeEngine for lightweight post-training workflows, GraphTrainer for graph-pass optimization and pipeline parallelism (GraphPP), and an Autoresearch Harness for autonomous kernel and graph-pass optimization.

Tokens
31.3K
Snippets
61
Records
160
Agent score
91%

What's inside torchtitan

  1. Overview of GraphTrainer

    main
    GraphTrainer is an experimental component built on top of torchtitan designed to address CPU-side kernel launch overhead in high-speed accelerators. It captures the entire training step—including forward, loss, backward, and optionally optimizer.step—as a single unified FX graph. This approach allows for compiler-driven training where optimizations like activation checkpointing, CUDAGraph, CPU offload, and communication overlap are implemented as graph passes rather than complex runtime hooks.

  3. Available features for Qwen3 models

    main

    The Qwen3 model implementation supports various parallelism and optimization strategies depending on the model type:

    Dense Models:

    • Supports FSDP/HSDP, TP (Tensor Parallelism), CP (Context Parallelism), and DDP (Distributed Data Parallelism).
    • Supports AC (Activation Checkpointing) and torch.compile.

    MoE (Mixture-of-Experts) Models:

    • Supports FSDP/HSDP, TP, CP, DDP, and EP (Expert Parallelism).
    • Supports AC and torch.compile.
    • Uses Token Choice routing with an auxiliary-loss-free load balancing algorithm.

  4. Understand the design and scope of overrides

    main

    The override mechanism is designed for three primary use cases:

    1. Hardware-specific kernels: Implementing efficient Triton/CUDA kernels (e.g., for rotary embeddings or fused attention) that target specific hardware.
    2. Larger fused regions: Replacing multiple Config nodes with a single, more efficient implementation (e.g., a custom MoE implementation that fuses dispatch, compute, and combine).
    3. Team experimentation: Testing non-trivial implementations without making intrusive changes to the core repository.

    Important distinction: This mechanism is not an operator-override API. To replace or add a specific operator (op), you should use PyTorch's custom-operator path and wrap that op in a Module.

  5. Quickstart TorchTitan Structured Logging

    main

    TorchTitan provides structured logging for distributed training, emitting per-rank JSONL events for phase timing, diagnostics, and analysis. To use it, call init_logger() and sl.init_structured_logger() once per process before any trace calls. Every record automatically includes rank, source, caller (file:line:function), time_us, step, relative_step, and step_tags.

    from torchtitan.tools.logging import init_logger
from torchtitan.observability import structured_logger as sl

# console logger (stdout, [titan] prefix)
init_logger()

# Register handlers (e.g., to save to a local jsonl)
sl.init_structured_logger(source="training", output_dir="./outputs")

# Register a point-in-time marker
sl.log_trace_instant("training_start")

loaded_step = 0
for step in range(loaded_step + 1, num_steps + 1):
    # Stamp subsequent records with `step` and `relative_step`
    sl.set_step(step, relative_step=step - loaded_step)

    if should_garbage_collect:
        # Annotate the current step; tags reset at the next set_step()
        sl.add_step_tag("gc")
        with sl.log_trace_span("gc_collect"):
            run_gc()

    with sl.log_trace_span("fwd_bwd"):
        output = model(batch)
        loss.backward()

    with sl.log_trace_span("Optimizer"):
        optimizer.step()

    # Register scalars for debugging
    sl.log_trace_scalar({
        "num_trainable_tokens": num_trainable_tokens,
         "batch_size": bsz
         })

  6. Run benchmarks and extract metrics

    main

    To measure the performance of a graph optimization, run the benchmark script and extract steady-state metrics from the final step of the log.

    Run the benchmark:

    bash torchtitan/experiments/graph_trainer/autoresearch/scripts/run_benchmark.sh

    Extract metrics: Use grep to pull the last step from run.log:

    grep "step:" run.log | tail -1

    Output format: The output provides step, loss, grad_norm, memory, tps (tokens per second), tflops, and mfu (Model Flops Utilization).

    bash torchtitan/experiments/graph_trainer/autoresearch/scripts/run_benchmark.sh

grep "step:" run.log | tail -1

  7. Perform validation during training with the Validator class

    main

    To perform validation directly within the training loop, use the Validator class. It can be configured via Validator.Config within your config_registry function. The Validator class is designed to reuse the trainer's existing parallelization and pipelining capabilities.

    validator=Validator.Config(
    freq=500,
    dataset="c4_validation",
    steps=-1,  # consumes the entire validation set
),

  8. Start an Autoresearch run for graph-pass optimization

    main

    The Autoresearch Harness is an autonomous loop for optimizing graph-passes or kernels in graph_trainer. An LLM agent iteratively edits torchtitan/experiments/graph_trainer/passes.py, benchmarks changes, and retains them only if training-step time improves and numerics remain bitwise-identical to the eager reference.

    To start a run, follow these steps from the repository root:

    1. Configure Benchmarking: Edit torchtitan/experiments/graph_trainer/autoresearch/scripts/run_benchmark.sh to specify your model, configuration, and parallelism settings.
    2. Configure Agent Setup: Fill in the [SETUP] sections in torchtitan/experiments/graph_trainer/autoresearch/autoresearch.md. You must define the target, starting graph, scaffolding level (e.g., curated ideas, reference access, or web), reading-scope restrictions, and the numerics-check command.
    3. Seed Ideas (Optional): Add optimization directions to torchtitan/experiments/graph_trainer/autoresearch/ideas.md to guide the agent.
    4. Establish Baseline: Record your baseline performance as the first row in torchtitan/experiments/graph_trainer/autoresearch/results.tsv and the first entry in torchtitan/experiments/graph_trainer/autoresearch/experiment_log.md.
    5. Execute: Point the LLM agent at autoresearch.md to begin the autonomous loop.

  9. Use Experimental AutoParallel Sharding

    main

    GraphTrainer can use AutoParallel to solve SPMD placement. This requires --compile.mode aot_fx_trace. Once placed, the model is traced and compiled through the standard aot_fx_trace pipeline.

    # Llama 3 with AutoParallel
MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_debugmodel ./run_train.sh \
  --compile.mode aot_fx_trace \
  --compile.enable_autoparallel \
  --parallelism.data_parallel_shard_degree 2 \
  --parallelism.tensor_parallel_degree 2

  10. Maintain checkpoint compatibility with module-level hooks

    main

    If an override changes a module's parameter layout (e.g., fusing weights), it will change the checkpoint FQNs. To maintain compatibility with the original (stock) layout, register two hooks in your replacement module:

    1. register_state_dict_post_hook: To split or rename the new internal parameters into the stock FQNs during saving.
    2. register_load_state_dict_pre_hook: To recombine the stock parameters into the new internal layout before the default load process.

    For highly complex mappings, use a model-level BaseStateDictAdapter.

  11. Implement custom component overrides

    main

    You can override any Configurable component within the Trainer.Config tree without modifying the model's original config_registry.py or __init__.py.

    Key Concepts

    • Activation: Overrides are opt-in via override.imports. Only overrides registered by the modules listed in override.imports (which are provenance-checked) will be applied.
    • Targeting: Use Fully Qualified Names (FQN) globs on the @override decorator to target specific instances.
    • Subclass Matching: By default, targets match subclasses. If a replacement only supports a specific concrete target configuration, use exact=True.
    • Conflicts: Overrides are resolved per-node. Multiple overrides can target the same class as long as their claimed nodes in the config tree are disjoint. Errors occur if there are same-node or ancestor/descendant claims.
    • Per-instance configuration: You can pass kwargs to an override to configure the same module differently across different config trees. In Python, this is a (module_path, kwargs) tuple; on the CLI, use module=<json>.

  12. Download Llama 3.1 tokenizer

    main

    To train Llama 3.1 models (8B, 70B, 405B), you must download the tokenizer. Ensure you have access to the Llama model weights on Hugging Face first. Use the scripts/download_hf_assets.py script with your Hugging Face token.

    # Get your HF token from https://huggingface.co/settings/tokens

# Llama 3.1 tokenizer
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...