Moonlight_16B_A3B_Recipe

from modal_training_gym.train_recipes.miles_recipe.moonlight_16b_a3b import Moonlight_16B_A3B_Recipe

Moonlight-16B-A3B DAPO recipe for 1 node with 8 H100 GPUs.

Attributes

docker_image str

Registry reference for the Miles image every container runs. Default: "radixark/miles:dev-202608120325"

gpu_type str

Modal GPU type for every node. Default: "H100"

memory int | tuple[int, int] | None

Modal Function memory request/limit in MiB.

cpu float | tuple[float, float] | None

Modal Function CPU request/limit in cores per container.

cloud str | None

Modal cloud provider to pin the cluster to.

region str | None

Modal region to pin the cluster to.

name str

Modal app title. The launcher derives it from the class when empty. Default: ""

app_tags dict

Extra tags merged into the Modal app metadata for the dashboard. Default: {}

image_overlay Callable[[Image], Image] | None

Function that modifies the Modal image.

image_run_commands list[str]

Extra shell commands run while building the image. Default: []

image_env dict[str, str]

Extra env vars baked into the image. Default: {}

local_miles str | None

Local Miles checkout mounted over the image copy without rebuilding it.

patch_files list[str]

Local patch scripts applied to Miles/Megatron sources at image build. Default: []

substep_timing Literal['auto', 'off']

Default: "auto"

environment dict[str, str]

Training-container environment variables such as Megatron PYTHONPATH and NCCL settings.

async_mode bool

Run Miles' train_async.py so rollout generation and training overlap. Default: False

miles_model_script str

Script in the Miles repository sourced for MODEL_ARGS instead of model-architecture flags. Default: ""

miles_model_name str

Name accepted by Miles' model_args_utils.py. Default: ""

source_hf_checkpoint str | None

Source checkpoint when it differs from the model's own.

megatron_conversion_hf_checkpoint str | None

HF weights used for the HF→Megatron conversion instead of the model's own.

metrics MetricConfig | None

Metric tracker settings; expands to Miles' W&B-compatible flags.

actor_num_nodes int

Megatron actor nodes. Default: 1

actor_num_gpus_per_node int

GPUs per actor node. Default: 8

rollout_num_gpus int | None

Rollout-engine GPUs when disaggregated; None lets the resolver size it.

colocate bool

Trainer and rollout engines share GPUs; False gives each its own. Default: True

use_critic bool

Train a separate critic model for PPO. GRPO does not use one. Default: False

critic_num_nodes int | None

Nodes for the critic when use_critic is set.

critic_num_gpus_per_node int | None

GPUs per critic node.

hf_checkpoint str

Checkpoint trained from; normally set from the attached ModelConfig. Default: ""

save str

Checkpoint output directory on the mounted /checkpoints volume. Default: "/checkpoints"

load str

Directory to resume from; empty starts from the converted HF weights. Default: ""

ref_load str

Checkpoint read by the reference model for KL terms. Default: "/checkpoints/Moonlight-16B-A3B-Instruct_torch_dist"

megatron_to_hf_mode str

Export mode for saved Megatron checkpoints; empty disables the export. Default: ""

model_name str

Default: ""

save_interval int

Save a checkpoint every N rollout steps. Default: 20

no_save_optim bool

Omit optimizer state from checkpoints. The resulting checkpoints cannot resume the optimizer exactly. Default: False

conversion_tensor_model_parallel_size int | None

conversion_pipeline_model_parallel_size int | None

conversion_expert_model_parallel_size int | None

conversion_expert_tensor_parallel_size int | None

convert_ephemeral_disk_mb int | None

use_fault_tolerance bool

Enable Miles' fault tolerance to recover from worker failures. Default: False

rollout_health_check_interval int

Seconds between rollout engine /health_generate checks. Default: 30

rollout_health_check_timeout int

Seconds to wait for /health_generate before killing the engine. Default: 30

rollout_health_check_first_wait int

Initial health-check delay in seconds before checking /health_generate. DeepGEMM compilation may require a longer delay. Default: 0

update_weight_buffer_size int | None

Byte size of the buffer broadcasting updated weights to the engines.

num_rollout int

Training and rollout steps for the run. Default: 3000

start_rollout_id int | None

Rollout step to start counting from. None continues from the iteration stored in load; TrainConfig(checkpoint=...) sets 0 so num_rollout counts the steps this run performs.

rollout_batch_size int

Prompts per rollout step, each expanded into a group of responses. Default: 128

n_samples_per_prompt int

Responses sampled per prompt as one GRPO group. Default: 8

rollout_max_response_len int

Max generated tokens per sample. Default: 4096

rollout_temperature float

Sampling temperature for rollout generation. Default: 1.0

rollout_shuffle bool

Shuffle the prompt dataset between epochs. Default: True

rollout_top_p float

Nucleus-sampling top-p for rollout generation. Default: 1.0

