Zhichenzzz · Oct 2, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026
diff --git a/docs/user-guide/training-backend.md b/docs/user-guide/training-backend.md
@@ -468,39 +468,42 @@ the config a torchtitan user would write by hand.
 --titan-model-name qwen3 \
 --titan-model-flavor 30B-A3B \
 --hf-checkpoint /root/models/Qwen3-30B-A3B \
---titan-seq-len 16384
+--seq-length 16384
 ```
 
 `--titan-model-name` is a package under `torchtitan/models/`, and `--titan-model-flavor` one
 of the sizes that package registers. `--hf-checkpoint` supplies the weights and tokenizer;
 torchtitan's own state-dict adapter converts them, so there is no offline conversion step.
 
-`--titan-seq-len` sizes the rotary tables and the buffers pipeline stages exchange, so it has
+`--seq-length` sizes the rotary tables and the buffers pipeline stages exchange, so it has
 to be at least as long as the longest sequence you will train on — prompt plus response.
 Miles rejects a value that leaves no room for a prompt rather than letting it fail inside a
 kernel later.
 
 ### 2. Choosing the parallelism
 
 ```bash
---titan-tensor-parallel-degree 2 \
---titan-pipeline-parallel-degree 2 \
---titan-context-parallel-degree 2 \
---titan-expert-parallel-degree 2 \
---titan-data-parallel-replicate-degree 1
+--tensor-model-parallel-size 2 \
+--pipeline-model-parallel-size 2 \
+--context-parallel-size 2 \
+--expert-model-parallel-size 2 \
+--dp-replicate-size 1
 ```
 
-The FSDP shard degree is not a flag: torchtitan infers it from what the other degrees leave
-over, so these five settings plus the GPU count fully determine the layout. All of them
-compose — tensor, pipeline, context, expert and FSDP have been run together on one job.
+These are the same flags Megatron takes, so a recipe moves between the two backends without
+renaming its parallelism. The FSDP shard degree is not a flag: torchtitan infers it from what
+the other degrees leave over, so these five settings plus the GPU count fully determine the
+layout. All of them compose — tensor, pipeline, context, expert and FSDP have been run together
+on one job. `--recompute-granularity full` and `--bf16` are accepted with Megatron's meaning as
+well (`--gradient-checkpointing` stays as the FSDP-side spelling).
 
 Two notes on how they behave here:
 
 - **Context parallelism is internal.** The trainer shards the sequence for attention and
   gathers the logits back before the loss sees them, so the RL loss and its metrics behave
   exactly as at `cp=1` and cost the same memory.
 - **Pipeline parallelism needs one shape for the whole run**, so every microbatch is padded
-  to `--titan-seq-len`. Weight-tied flavors (qwen3 0.6B / 1.7B / 4B) cannot be pipelined —
+  to `--seq-length`. Weight-tied flavors (qwen3 0.6B / 1.7B / 4B) cannot be pipelined —
   torchtitan refuses, by design.
 
 ### 3. Fitting it in memory

diff --git a/examples/infra_features/train_infer_mismatch_helper/mis.py b/examples/infra_features/train_infer_mismatch_helper/mis.py
@@ -333,7 +333,7 @@ def compute_mis_weights_with_cp(
         is_metrics: The metrics for the importance sampling weights, a dict of flattened tensors.
     """
     # Lazy import to avoid importing Megatron dependencies when only `compute_mis_weights` is used.
-    from miles.backends.training_utils.cp_utils import all_gather_with_cp, slice_log_prob_with_cp
+    from miles.backends.training_utils.data.context_parallel import all_gather_with_cp, slice_log_prob_with_cp
 
     qkv_format = args.qkv_format
     max_seq_lens = kwargs.get("max_seq_lens", None)

@@ -18,7 +18,7 @@
 
 import torch.nn as nn
 
-from miles.backends.training_utils.torch_native.routing_replay import (
+from miles.backends.training_utils.replay.routing_replay import (
     FALLTHROUGH,
     RECORD,
     REPLAY_BACKWARD,

@@ -34,8 +34,11 @@ class FSDPArgs:
     wandb_run_name: str | None = None
 
     # Precision
-    gradient_checkpointing: bool = False
+    bf16: bool = True
     fp16: bool = False
+    # Activation recompute: --recompute-granularity full is Megatron's spelling of the same switch
+    gradient_checkpointing: bool = False
+    recompute_granularity: str | None = None
     keep_fp32_master: bool = True
 
     # FSDP configuration
@@ -126,6 +129,12 @@ def load_args_from_parser(parser: argparse.ArgumentParser):
         parser.set_defaults(**data)
         args = parser.parse_args()
     args.bf16 = not args.fp16
+    if args.recompute_granularity is not None:
+        if args.recompute_granularity != "full":
+            raise ValueError(
+                f"--recompute-granularity {args.recompute_granularity!r}: this backend recomputes whole layers only; use full"
+            )
+        args.gradient_checkpointing = True
     return args
 
 

@@ -43,16 +43,16 @@
 
 from ...utils.profile_utils import TrainProfiler
 from ...utils.tensor_backper import TensorBackuper
-from ..training_utils.data import DataIterator, get_data_iterator, get_num_rollouts, get_rollout_data
-from ..training_utils.log_utils import log_cpu_memory, log_perf_data, log_rollout_data
-from ..training_utils.loss import (
+from ..training_utils.data.rollout import DataIterator, get_data_iterator, get_num_rollouts, get_rollout_data
+from ..training_utils.loss.objective import (
     compute_advantages_and_returns,
     get_log_probs_and_entropy,
     get_values,
     log_train_advantage_computation_event,
 )
+from ..training_utils.metrics.log_utils import log_cpu_memory, log_perf_data, log_rollout_data
 from ..training_utils.parallel import get_parallel_state
-from ..training_utils.replay_data import fill_replay_data, register_replay_list_sequential
+from ..training_utils.replay.replay_data import fill_replay_data, register_replay_list_sequential
 from .checkpoint import load_checkpoint
 from .ft.checkpoint_transfer import recv_ckpt
 from .ft.checkpoint_transfer import send_ckpt as _send_ckpt

@@ -11,7 +11,7 @@
 from miles.utils.ft_utils.process_group_utils import GeneralPGUtil, GroupInfo, collective_bool_and
 from miles.utils.tracking_utils.structured_log import log_structured
 
-from ...training_utils.log_utils import aggregate_train_losses
+from ...training_utils.metrics.log_utils import aggregate_train_losses
 from ...training_utils.parallel import ParallelState
 
 if TYPE_CHECKING:

@@ -4,7 +4,7 @@
 from miles.backends.megatron_utils.lora import checkpoint as lora_checkpoint
 from miles.backends.megatron_utils.lora import model as lora_model
 from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer
-from miles.backends.training_utils.data import get_rollout_data
+from miles.backends.training_utils.data.rollout import get_rollout_data
 from miles.utils.lora.utils import AdapterSpec
 from miles.utils.ray_utils import Box
 from miles.utils.tracking_utils.structured_log import with_logs

@@ -11,8 +11,8 @@
     step_slot_optimizers,
 )
 from miles.backends.megatron_utils.model import run_forward_backward_pass, setup_train_iteration_config
-from miles.backends.training_utils.data import get_data_iterator
-from miles.backends.training_utils.log_utils import aggregate_train_losses
+from miles.backends.training_utils.data.rollout import get_data_iterator
+from miles.backends.training_utils.metrics.log_utils import aggregate_train_losses
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase
 from miles.utils.types import RolloutBatch

@@ -29,7 +29,7 @@
 from miles.backends.megatron_utils.ft.types import TrainStepOutcome
 from miles.backends.megatron_utils.local_weight_checksum import dump_local_weight_checksums
 from miles.backends.megatron_utils.optimizer_state_reset import reset_optimizer_states
-from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
+from miles.backends.training_utils.data.sampling_mask import get_rollout_sampling_masks
 from miles.backends.training_utils.weight_update.snapshot_publisher import SnapshotPublisher
 from miles.utils.audit_utils.witness.allocator import WitnessInfo
 from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale
@@ -40,10 +40,10 @@
 from miles.utils.tracking_utils.structured_log import log_structured
 
 from ...utils.misc import filter_keys
-from ..training_utils.ci_utils import check_grad_norm, check_kl
-from ..training_utils.data import DataIterator, get_batch
-from ..training_utils.log_utils import aggregate_forward_results, aggregate_train_losses, log_train_step
-from ..training_utils.loss import loss_function
+from ..training_utils.data.rollout import DataIterator, get_batch
+from ..training_utils.loss.checks import check_grad_norm, check_kl
+from ..training_utils.loss.objective import loss_function
+from ..training_utils.metrics.log_utils import aggregate_forward_results, aggregate_train_losses, log_train_step
 from ..training_utils.parallel import get_parallel_state
 from .checkpoint import load_checkpoint, save_checkpoint, save_checkpoint_with_lora
 from .ci_utils import (

@@ -12,8 +12,8 @@
 from miles.backends.torchtitan_utils.routing_replay import install as install_routing_replay
 from miles.backends.torchtitan_utils.trainer import TitanTrainer
 from miles.backends.training_utils.parallel import get_parallel_state, set_parallel_state
+from miles.backends.training_utils.replay.routing_replay import enable as enable_routing_replay
 from miles.backends.training_utils.torch_native.actor import TorchNativeTrainRayActor
-from miles.backends.training_utils.torch_native.routing_replay import enable as enable_routing_replay
 from miles.utils.context_utils import with_defer
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
 from miles.utils.memory_utils import clear_memory

@@ -11,15 +11,11 @@ class TorchtitanArgs(FSDPArgs):
     titan_model_name: str = "qwen3"
     titan_model_flavor: str = "0.6B"
 
-    # sizes torchtitan's rotary tables; must cover the longest prompt + response
-    titan_seq_len: int = 4096
+    seq_length: int = 4096
 
-    # torchtitan's ParallelismConfig fields, same names, defaults and semantics
-    titan_data_parallel_replicate_degree: int = 1
-    titan_tensor_parallel_degree: int = 1
-    titan_pipeline_parallel_degree: int = 1
-    titan_context_parallel_degree: int = 1
-    titan_expert_parallel_degree: int = 1
+    tensor_model_parallel_size: int = 1
+    pipeline_model_parallel_size: int = 1
+    expert_model_parallel_size: int = 1
 
 
 def build_torchtitan_parser(extra_args_provider=None):
@@ -34,19 +30,19 @@ def validate_torchtitan_args(args) -> None:
     torch_version = tuple(int(part) for part in torch.__version__.split(".")[:2])
     if torch_version < (2, 13):
         raise ValueError(f"the torchtitan backend needs torch>=2.13; this environment runs {torch.__version__}")
-    if args.titan_context_parallel_degree != 1 and args.titan_model_name == "qwen3_5":
+    if args.context_parallel_size != 1 and args.titan_model_name == "qwen3_5":
         raise ValueError("torchtitan does not support context parallelism for qwen3_5")
 
     if args.rollout_max_context_len is not None:
-        if args.titan_seq_len < args.rollout_max_context_len:
+        if args.seq_length < args.rollout_max_context_len:
             raise ValueError(
-                f"--titan-seq-len {args.titan_seq_len} is shorter than "
+                f"--seq-length {args.seq_length} is shorter than "
                 f"--rollout-max-context-len {args.rollout_max_context_len}: torchtitan builds its "
                 "rotary embeddings for the former, so a longer sequence would index past them"
             )
-    elif args.titan_seq_len <= args.rollout_max_response_len:
+    elif args.seq_length <= args.rollout_max_response_len:
         raise ValueError(
-            f"--titan-seq-len {args.titan_seq_len} leaves no room for a prompt ahead of "
+            f"--seq-length {args.seq_length} leaves no room for a prompt ahead of "
             f"--rollout-max-response-len {args.rollout_max_response_len}: torchtitan builds its "
             "rotary embeddings for the former, and a prompt-plus-response beyond them asserts "
             "inside the rope kernel"

@@ -44,11 +44,11 @@ def build_trainer_config(args: Namespace, *, hf_assets_path: str, lr_total_steps
         raise ValueError(f"torchtitan backend supports --optimizer adam, got {args.optimizer!r}")
 
     ties_embeddings = load_hf_config(hf_assets_path).tie_word_embeddings
-    if ties_embeddings and args.titan_pipeline_parallel_degree > 1:
+    if ties_embeddings and args.pipeline_model_parallel_size > 1:
         raise ValueError(
             "the checkpoint ties lm_head to the embedding, which torchtitan cannot do across pipeline "
             "stages: it would train a separate lm_head that the HF export then has no tensor to ship "
-            "into. Use --titan-pipeline-parallel-degree 1 or an untied checkpoint."
+            "into. Use --pipeline-model-parallel-size 1 or an untied checkpoint."
         )
 
     config = Trainer.Config()
@@ -62,16 +62,16 @@ def build_trainer_config(args: Namespace, *, hf_assets_path: str, lr_total_steps
         args.save or tempfile.mkdtemp(prefix="miles-torchtitan-"), "torchtitan", dump_subdir
     )
 
-    config.parallelism.data_parallel_replicate_degree = args.titan_data_parallel_replicate_degree
-    config.parallelism.tensor_parallel_degree = args.titan_tensor_parallel_degree
-    config.parallelism.pipeline_parallel_degree = args.titan_pipeline_parallel_degree
-    config.parallelism.context_parallel_degree = args.titan_context_parallel_degree
-    config.parallelism.expert_parallel_degree = args.titan_expert_parallel_degree
+    config.parallelism.data_parallel_replicate_degree = args.dp_replicate_size
+    config.parallelism.tensor_parallel_degree = args.tensor_model_parallel_size
+    config.parallelism.pipeline_parallel_degree = args.pipeline_model_parallel_size
+    config.parallelism.context_parallel_degree = args.context_parallel_size
+    config.parallelism.expert_parallel_degree = args.expert_model_parallel_size
     config.parallelism.pipeline_parallel_microbatch_size = 1
     parallel_dims = parallel_dims_from_config(config.parallelism)
     dp_size = parallel_dims.dp_replicate * parallel_dims.dp_shard
 
-    config.training.seq_len = args.titan_seq_len
+    config.training.seq_len = args.seq_length
     if parallel_dims.pp_enabled and args.global_batch_size % (dp_size * args.micro_batch_size):
         raise ValueError(
             f"--global-batch-size {args.global_batch_size} must be a multiple of dp * micro_batch_size "

@@ -8,12 +8,7 @@
 import torch.nn.functional as F
 from torch.distributed.tensor import DTensor
 
-from miles.backends.training_utils.torch_native.routing_replay import (
-    FALLTHROUGH,
-    REPLAY_BACKWARD,
-    REPLAY_FORWARD,
-    stage,
-)
+from miles.backends.training_utils.replay.routing_replay import FALLTHROUGH, REPLAY_BACKWARD, REPLAY_FORWARD, stage
 from miles.utils.replay_base import routing_replay_manager
 
 logger = logging.getLogger(__name__)

@@ -35,7 +35,7 @@ def padded_length(self, n_tokens: int) -> int:
             target = self.config.training.seq_len
             if n_tokens > target:
                 raise ValueError(
-                    f"packed microbatch of {n_tokens} tokens exceeds --titan-seq-len "
+                    f"packed microbatch of {n_tokens} tokens exceeds --seq-length "
                     f"{target}, which is the fixed shape PP stages exchange"
                 )
             return target

@@ -5,7 +5,7 @@
 import torch.distributed as dist
 import torch.nn.functional as F
 
-from .parallel import get_parallel_state
+from ..parallel import get_parallel_state
 
 logger = logging.getLogger(__name__)
 

@@ -15,8 +15,8 @@
 
 from miles.utils.types import RolloutBatch
 
-from .cp_utils import all_gather_with_cp, slice_log_prob_with_cp
-from .parallel import get_parallel_state
+from ..parallel import get_parallel_state
+from .context_parallel import all_gather_with_cp, slice_log_prob_with_cp
 
 logger = logging.getLogger(__name__)
 

@@ -7,17 +7,16 @@
 import torch.nn.functional as F
 
 from miles.utils.audit_utils.witness.allocator import WitnessInfo
-from miles.utils.data import get_minimum_num_micro_batch_size
+from miles.utils.data import get_minimum_num_micro_batch_size, process_rollout_data
 from miles.utils.ft_utils.process_group_utils import GeneralPGUtil
 from miles.utils.object_store import ObjectStoreGetResult
+from miles.utils.ray_utils import Box
 from miles.utils.seqlen_balancing import get_seqlen_balanced_partitions
 from miles.utils.types import RolloutBatch
 
-from ...utils.data import process_rollout_data
-from ...utils.ray_utils import Box
-from .cp_utils import slice_log_prob_with_cp, slice_with_cp
-from .mm_data import expand_multimodal_rollout_data_in_place
-from .parallel import get_parallel_state
+from ..parallel import get_parallel_state
+from .context_parallel import slice_log_prob_with_cp, slice_with_cp
+from .multimodal import expand_multimodal_rollout_data_in_place
 
 logger = logging.getLogger(__name__)
 

diff --git a/miles/backends/training_utils/ci_utils.py → miles/backends/training_utils/loss/checks.py b/miles/backends/training_utils/ci_utils.py → miles/backends/training_utils/loss/checks.py
@@ -2,8 +2,8 @@
 
 import torch
 
-from miles.backends.training_utils.cp_utils import get_logits_and_tokens_offset_with_cp
-from miles.backends.training_utils.loss_hub.math_utils import (
+from miles.backends.training_utils.data.context_parallel import get_logits_and_tokens_offset_with_cp
+from miles.backends.training_utils.loss.hub.math_utils import (
     get_advantages_and_returns_batch,
     get_grpo_returns,
     get_reinforce_plus_plus_baseline_advantages,

@@ -3,10 +3,13 @@
 
 import torch
 
-from miles.backends.training_utils.cp_utils import allgather_cp_redistribute, get_logits_and_tokens_offset_with_cp
-from miles.backends.training_utils.loss_hub.math_utils import calculate_log_probs_and_entropy
+from miles.backends.training_utils.data.context_parallel import (
+    allgather_cp_redistribute,
+    get_logits_and_tokens_offset_with_cp,
+)
+from miles.backends.training_utils.data.sampling_mask import build_local_sampling_mask
+from miles.backends.training_utils.loss.hub.math_utils import calculate_log_probs_and_entropy
 from miles.backends.training_utils.parallel import get_parallel_state
-from miles.backends.training_utils.sampling_mask import build_local_sampling_mask
 from miles.utils.sampling_mask import RolloutSamplingMask