Slime RL Training: LLM Post-Training with Megatron and SGLang

{ "title": "SLiME Reference Document", "body": "---\n# SLiME Reference Document\n\nSLiME is an LLM post-training framework from Tsinghua THUDM designed for reinforcement learning (RL) scaling. It powers the GLM-4.5, GLM-4.6, and GLM-4.7 model series and connects Megatron-LM for training with SGLang for rollout generation. The framework is backed by production deployment at Z.ai and supports training GLM, Qwen3, DeepSeek V3, and Llama 3 models.\n\n## When to Use SLiME\n\nUse SLiME when you need Megatron-LM native training with SGLang inference, require custom data generation workflows with flexible data buffers, or want to train GLM, Qwen3, DeepSeek V3, or Llama 3 models. It is a research-grade framework with production backing. For enterprise-grade stability, consider miles. For flexible backend swapping, consider verl. For PyTorch-native abstractions, consider torchforge.\n\n## Capabilities\n\nSLiME provides training with Megatron-LM full parallelism (tensor, pipeline, data, and sequence parallelism). Rollout generation uses SGLang with a router for high throughput. The Data Buffer handles flexible prompt management and sample storage. The framework supports GLM-4.x, Qwen3, DeepSeek V3/R1, and Llama 3. Asynchronous training mode overlaps rollout and training. Multi-turn agentic training with custom generate functions and tool use is supported. Co-location mode shares GPUs between training and inference. Custom reward model integration is available. Multi-task evaluation uses separate prompt datasets. Fault tolerance handles SGLang engine crashes. Additional features include gradient checkpointing, sequence parallelism, and streaming data.\n\n## Architecture\n\nThe architecture centers on a Data Buffer that connects training and rollout components.\n\n\n┌─────────────────────────────────────────────────────────┐\n│ Data Buffer │\n│ - Prompt initialization and management │\n│ - Custom data generation and filtering │\n│ - Rollout sample storage │\n└─────────────┬───────────────────────────┬───────────────┘\n │ │\n┌─────────────▼───────────┐ ┌─────────────▼───────────────┐\n│ Training (Megatron-LM) │ │ Rollout (SGLang + Router) │\n│ - Actor model training │ │ - Response generation │\n│ - Critic (optional) │ │ - Reward/verifier output │\n│ - Weight sync to rollout│ │ - Multi-turn support │\n└─────────────────────────┘ └─────────────────────────────┘\n\n\n## Installation\n\n### Docker (Recommended)\n\n\n# Recommended: Docker\ndocker pull slimerl/slime:latest\ndocker run --rm --gpus all --ipc=host --shm-size=16g \\\n -it slimerl/slime:latest /bin/bash\n\n# Inside container\ncd /root/slime && pip install -e . --no-deps\n\n\n### From Source\n\n\ngit clone https://github.com/THUDM/slime.git\ncd slime\npip install -r requirements.txt\npip install -e .\n\n\n## Prerequisites\n\nYou need a Docker environment or Megatron-LM plus SGLang installed. A model checkpoint in HuggingFace or Megatron format is required. Training data must be in JSONL format. For multi-turn agentic training, you need a custom generate function and a tool or environment interface.\n\n## Data Format\n\nTraining data uses JSONL format. Simple prompt-label pairs look like this:\n\n\n# data.jsonl format\n{\"prompt\": \"What is 2 + 2?\", \"label\": \"4\"}\n{\"prompt\": \"Solve: 3x = 12\", \"label\": \"x = 4\"}\n\n\nChat format with system, user, and tool roles is also supported:\n\n\n{\n \"prompt\": [\n {\"role\": \"system\", \"content\": \"You are a math tutor.\"},\n {\"role\": \"user\", \"content\": \"What is 15 + 27?\"}\n ],\n \"label\": \"42\"\n}\n\n\n## Model Configuration\n\nPre-configured model scripts are available in scripts/models/. List available models and source the one you need:\n\n\n# List available models\nls scripts/models/\n# glm4-9B.sh, qwen3-4B.sh, qwen3-30B-A3B.sh, deepseek-v3.sh, llama3-8B.sh, ...\n\n# Source your model\nsource scripts/models/qwen3-4B.sh\n\n\n## Procedures\n\n### Standard GRPO Training\n\n1. Prepare data in JSONL format (simple prompt-label or chat format with system/user/tool roles).\n2. Choose a pre-configured model script from scripts/models/.\n3. Launch training with python train.py, specifying actor nodes and GPUs, rollout GPUs, advantage estimator, KL loss, batch sizes, prompt data path, input and label keys, chat template flag, rollout count, save and eval intervals, and model args.\n4. Monitor training via TensorBoard using tensorboard --logdir outputs/, verify reward curves are increasing, and check GPU utilization.\n\nSource the model configuration and launch training:\n\n\n# Source model configuration\nsource scripts/models/qwen3-4B.sh\n\n# Launch training\npython train.py \\\n --actor-num-nodes 1 \\\n --actor-num-gpus-per-node 4 \\\n --rollout-num-gpus 4 \\\n --advantage-estimator grpo \\\n --use-kl-loss --kl-loss-coef 0.001 \\\n --rollout-batch-size 32 \\\n --n-samples-per-prompt 8 \\\n --global-batch-size 256 \\\n --num-rollout 3000 \\\n --prompt-data /path\n\n\n### Asynchronous Training\n\nUse asynchronous training when working with large models that have long generation times or when you see high GPU idle time in synchronous mode. Ensure sufficient memory for buffering.\n\n\npython train_async.py \\\n --actor-num-nodes 1 \\\n --actor-num-gpus-per-node 8 \\\n --rollout-num-gpus 8 \\\n --advantage-estimator grpo \\\n --async-buffer-size 4 \\\n --prompt-data /path/to/train.jsonl \\\n ${MODEL_ARGS[@]}\n\n\nKey parameters for async mode:\n\n\n--async-buffer-size 4 # Number of rollouts to buffer\n--update-weights-interval 2 # Sync weights every N rollouts\n\n\n### Multi-Turn Agentic Training\n\n1. Define a custom async generate function (e.g., in custom_generate.py) that handles multi-turn logic with tool calling, extracts tool calls, executes tools, appends tool results, and computes rewards.\n2. Launch with python train.py, specifying the custom generate function path, max turns, prompt data, and model args.\n3. See examples/search-r1/ for a complete multi-turn search example.\n\n\n# custom_generate.py\nasync def custom_generate(args, samples, evaluation=False):\n \"\"\"Multi-turn generation with tool calling.\"\"\"\n for sample in samples:\n conversation = sample.prompt\n\n for turn in range(args.max_turns):\n # Generate response\n response = await generate_single(conversation)\n\n # Check for tool call\n tool_call = extract_tool_call(response)\n if tool_call:\n tool_result = execute_tool(tool_call)\n conversation.append({\"role\": \"assistant\", \"content\": response})\n conversation.append({\"role\": \"tool\", \"content\": tool_result})\n else:\n break\n\n sample.response = response\n sample.reward = compute_reward(sample)\n\n return samples\n\n\n\npython train.py \\\n --custom-generate-function-path custom_generate.py \\\n --max-turns 5 \\\n --prompt-data /path/to/agent_data.jsonl \\\n ${MODEL_ARGS[@]}\n\n\n## Parameters\n\n### Required Parameters\n\n| Parameter | Meaning |\n|---|---|\n| --actor-num-nodes | Number of nodes for actor model training. |\n| --actor-num-gpus-per-node | Number of GPUs per node for actor training. |\n| --rollout-num-gpus | Number of GPUs for rollout generation. |\n| --advantage-estimator | Algorithm for advantage estimation (grpo, gspo, ppo, reinforce_plus_plus). |\n| --rollout-batch-size | Number of prompts per rollout batch. |\n| --n-samples-per-prompt | Number of responses to generate per prompt. |\n| --global-batch-size | Total batch size for training (must equal rollout_batch_size * n_samples_per_prompt * num_steps_per_rollout). |\n| --num-rollout | Total number of rollout steps. |\n| --prompt-data | Path to training data JSONL file. |\n\n### Optional Parameters\n\n| Parameter | Meaning | Required When |\n|---|---|---|\n| --use-kl-loss | Enable KL divergence loss. | No |\n| --kl-loss-coef | Coefficient for KL loss (e.g., 0.001). | If --use-kl-loss |\n| --input-key | Key in JSONL for prompt (default: prompt). | No |\n| --label-key | Key in JSONL for label (default: label). | No |\n| --apply-chat-template | Apply chat template to prompts. | No |\n| --save-interval | Steps between saving checkpoints (e.g., 100). | No |\n| --eval-interval | Steps between evaluations (e.g., 50). | No |\n| --async-buffer-size | Number of rollouts to buffer in async mode (e.g., 4). | For async training |\n| --update-weights-interval | Sync weights every N rollouts (e.g., 2). | For async training |\n| --custom-generate-function-path | Path to Python file with custom async generate function. | For multi-turn agentic training |\n| --max-turns | Maximum number of turns in multi-turn generation (e.g., 5). | For multi-turn agentic training |\n| --colocate | Share GPUs between training and inference. | No |\n| --sglang-mem-fraction-static | Fraction of GPU memory for SGLang (e.g., 0.8). | No |\n| --sglang-context-length | Context length for SGLang (e.g., 8192). | No |\n| --sglang-log-level | Log level for SGLang (e.g., INFO). | No |\n| --tensor-model-parallel-size | Tensor model parallelism size (Megatron arg). | No |\n| --pipeline-model-parallel-size | Pipeline model parallelism size (Megatron arg). | No |\n| --num-layers | Number of transformer layers (Megatron arg). | No |\n| --hidden-size | Hidden size (Megatron arg). | No |\n| --use-fault-tolerance | Enable fault tolerance for SGLang engine crashes. | No |\n| --recompute-activations | Enable gradient checkpointing to reduce memory. | No |\n| --micro-batch-size | Micro-batch size for training (e.g., 1). | No |\n| --sequence-parallel | Enable sequence parallelism. | No |\n| --num-data-workers | Number of data loading workers (e.g., 4). | No |\n| --streaming-data | Use streaming dataset for data loading. | No |\n| --custom-rm-path | Path to custom reward model Python file. | No |\n| --eval-prompt-data | Specify evaluation datasets with name and path (e.g., aime /path/to/aime.jsonl). Can be repeated. | No |\n| --n-samples-per-eval-prompt | Number of samples per evaluation prompt (e.g., 16). | No |\n\n## Constraints and Caveats\n\nThe product rollout_batch_size * n_samples_per_prompt must equal global_batch_size * num_steps_per_rollout.\n\n\nrollout_batch_size × n_samples_per_prompt = global_batch_size × num_steps_per_rollout\n\n\nThe SGLang engine may crash mid-training. Use --use-fault-tolerance, increase --sglang-mem-fraction-static, or reduce --rollout-batch-size.\n\nWeight sync timeout can cause training hangs. Increase --update-weights-interval or use --colocate.\n\nOOM during backward pass can occur. Enable --recompute-activations, reduce --micro-batch-size, or enable --sequence-parallel.\n\nSlow data loading can leave GPUs idle. Increase --num-data-workers or use --streaming-data.\n\nCo-location mode reduces memory but requires careful --sglang-mem-fraction-static setting.\n\nThe custom generate function must be async and return samples with response and reward.\n\nModel scripts are pre-configured. See scripts/models/ for available models (GLM, Qwen, DeepSeek, Llama, Kimi K2, Moonlight-16B). See examples/ for usage examples.\n\n## Data Buffer Customization\n\nThe default RolloutDataSource fetches prompts from the dataset:\n\n\nclass RolloutDataSource:\n def get_samples(self, num_samples):\n \"\"\"Fetch prompts from dataset.\"\"\"\n return self.dataset.sample(num_samples)\n\n def add_samples(self, samples):\n \"\"\"Called after generation (no-op by default).\"\"\"\n pass\n\n\nFor custom selection logic, extend RolloutDataSourceWithBuffer:\n\n\nclass RolloutDataSourceWithBuffer(RolloutDataSource):\n def __init__(self):\n self.buffer = []\n\n def add_samples(self, samples):\n \"\"\"Store generated samples for reuse.\"\"\"\n self.buffer.extend(samples)\n\n def buffer_filter(self, args, buffer, num_samples):\n \"\"\"Custom selection logic (prioritized, stratified, etc.).\"\"\"\n return select_best(buffer, num_samples)\n\n\n## Failure Modes and Mitigations\n\n### SGLang Engine Crash\n\nThe inference engine dies mid-training. Mitigate by enabling fault tolerance, increasing memory allocation, or reducing batch size:\n\n\n# Enable fault tolerance\n--use-fault-tolerance\n\n# Increase memory allocation\n--sglang-mem-fraction-static 0.85\n\n# Reduce batch size\n--rollout-batch-size 16\n\n\n### Weight Sync Timeout\n\nTraining hangs after rollout. Mitigate by increasing the sync interval or using co-located mode:\n\n\n# Increase sync interval\n--update-weights-interval 5\n\n# Use colocated mode (no network transfer)\n--colocate\n\n\n### OOM During Training\n\nCUDA OOM in backward pass. Mitigate by enabling gradient checkpointing, reducing micro-batch size, or enabling sequence parallelism:\n\n\n# Enable gradient checkpointing\n--recompute-activations\n\n# Reduce micro-batch size\n--micro-batch-size 1\n\n# Enable sequence parallelism\n--sequence-parallel\n\n\n### Slow Data Loading\n\nGPU idle during data fetch. Mitigate by increasing data workers or using streaming:\n\n\n# Increase data workers\n--num-data-workers 4\n\n# Use streaming dataset\n--streaming-data\n\n\n## Advanced Features\n\n### Co-Location Mode\n\nShare GPUs between training and inference to reduce network transfer:\n\n\npython train.py \\\n --colocate \\\n --actor-num-gpus-per-node 8 \\\n --sglang-mem-fraction-static 0.4 \\\n ${MODEL_ARGS[@]}\n\n\n### Custom Reward Model\n\nIntegrate a custom reward model by defining a class and specifying its path:\n\n\n# custom_rm.py\nclass CustomRewardModel:\n def __init__(self, model_path):\n self.model = load_model(model_path)\n\n def compute_reward(self, prompts, responses):\n inputs = self.tokenize(prompts, responses)\n scores = self.model(inputs)\n return scores.tolist()\n\n\n\n--custom-rm-path custom_rm.py\n\n\n### Multi-Task Evaluation\n\nEvaluate on multiple datasets by repeating the --eval-prompt-data parameter:\n\n\n--eval-prompt-data aime /path/to/aime.jsonl \\\n--eval-prompt-data gsm8k /path/to/gsm8k.jsonl \\\n--n-samples-per-eval-prompt 16\n\n\n### Megatron Parallelism Arguments\n\nSpecify Megatron parallelism and model architecture:\n\n\n--tensor-model-parallel-size 2\n--pipeline-model-parallel-size 1\n--num-layers 32\n--hidden-size 4096\n\n\n### SGLang Configuration\n\nConfigure SGLang memory, context length, and logging:\n\n\n--sglang-mem-fraction-static 0.8\n--sglang-context-length 8192\n--sglang-log-level INFO\n\n\n## Quick Reference\n\n\n# Resource allocation\n--actor-num-nodes 1\n--actor-num-gpus-per-node 8\n--rollout-num-gpus 8\n--colocate # Share GPUs between training/inference\n\n# Data\n--prompt-data /path/to/data.jsonl\n--input-key prompt\n--label-key label\n\n# Training loop\n--num-rollout 3000\n--rollout-batch-size 32\n--n-samples-per-prompt 8\n--global-batch-size 256\n\n# Algorithm\n--advantage-estimator grpo # or: gspo, ppo, reinforce_plus_plus\n--use-kl-loss\n--kl-loss-coef 0.001\n\n\n## Version Information\n\nSLiME was first released in 2025.\n---" }