/slime-user

SLIME Source Code Quick Reference

This document provides quick access to important source code files, functions, and classes in the SLIME codebase.

Core Entry Points

Main Training Scripts

Path: train.py Purpose: Synchronous training entry point Key Function: train(args) - Main training loop

Path: train_async.py Purpose: Asynchronous training entry point for overlapping rollout and training Usage: Use instead of train.py for higher GPU utilization

Configuration and Arguments

Main Argument Parser

Path: slime/utils/arguments.py Key Functions:

• get_slime_extra_args_provider(): Returns arg provider for SLIME-specific args
• reset_arg(parser, name, **kwargs): Resets/adds Megatron argument defaults
• Internal arg groups:
  • add_cluster_arguments(): Ray cluster configuration
  • add_data_arguments(): Dataset and data processing
  • add_rollout_arguments(): Rollout/inference configuration
  • add_training_arguments(): RL training hyperparameters
  • add_reward_model_arguments(): Reward model setup
  • add_evaluation_arguments(): Evaluation configuration

Key Parameters Defined:

• Cluster: --actor-num-nodes, --actor-num-gpus-per-node, --rollout-num-gpus, --colocate
• Data: --prompt-data, --input-key, --label-key, --metadata-key, --apply-chat-template
• Rollout: --rollout-batch-size, --n-samples-per-prompt, --rollout-max-response-len, --rollout-temperature
• Training: --num-rollout, --num-steps-per-rollout, --global-batch-size, --advantage-estimator
• Reward: --rm-type, --custom-rm-path
• Evaluation: --eval-interval, --eval-prompt-data, --n-samples-per-eval-prompt

SGLang Arguments Integration

Path: slime/backends/sglang_utils/arguments.py Key Functions:

• add_sglang_arguments(parser): Adds SGLang parameters with --sglang- prefix
• validate_args(args): Validates SGLang-specific arguments

Usage: SGLang params passed as --sglang-<param_name> (e.g., --sglang-mem-fraction-static)

Data Types and Structures

Sample Type

Path: slime/utils/types.py Class: Sample Key Fields:

• prompt: str - Input prompt text
• response: str - Generated response text
• tokens: list[int] - Concatenated prompt + response token IDs
• response_length: int - Length of response in tokens
• reward: float - Reward score from RM
• loss_mask: list[int] - Binary mask for loss calculation (1=compute, 0=mask)
• truncated: bool - Whether response was truncated
• aborted: bool - Whether generation was aborted
• label: str - Ground truth label (for evaluation)
• metadata: dict - Custom metadata (session_id, tool_code, etc.)

Usage: Core data structure passed through rollout → reward → training pipeline

Rollout Function Types

Path: slime/rollout/base_types.py Classes:

• RolloutFnTrainOutput: Return type for training rollout
• RolloutFnEvalOutput: Return type for evaluation rollout

Rollout and Generation

Main Rollout Function

Path: slime/rollout/sglang_rollout.py Key Function: generate_rollout(args, rollout_id, data_buffer, evaluation=False) -> list[list[Sample]] Purpose: Main rollout orchestration using SGLang Key Features:

• Asynchronous generation with asyncio
• Dynamic sampling support
• Partial rollout caching
• DP rank balancing

Key Classes:

• GenerateState: Singleton for global generation state
  • tokenizer: HF tokenizer instance
  • processor: HF processor for VLMs
  • semaphore: Concurrency control
  • sampling_params: Default sampling configuration
  • dp_counts: DP rank load balancing
  • Methods:
    • reset(): Reset state for new rollout
    • submit_generate_tasks(samples): Submit async generation tasks
    • dp_rank_context(): Context manager for DP rank selection

Key Functions:

• async def generate_and_rm_group(args, samples, sampling_params, evaluation): Generate and compute rewards for sample group
• async def generate_single(args, sample, sampling_params, evaluation): Generate single sample
• async def call_sglang(prompt, sampling_params, ...): Call SGLang API

Custom Generation Function Signature:

async def generate(args, sample: Sample, sampling_params) -> Sample:
    # Custom generation logic
    # Must set: sample.tokens, sample.response_length, sample.response, sample.truncated
    return sample

Reward Models

Path: slime/rollout/rm_hub/ Key Functions:

• async_rm(args, sample: Sample, **kwargs) -> float: Async reward computation for single sample
• batched_async_rm(args, samples: list[Sample], **kwargs) -> list[float]: Batched reward computation

Custom RM Function Signature:

async def reward_func(args, sample: Sample, **kwargs) -> float:
    # Custom reward logic
    return reward_score

Filters and Metrics

Path: slime/rollout/filter_hub/ Key Files:

• dynamic_sampling_filters.py: Built-in filters
  • check_reward_nonzero_std(args, samples, **kwargs): Filter for reward diversity
• base_types.py: Filter interface definitions
  • MetricGatherer: Base class for metric collection
  • call_dynamic_filter(filter_fn, args, samples): Execute filter function

Training Backends

Megatron Backend

Path: slime/backends/megatron/ Key Files:

• arguments.py: Megatron-specific arguments
• actor.py: Actor model implementation
• trainer.py: Training loop implementation

Integration Points:

• Uses PYTHONPATH to find Megatron-LM installation
• Imports megatron.training.arguments.parse_args
• Checkpoint loading via load_checkpoint(), saving via save_checkpoint()
• Custom hooks:
  • --custom-megatron-init-path: Initialization hook
  • --custom-megatron-before-log-prob-hook-path: Pre-log-prob hook
  • --custom-megatron-before-train-step-hook-path: Pre-train-step hook

