Skip to content
GitHub
View on GitHub

Qwen3_6_35b_Recipe

Qwen3.6-35B-A3B (MoE) on 1×8×H100 with TP2/PP2/CP1/EP4.

from modal_training_gym.train_recipes.slime_recipe.qwen3_6_35b import Qwen3_6_35b_Recipe

Qwen3.6-35B-A3B (MoE) on 1×8×H100 with TP2/PP2/CP1/EP4.

Inherits from: SlimeRecipe, BaseTrainRecipe

FieldTypeDefaultDescription
recipe_typeRecipeTypeslimeDiscriminator marking this recipe as slime; never override.
namestr""Modal app title; when empty the launcher derives one from the recipe class.
app_tagsdict{}Extra tags merged into the Modal app metadata for dashboard auto-discovery.
FieldTypeDefaultDescription
environmentdict{'PYTHONPATH': '/root/Megatron-LM/', 'CUDA_DEVICE_MAX_CONNECTIONS': '1', 'NCCL_NVLS_ENABLE': '1'}Env vars set in the training containers (defaults include PYTHONPATH for Megatron and NCCL tuning).
async_modeboolFalseRun slime’s train_async.py so rollout generation and training overlap (one-step off-policy) instead of alternating.
wandbWandbConfig | NoneNoneW&B settings; expands to slime’s --use-wandb/--wandb-project/ --wandb-group flags.
image_overlaycollections.abc.Callable[[modal.image.Image], modal.image.Image] | NoneNoneCallable that customizes the Modal image (e.g. lambda img: img.pip_install("pkg")).
local_slimestr | NoneNonePath to a local slime checkout mounted over the image’s copy — dev overlay for testing slime changes without an image rebuild.
memoryint | tuple[int, int] | NoneNoneModal Function memory request/limit in MiB.
cloudstr | NoneNoneModal cloud provider to pin the cluster to.
regionstr | NoneNoneModal region to pin the cluster to.
slime_model_scriptstr"scripts/models/qwen3.5-35B-A3B.sh"Script path relative to the slime repo, sourced before train.py to provide MODEL_ARGS; when set, model-architecture flags are not emitted from the attached ModelConfig.
source_hf_checkpointstr | NoneNoneHF repo fetched as the source checkpoint when it differs from the model’s own (used by some recipes).
megatron_conversion_hf_checkpointstr | NoneNoneHF checkpoint used for the HF→Megatron conversion step instead of the training model’s own weights.
patch_fileslist[str][]Local patch scripts copied into the image at build time (applied to slime/Megatron sources).
image_run_commandslist[str][]Extra shell commands run while building the image.
image_envdict[str, str]{}Extra env vars baked into the image.
train_function_kwargsdict[str, int]{'ephemeral_disk': 1048576}Extra Modal Function options for the train function; supported keys: secrets, experimental_options, ephemeral_disk.
capture_traceboolFalseAttach slime’s per-sample execution trace (generate/reward/tool-call timeline) to recorded rollouts for the dashboard.
trace_sample_limitint16With capture_trace, number of samples per rollout that get a trace attached (sampling keeps the added data volume small).
FieldTypeDefaultDescription
gpu_typestr"H100"Modal GPU type for every node, e.g. "H100" or "B200".
colocateboolTrueTrainer and rollout engines share the same GPUs, shifting memory between phases; False gives each its own GPUs (disaggregated).
actor_num_nodesint1Number of nodes for the Megatron actor (trainer).
actor_num_gpus_per_nodeint8GPUs per actor node.
rollout_num_gpusint | NoneNoneTotal GPUs for rollout engines when disaggregated; None lets the allocation resolver size it.
rollout_num_gpus_per_engineint4GPUs per sglang engine — its tensor-parallel size.
tensor_model_parallel_sizeint2Megatron tensor-parallel size for the actor.
sequence_parallelboolTrueMegatron sequence parallelism (requires TP > 1).
use_criticboolFalseTrain a separate critic model (PPO-style; GRPO runs without one).
critic_num_nodesint | NoneNoneNodes for the critic when use_critic is set.
critic_num_gpus_per_nodeint | NoneNoneGPUs per critic node.
FieldTypeDefaultDescription
num_rolloutint3000Total rollout steps (= training steps) for the run.
rollout_batch_sizeint16Prompts sampled per rollout step; each prompt is expanded into a group of sampled responses.
rollout_max_response_lenint16384Max generated tokens per sample.
rollout_temperaturefloat1.0Sampling temperature for rollout generation.
rollout_shuffleboolTrueShuffle the prompt dataset between epochs.
rollout_top_pfloat1.0Nucleus-sampling top-p for rollout generation.
rollout_stop_token_idslist[int] | NoneNoneExtra token ids that terminate generation.
FieldTypeDefaultDescription
use_fault_toleranceboolTrueEnable slime’s fault tolerance so the run can recover from worker failures.
rollout_health_check_intervalint30Interval in seconds between rollout engine /health_generate checks during generate/eval.
rollout_health_check_timeoutint30Timeout in seconds to wait for a rollout engine /health_generate response before killing it.
rollout_health_check_first_waitint300Initial grace period (in seconds) before starting health checks. This allows time for model compilation and initialization. Increase this value significantly when using deepgemm.
FieldTypeDefaultDescription
savestr"/checkpoints"Checkpoint output directory (the mounted /checkpoints volume).
save_intervalint20Save a checkpoint every N rollout steps.
loadstr""Checkpoint directory to resume from; empty starts from the converted HF weights.
no_save_optimboolTrueOmit optim state from checkpoints (smaller, but no exact resume).
megatron_to_hf_modestr""Mode used to export saved Megatron checkpoints back to HF format; empty disables the export step.
freeze_params_name_listlist[str] | NoneNoneRegex patterns (matched with re.search) of parameter names to freeze, e.g. a VL model’s vision tower so RL only updates the language backbone.
FieldTypeDefaultDescription
advantage_estimatorstr"grpo"Advantage estimator, e.g. "grpo".
n_samples_per_promptint8Responses sampled per prompt (the GRPO group size).
eps_clipfloat0.2PPO clip lower bound.
eps_clip_highfloat0.28PPO clip upper bound (asymmetric DAPO-style clipping).
use_kl_lossboolTrueAdd a per-token KL loss term against the reference model.
kl_loss_typestr"low_var_kl"KL formulation, e.g. "low_var_kl".
kl_loss_coeffloat0.0Coefficient of the KL loss term.
kl_coeffloat0.0KL penalty coefficient applied in the reward.
entropy_coeffloat0.0Entropy bonus coefficient.
calculate_per_token_lossboolFalseAverage the loss over tokens instead of over samples.
ref_loadstr"/checkpoints/Qwen3.6-35B-A3B_torch_dist_tp2pp2"Checkpoint path the reference model is read from (for KL terms).
FieldTypeDefaultDescription
over_sampling_batch_sizeint | NoneNonePrompts sampled beyond rollout_batch_size so groups rejected by the filter can be replaced (DAPO).
dynamic_sampling_filter_pathstr | NoneNoneImport path of the predicate deciding which sample groups to keep, e.g. dropping all-equal-reward groups.
balance_databoolTrueRebalance kept samples across data-parallel ranks.
FieldTypeDefaultDescription
global_batch_sizeint128Training samples per optim step.
lrfloat1e-06Learning rate.
lr_decay_stylestr"constant"Schedule, e.g. "constant" or "cosine".
weight_decayfloat0.1Weight decay.
adam_beta1float0.9Adam beta1.
adam_beta2float0.98Adam beta2.
optimizerstr"adam"Optimizer name, e.g. "adam".
FieldTypeDefaultDescription
attention_dropoutfloat0.0Attention dropout probability.
hidden_dropoutfloat0.0Hidden-layer dropout probability.
attention_softmax_in_fp32boolTrueCompute attention softmax in fp32.
accumulate_allreduce_grads_in_fp32boolTrueAccumulate and all-reduce gradients in fp32.
use_distributed_optimizerboolFalseShard optim state across data-parallel ranks (Megatron distributed optimizer).
recompute_granularitystr"full"Activation recomputation granularity ("full" or "selective").
recompute_methodstr"uniform"Recomputation method ("uniform" or "block").
recompute_num_layersint1Layers per recomputation chunk.
qkv_formatstr"thd"QKV layout for the Megatron backend ("thd" or "bshd"), emitted as --qkv-format.
FieldTypeDefaultDescription
use_dynamic_batch_sizeboolTruePack variable-length samples into micro-batches up to max_tokens_per_gpu instead of a fixed micro batch size.
max_tokens_per_gpuint8192Token budget per GPU per micro-batch when dynamic batching is on.
FieldTypeDefaultDescription
eval_intervalint | NoneNoneRun eval every N rollout steps; None disables eval.
n_samples_per_eval_promptint4Responses sampled per eval prompt.
eval_max_response_lenint16384Max generated tokens per eval sample.
eval_top_pfloat1.0Nucleus-sampling top-p for eval generation.
eval_configdict | NoneNoneInline dict materialized to a YAML file and passed as --eval-config; holds eval defaults and the eval dataset list.
FieldTypeDefaultDescription
update_weight_modestr"full""full" rebroadcasts all weights each sync; "delta" pin-snapshots the last broadcast on CPU and ships only byte-level changes (~5-10x faster for large MoE models whose weights barely move per rollout).
update_weight_transportstr"nccl""nccl" or "disk"; disk requires trainer and rollout engines to share a filesystem.
update_weight_encodingstr"indices"Encoding for delta payloads, e.g. "indices".
update_weight_disk_dirstr""Shared directory used by the disk transport.
FieldTypeDefaultDescription
rm_typestr | NoneNoneName of a slime built-in reward function (e.g. "deepscaler"); leave None when shipping a reward callable instead.
FieldTypeDefaultDescription
custom_rm_functioncollections.abc.Callable | NoneNoneReward callable shipped by value to the containers and registered as slime’s --custom-rm-path.
custom_generate_functioncollections.abc.Callable | NoneNoneCallable replacing slime’s generate step; shipped by value and registered via its resolved import path.
custom_reward_post_process_functioncollections.abc.Callable | NoneNoneCallable applied to rewards after generation. Prefer this over setting a raw dotted path yourself: functions defined in a __main__ tutorial script have no reliably importable module name, so slime’s own importlib.import_module on that path fails inside the Ray actor.
rollout_functioncollections.abc.Callable | str | NoneNoneReplaces slime’s entire rollout loop (--rollout-function-path).
custom_rollout_log_functioncollections.abc.Callable | str | NoneNoneCalled with each rollout’s data for logging; the gym wraps it so phase reporting and dashboard capture still run.
custom_eval_rollout_log_functioncollections.abc.Callable | str | NoneNoneSame as above, for eval rollouts.
custom_megatron_before_log_prob_hookcollections.abc.Callable | str | NoneNoneHook run in the Megatron trainer before log-prob computation.
custom_megatron_before_train_step_hookcollections.abc.Callable | str | NoneNoneHook run in the Megatron trainer before each train step.
FieldTypeDefaultDescription
extra_configdict | NoneNoneThe primary escape hatch: dict written to YAML and passed as --custom-config-path. Keys become attributes on slime’s parsed args and always override same-named recipe fields.
sglang_configdict | NoneNoneDict written to YAML and passed as --sglang-config — structured sglang engine config that isn’t a flat flag (e.g. PD-disaggregation server_groups).
sglang_request_paramsdict | NoneNoneExtra request parameters injected into sglang generate calls (shipped via extra_config, read as args.sglang_request_params by generate paths such as on-policy distillation).
apply_chat_template_kwargsdict | str""Kwargs forwarded to the tokenizer’s apply_chat_template, passed as inline JSON.
train_env_varsdict | str | NoneNoneEnv vars for the training processes, passed as inline JSON.
multimodal_keysdict | str | NoneNoneDataset columns holding multimodal inputs, passed as inline JSON; auto-filled from the attached DatasetConfig.
FieldTypeDefaultDescription
sglang_mem_fraction_staticfloat0.75Fraction of GPU memory sglang reserves for weights + KV cache.
sglang_enable_dp_attentionboolTrueEnable data-parallel attention across engine ranks.
sglang_dp_sizeint | None4Data-parallel size for the engines.
sglang_ep_sizeint | None4Expert-parallel size for MoE models.
sglang_enable_dp_lm_headboolTrueData-parallel LM head (pairs with DP attention).
sglang_disable_custom_all_reduceboolFalseFall back to NCCL all-reduce instead of sglang’s custom kernel.
sglang_cuda_graph_bslist[int] | None[1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256]Batch sizes to capture CUDA graphs for.
sglang_max_running_requestsint | None512Cap on concurrent in-flight requests per engine.
sglang_tool_call_parserstr | NoneNoneParser for tool-call output, e.g. "qwen25".
sglang_reasoning_parserstr | NoneNoneParser for reasoning/thinking output.
FieldTypeDefaultDescription
sglang_cuda_graph_backend_prefillstr | NoneNone
hf_checkpointstr"Qwen/Qwen3.6-35B-A3B"
pipeline_model_parallel_sizeint2
context_parallel_sizeint2
expert_model_parallel_sizeint4
expert_tensor_parallel_sizeint1
sglang_speculative_algorithmstr | None"EAGLE"
sglang_speculative_num_stepsint | None3
sglang_speculative_eagle_topkint | None1
sglang_speculative_num_draft_tokensint | None4
sglang_mamba_scheduler_strategystr"extra_buffer"
moe_token_dispatcher_typestr"flex"
moe_enable_deepepboolTrue
optimizer_cpu_offloadboolTrue
overlap_cpu_optimizer_d2h_h2dboolTrue
use_precision_aware_optimizerboolTrue
attention_backendstr"flash"

cli_args(self, dataset: 'DatasetConfig | None' = None, model: 'ModelConfig | None' = None) -> list[str]

Section titled “cli_args(self, dataset: 'DatasetConfig | None' = None, model: 'ModelConfig | None' = None) -> list[str]”

get_base_recipe(model_config: modal_training_gym.common.models.base.ModelConfig) -> 'SlimeRecipe'

Section titled “get_base_recipe(model_config: modal_training_gym.common.models.base.ModelConfig) -> 'SlimeRecipe'”

validate_model_parallelism(self, model: 'ModelConfig') -> None

Section titled “validate_model_parallelism(self, model: 'ModelConfig') -> None”

Source: modal_training_gym/train_recipes/slime_recipe/qwen3_6_35b.py