Inkling_Small_LoRA_Recipe

from modal_training_gym.train_recipes.miles_recipe.inkling import Inkling_Small_LoRA_Recipe

Inkling-Small rank-32 LoRA GRPO recipe for 4 nodes with 8 H200 GPUs each.

Attributes

docker_image str

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

gpu_type str

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

memory tuple[int, int]

Modal Function memory request/limit in MiB. Default: (1024, 2097152)

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: "scripts/models/inkling-small.sh"

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: 4

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/Inkling-Small_torch_dist"

megatron_to_hf_mode str

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

model_name str

Default: "inkling"

save_interval int

Save a checkpoint every N rollout steps. Default: 10

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

Default: 8

conversion_pipeline_model_parallel_size int

Default: 1

conversion_expert_model_parallel_size int

Default: 8

conversion_expert_tensor_parallel_size int

Default: 1

convert_ephemeral_disk_mb int | None

Default: 1048576

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: 100

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: 64

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: 2048

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: 16

use_miles_router bool

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

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: True

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: 8

context_parallel_size int | None

Megatron context-parallel size; multiplies the effective context length.

expert_model_parallel_size int

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

expert_tensor_parallel_size int

Tensor-parallel size within each expert. Default: 1

decoder_last_pipeline_num_layers int

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

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

Training samples per optim step. Default: 128

lr float

Learning rate. Default: 0.0002

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: True

optimizer_cpu_offload bool

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

overlap_cpu_optimizer_d2h_h2d bool

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

use_precision_aware_optimizer bool

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

lora_rank int | None

LoRA rank; None trains full weights. Default: 32

lora_alpha int | None

LoRA scaling factor. Default: 32

lora_dropout float | None

Dropout applied to LoRA layers.

target_modules str | None

Comma-separated module names LoRA adapters attach to. Default: "all-linear"

experts_shared_outer_loras bool

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

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. Default: "triton"

sglang_lora_use_virtual_experts bool

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

use_tis bool

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

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: False

calculate_per_token_loss bool

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

rm_type str

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.

dynamic_sampling_filter_path str | None

Import path of the predicate deciding which sample groups to keep.

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: True

recompute_granularity str

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

recompute_method str

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

recompute_num_layers int

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: 4096

eval_interval int | None

Run eval every N rollout steps; None disables eval.

n_samples_per_eval_prompt int

Responses sampled per eval prompt. Default: 4

eval_max_response_len int

Max generated tokens per eval sample. Default: 16384

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: True

sglang_mem_fraction_static float

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

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

Expert-parallel size for MoE models. Default: 16

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: True

sglang_cuda_graph_bs list[int] | None

Batch sizes to capture CUDA graphs for.

sglang_attention_backend str

SGLang attention kernel backend. The server selects one when unset. Default: "fa4"

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: True

no_offload_rollout bool

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

sglang_moe_runner_backend str

SGLang MoE GEMM runner. The server selects one when unset. Default: "triton"

sglang_max_running_requests int

Cap on concurrent in-flight requests per engine. Default: 32

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

modality Literal['text', 'vision']

Default: "text"

transformer_impl str

Default: "transformer_engine"

bf16 bool

Default: True

no_bias_dropout_fusion bool

Default: True

distributed_timeout_minutes int

Default: 30

eps_clip_c float

Default: 3.0

sglang_mamba_scheduler_strategy str

Default: "extra_buffer"

sglang_enable_multimodal bool

Default: True

sglang_context_length int

Default: 4096

custom_model_provider_path str | None

sglang_max_loras_per_batch int

Default: 1

sglang_max_lora_rank int

Default: 32

sglang_max_total_tokens int

Default: 320000

sglang_cuda_graph_max_bs int

Default: 64

sglang_max_mamba_cache_size int

Default: 256

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.