FSDP Backend

Path: slime/backends/fsdp/ Key Features:

• Native HuggingFace format support (no weight conversion)
• PyTorch FSDP2 implementation
• Auto-reads architecture from AutoConfig

Key Parameters:

• --train-backend fsdp: Enable FSDP backend
• --hf-checkpoint: HF model path (required, used for both rollout and training)
• --gradient-checkpointing: Enable gradient checkpointing
• --fsdp-cpu-offload: Offload to CPU
• --context-parallel-size: Context parallelism (supported)

Utilities

Data Management

Path: slime/utils/data.py Class: Dataset Purpose: Manage JSONL datasets with streaming support

Evaluation Configuration

Path: slime/utils/eval_config.py Classes:

• EvalDatasetConfig: Configuration for eval datasets
• Functions:
  • build_eval_dataset_configs(args): Build eval configs from args
  • ensure_dataset_list(data_arg): Parse dataset specifications

Processing Utilities

Path: slime/utils/processing_utils.py Key Functions:

• load_tokenizer(checkpoint, trust_remote_code): Load HF tokenizer
• load_processor(checkpoint, trust_remote_code): Load HF processor (for VLMs)
• encode_image_for_rollout_engine(image): Encode images for rollout

HTTP Utilities

Path: slime/utils/http_utils.py Key Functions:

• async def get(url, **kwargs): Async HTTP GET
• async def post(url, data, **kwargs): Async HTTP POST Usage: Communicate with SGLang servers and router

Async Utilities

Path: slime/utils/async_utils.py Key Functions:

• run(coro): Run async coroutine in sync context Usage: Bridge between sync and async code

Logging

Path: slime/utils/logging_utils.py Key Functions:

• configure_logger(name, level): Configure logger instance

Misc Utilities

Path: slime/utils/misc.py Key Items:

• SingletonMeta: Metaclass for singleton pattern
• load_function(path): Dynamically load function from module path

Conversion Tools

HF to Megatron Conversion

Path: tools/convert_hf_to_torch_dist.py Usage:

PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /path/to/hf_model \
    --save /path/to/torch_dist_output

Megatron to HF Conversion

Path: tools/convert_torch_dist_to_hf.py Usage:

PYTHONPATH=/root/Megatron-LM python tools/convert_torch_dist_to_hf.py \
    --input-dir /path/to/torch_dist/iter_xxx/ \
    --output-dir /path/to/hf_output \
    --origin-hf-dir /path/to/original_hf_model

FSDP to HF Conversion

Path: tools/convert_fsdp_to_hf.py Usage:

python tools/convert_fsdp_to_hf.py \
    --input-dir /path/to/fsdp_ckpt/iter_xxx \
    --output-dir /path/to/hf_output \
    --origin-hf-dir /path/to/original_hf_model

Model Configuration Files

Path: scripts/models/ Files:

• qwen3-4B.sh: Defines MODEL_ARGS array for Qwen3-4B
• glm4-9B.sh: Defines MODEL_ARGS array for GLM4-9B
• qwen3-30B-A3B.sh: MOE configuration for Qwen3-30B
• glm45-355B-A32B.sh: Large MOE configuration

Key Parameters in Model Configs:

• Architecture: --num-layers, --hidden-size, --ffn-hidden-size
• Attention: --num-attention-heads, --num-query-groups, --kv-channels
• Normalization: --normalization RMSNorm, --norm-epsilon
• Position: --use-rotary-position-embeddings, --rotary-base
• Vocab: --vocab-size, --disable-bias-linear
• Activation: --swiglu

Quick Code Navigation

Finding Reward Model Implementations

1. Check slime/rollout/rm_hub/ for built-in RMs
2. Use --rm-type <type> for built-in, --custom-rm-path for custom

Adding Custom Generation Logic

1. Implement async def generate(args, sample: Sample, sampling_params) -> Sample
2. Set via --custom-generate-function-path module.path:function_name
3. Reference implementation: slime/rollout/sglang_rollout.py (default)
4. Multi-turn example: examples/search-r1/search_r1_logic.py

Adding Custom Filters

1. Implement filter function: def filter_fn(args, samples: list[Sample], **kwargs) -> bool
2. Set via --dynamic-sampling-filter-path module.path:function_name
3. Built-in filters: slime/rollout/filter_hub/dynamic_sampling_filters.py

Customizing Rollout Function

1. Implement def generate_rollout(args, rollout_id, data_buffer, evaluation) -> list[list[Sample]]
2. Set via --rollout-function-path module.path:function_name
3. Default implementation: slime/rollout/sglang_rollout.py:generate_rollout

Adding Megatron Hooks

• Init hook: --custom-megatron-init-path module.path:function_name
• Pre-log-prob hook: --custom-megatron-before-log-prob-hook-path module.path:function_name
• Pre-train-step hook: --custom-megatron-before-train-step-hook-path module.path:function_name

Common Import Patterns

# Load SLIME utilities
from slime.utils.types import Sample
from slime.utils.http_utils import post
from slime.utils.processing_utils import load_tokenizer
from slime.rollout.rm_hub import async_rm

# Custom generation function
async def my_generate(args, sample: Sample, sampling_params) -> Sample:
    tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
    # ... generation logic
    return sample

# Custom reward function
async def my_reward(args, sample: Sample, **kwargs) -> float:
    # ... reward computation
    return reward