TorchSpec

TorchSpec is organized into several functional domains to support asynchronous training and inference pipelines:

• config/: Hierarchical configuration using OmegaConf (e.g., train_config.py, inference_config.py).
• ray/: Ray infrastructure for GPU setup and resource management.
• controller/: Orchestration of the async pipeline via AsyncTrainingController and AsyncInferenceManager.
• inference/: Inference engine layer providing unified interfaces for HFEngine, SglEngine, and VllmEngine.
• models/: Model definitions including draft/ (lightweight transformers), target/ (target model abstractions), and ops/ (custom operations like KL loss).
• training/: Training logic including TrainerActor (Ray actor), Eagle3Trainer, and FSDP2 helpers.
• transfer/mooncake/: Distributed tensor transfer using Mooncake for multi-node RDMA-based communication.
• data/: Data pipelines for loading, parsing, and preprocessing datasets.
• utils/: Shared utilities for distributed setup, logging, and memory estimation.

Placement groups reserve GPUs and place them on specific nodes using create_placement_groups(args, roles=...). The mode selected determines how training and inference GPUs are allocated:

Note: placement_strategy only applies when online training creates both roles; single-role workflows use PACK placement.

The training dataset must be a JSONL file where each row contains a conversations field following the OpenAI-style conversation format. This supports text-only, multimodal (single/multi-image), system messages, multi-turn, and tool calls.

Multimodal Requirements

• Config: Set chat_template=kimi-k25-vlm and inference.sglang.enable_multimodal=true in your configuration.
• Content Structure: For multimodal samples, use a list-of-dicts content format with image_url or image fields. Do not use plain-string placeholders for images, as they will not send pixel data to the engine.
• Image URLs: For images stored in remote storage (S3, GCS, etc.), use pre-signed URLs so the SGLang engine can fetch them without credentials.
• Assistant Responses: Assistant content should include <think>...</think> blocks.

Example of a multimodal JSONL row:

{
  "id": "kimi_mm_001",
  "conversations": [
    {
      "role": "user",
      "content": [
        {"type": "image_url", "image_url": {"url": "https://upload.wikimedia.org/...jpg"}},
        {"type": "text", "text": "What animal is in this image?"}
      ]
    },
    {
      "role": "assistant",
      "content": "<think>The image shows a domestic cat.</think>This is a cat."
    }
  ]
}

A 'tap' is a probe used to read the signal from the residual stream without interrupting the forward pass. To ensure a tap is a 'pure read' and does not corrupt the model state, it must be called with specific parameters that prevent side effects (like folding the delta or writing to the bank).

To perform a pure read, use the attn_res call with the following configuration:

• Set output_norm_weight = None: This ensures the kernel returns the mixed (softmax mixture) value rather than the normed value.
• Set delta = None: Prevents the kernel from folding the delta back into the prefix.
• Set block_write_idx = -1: Prevents the kernel from committing a new block into the bank.
• Pass a freshly computed prefix (calculated as prefix_sum + pending_mlp_out) rather than the live prefix_sum used by the model.

# Example of a pure tap call
attn_res(prefix, None, bank, score_norm, score_proj, None, num_blocks=...,
         block_write_idx=-1, ...)
# Note: output_norm_weight = None returns the `mixed` value

Datasets must be in JSONL format where each row contains a conversations field following an OpenAI-style structure.

Multimodal Requirements:

• Set chat_template=kimi-k25-vlm and inference.sglang.enable_multimodal=true in your configuration.
• Assistant responses must include <think>...</think> blocks.
• For multimodal samples, use a list-of-dicts content structure with image_url or image fields. Do not use plain-string `` placeholders for multimodal samples, as they will not send pixel data to the engine.
• For remote images (S3/GCS), use pre-signed URLs so the engine can fetch them without credentials.

Example JSONL row (multimodal):

{
  "id": "kimi_mm_001",
  "conversations": [
    {
      "role": "user",
      "content": [
        {"type": "image_url", "image_url": {"url": "https://upload.wikimedia.org/...jpg"}},
        {"type": "text", "text": "What animal is in this image?"}
      ]
    },
    {
      "role": "assistant",
      "content": "<think>The image shows a domestic cat.</think>This is a cat."
    }
  ]
}

The training entry point (torchspec/train_entry.py → train_async_no_generation) follows a specific sequence. If a job appears stuck, use this guide to identify the current phase:

Note on Phase 5: This is the most common phase where jobs appear stuck. After shards reach 100%, the system silently allocates KV cache and captures CUDA graphs, which can take several minutes without any log output.