rollout_stop_token_ids list[int] | None

Extra token ids that terminate generation.

rollout_num_gpus_per_engine int

GPUs and tensor-parallel size per SGLang engine. Default: 8

use_miles_router bool

Route rollout requests through Miles' router instead of directly to engines. Default: True

rollout_top_k int | None

Top-k for rollout generation; None leaves Miles' own default.

use_rollout_routing_replay bool

Reuse the rollout's MoE expert routing in training. Default: False

tensor_model_parallel_size int

Megatron tensor-parallel size for the actor. Default: 4

pipeline_model_parallel_size int

Megatron pipeline-parallel size for the actor. Default: 1

context_parallel_size int | None

Megatron context-parallel size; multiplies the effective context length. Default: 1

expert_model_parallel_size int | None

Expert-parallel size for MoE; must divide the model's num_experts. Default: 8

expert_tensor_parallel_size int | None

Tensor-parallel size within each expert. Default: 1

decoder_last_pipeline_num_layers int | None

Layers placed on the last pipeline stage, to rebalance an uneven split.

sequence_parallel bool

Megatron sequence parallelism. Requires tensor parallelism greater than one. Default: True

train_backend str

Training backend. Default: "megatron"

global_batch_size int | None

Training samples per optim step.

lr float

Learning rate. Default: 1e-06

lr_decay_style str

Learning-rate schedule. Default: "constant"

weight_decay float

Weight decay. Default: 0.1

adam_beta1 float

Adam beta1. Default: 0.9

adam_beta2 float

Adam beta2. Default: 0.98

optimizer str

Optimizer name. Default: "adam"

use_distributed_optimizer bool

Shard optimizer state across data-parallel ranks with Megatron's distributed optimizer. Default: False

optimizer_cpu_offload bool

Keep optimizer state on CPU to reduce GPU memory use at the cost of slower steps. Default: True

overlap_cpu_optimizer_d2h_h2d bool

Overlap the offloaded optimizer's device↔host copies with compute. Default: True

use_precision_aware_optimizer bool

Use Megatron's precision-aware optimizer with lower-precision state. Default: True

lora_rank int | None

LoRA rank; None trains full weights.

lora_alpha int | None

LoRA scaling factor.

lora_dropout float | None

Dropout applied to LoRA layers.

target_modules str | None

Comma-separated module names LoRA adapters attach to.

experts_shared_outer_loras bool

Share one outer LoRA across MoE experts instead of one per expert. Default: False

lora_base_cpu_backup bool

Keep frozen base weights on CPU to free GPU memory. Default: False

no_gradient_accumulation_fusion bool

Disable fused gradient accumulation for incompatible LoRA paths. Default: False

sglang_lora_backend str | None

SGLang LoRA kernel backend.

sglang_lora_use_virtual_experts bool

Serve MoE LoRA adapters as virtual experts in sglang. Default: False

use_tis bool

Correct rollout and trainer mismatch with truncated importance sampling. Default: False

advantage_estimator str

Advantage estimator. Default: "grpo"

eps_clip float

PPO clip lower bound. Default: 0.2

eps_clip_high float

Upper PPO clip bound for asymmetric DAPO clipping. Default: 0.28

kl_loss_type str

KL formulation. Default: "low_var_kl"

kl_loss_coef float

Coefficient of the KL loss term. Default: 0.0

kl_coef float

KL penalty coefficient applied in the reward. Default: 0.0

entropy_coef float

Entropy bonus coefficient. Default: 0.0

use_kl_loss bool

Add a per-token KL loss term against the reference model. Default: True

calculate_per_token_loss bool

Average the loss over tokens instead of over samples. Default: False

rm_type str | None

Built-in reward function name. Leave unset for a custom reward. Default: "math"

over_sampling_batch_size int | None

Extra DAPO prompts sampled to replace filtered groups. Default: 256

dynamic_sampling_filter_path str | None

Import path of the predicate deciding which sample groups to keep. Default: "miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std"

balance_data bool

Rebalance kept samples across data-parallel ranks. Default: True

attention_dropout float

Attention dropout probability. Default: 0.0

hidden_dropout float

Hidden-layer dropout probability. Default: 0.0

attention_softmax_in_fp32 bool

Compute attention softmax in fp32. Default: True

accumulate_allreduce_grads_in_fp32 bool

Accumulate and all-reduce gradients in fp32. Default: True

attention_backend str | None

Megatron attention kernel backend.

no_check_for_nan_in_loss_and_grad bool

Skip the NaN check on loss and gradients to avoid a synchronization per step. Default: False

recompute_granularity str | None

Activation recomputation granularity: "full" or "selective". Default: "full"

recompute_method str | None

Recomputation method: "uniform" or "block". Default: "uniform"

recompute_num_layers int | None

Layers per recomputation chunk. Default: 1

qkv_format str

QKV layout for the Megatron backend: "thd" or "bshd". Default: "thd"

use_dynamic_batch_size bool

Pack samples up to max_tokens_per_gpu instead of a fixed micro batch. Default: True

micro_batch_size int | None

Fixed micro-batch size when dynamic batching is off; None leaves Miles' own default.

max_tokens_per_gpu int

Token budget per GPU per micro-batch when dynamic batching is on. Default: 8192

eval_interval int | None

Run eval every N rollout steps; None disables eval. Default: 20

n_samples_per_eval_prompt int

Responses sampled per eval prompt. Default: 8

eval_max_response_len int

Max generated tokens per eval sample. Default: 4096

eval_top_p float

Nucleus-sampling top-p for eval generation. Default: 1.0

eval_config dict | str | None

Evaluation defaults and datasets written to --eval-config as YAML.

skip_eval_before_train bool

Skip the eval pass before the first train step. Default: False

sglang_mem_fraction_static float

Fraction of GPU memory sglang reserves for weights + KV cache. Default: 0.7

sglang_enable_dp_attention bool

Enable data-parallel attention across engine ranks. Default: False

sglang_dp_size int | None

Data-parallel size for the engines.

sglang_ep_size int | None

Expert-parallel size for MoE models.

sglang_enable_dp_lm_head bool

Data-parallel LM head paired with DP attention. Default: False

sglang_disable_custom_all_reduce bool

Fall back to NCCL all-reduce instead of sglang's custom kernel. Default: False

sglang_cuda_graph_bs list[int]

Batch sizes to capture CUDA graphs for.

sglang_attention_backend str | None

SGLang attention kernel backend. The server selects one when unset.

sglang_disable_cuda_graph bool

Run the engines in eager mode instead of capturing CUDA graphs. Default: False

sglang_disable_overlap_schedule bool

Disable sglang's overlapped scheduler. Default: False

sglang_disable_radix_cache bool

Disable prefix (radix) caching across requests. Default: False

no_offload_train bool

Keep training weights and optimizer resident between rollout and training phases for colocated runs. Default: False

no_offload_rollout bool

Keep the rollout engines resident instead of offloading them. Default: False

sglang_moe_runner_backend str | None

SGLang MoE GEMM runner. The server selects one when unset.

sglang_max_running_requests int | None

Cap on concurrent in-flight requests per engine.

sglang_server_concurrency int | None

Cap on concurrent requests Miles sends to each engine.

sglang_tool_call_parser str | None

Tool-call output parser.

sglang_reasoning_parser str | None

Parser for reasoning/thinking output.

extra_config dict | None

Custom configuration written to YAML at --custom-config-path. Keys become Miles arguments and override same-named fields.

sglang_config dict | str | None

SGLang engine settings written to --sglang-config as YAML.

apply_chat_template_kwargs str | dict

Keyword arguments for tokenizer apply_chat_template, passed as JSON. Default: ""

train_env_vars dict | str | None

Env vars for the training processes, passed as inline JSON.

multimodal_keys dict | str | None

Multimodal dataset columns passed as JSON.

custom_rm_function Callable | None

Reward callable shipped by value as Miles' custom_rm_path.

custom_generate_function Callable | None

Custom Miles generation step shipped by value.

custom_reward_post_process_function Callable | None

Function applied to rewards after generation and shipped by value.

rollout_function Callable | str | None

Custom rollout loop passed through --rollout-function-path.

custom_rollout_log_function Callable | str | None

Function called with each rollout's data after dashboard and phase reporting.

custom_eval_rollout_log_function Callable | str | None

Function called with each evaluation rollout's data.

custom_megatron_before_log_prob_hook Callable | str | None

Hook run in the Megatron trainer before log-prob computation.

custom_megatron_before_train_step_hook Callable | str | None

Hook run in the Megatron trainer before each train step.

train_function_kwargs dict[str, Any]

Additional Modal Function keyword arguments for the training function. Default: {}

capture_trace bool

Attach sampled per-request execution traces to recorded rollouts. Default: False

trace_sample_limit int

Maximum traced samples per rollout when capture_trace is enabled. Default: 16

num_steps_per_rollout int | None

Default: 4

moe_enable_deepep bool

Default: True

moe_token_dispatcher_type str

Default: "flex"

cli_args

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

download_model

download_model() -> None

get_base_recipe

get_base_recipe(model_config: ModelConfig) -> MilesRecipe | None

Return the model preset for model_config.

Call this explicitly. TrainConfig uses the recipe it receives without applying a preset.

Returns

The model preset, or None when no preset is registered.

Raises

• TrainingGymConfigError: The recipe accepts only registered models and model_config is not registered.

gpu_allocation

gpu_allocation: GpuAllocation

post_process_data

post_process_data() -> None

post_process_model

post_process_model() -> None

total_nodes

total_nodes: int

validate_model_parallelism

validate_model_parallelism(model: ModelConfig) -> None

Validate the model’s parallelism settings.