Zhichenzzz · Oct 2, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026
diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md
@@ -107,7 +107,7 @@ from the trainer loop and uses a continuously-running worker.
 
 | You want to … | Edit |
 |---|---|
-| Add a new RL algorithm | `miles/backends/training_utils/loss.py` and `loss_hub/`, plus the enum in `miles/utils/arguments.py` |
+| Add a new RL algorithm | `miles/backends/training_utils/loss/objective.py` and `loss/hub/`, plus the enum in `miles/utils/arguments.py` |
 | Add a new built-in reward type | `miles/rollout/rm_hub/` (the `rm_type` dispatch lives in its `__init__.py`) |
 | Add a new built-in filter | `miles/rollout/filter_hub/` |
 | Support a new architecture on Megatron | `miles_plugins/models/<model>.py` + a bridge in `miles_plugins/mbridge/` |
@@ -155,7 +155,7 @@ If you have 30 minutes and want to understand Miles end-to-end:
 
 1. `train.py` — the loop, top-to-bottom.
 2. `miles/rollout/sglang_rollout.py:generate_rollout` — how prompts become samples.
-3. `miles/backends/training_utils/loss.py` — the loss and advantage computation.
+3. `miles/backends/training_utils/loss/objective.py` — the loss and advantage computation.
 4. `miles/router/router.py` — the FastAPI proxy.
 5. `miles/backends/training_utils/weight_update/` — how trained weights reach the engines:
    the session handshake, the transport protocols, and the per-backend HF weight iterators.

diff --git a/docs/user-guide/training-backend.md b/docs/user-guide/training-backend.md
@@ -554,39 +554,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)

@@ -8,10 +8,10 @@
 
 from miles.backends.fsdp_utils import checkpoint
 from miles.backends.fsdp_utils.adaptations import routing_replay
+from miles.backends.training_utils.metrics import train_dump
 from miles.backends.training_utils.parallel import get_parallel_state, set_parallel_state
 from miles.backends.training_utils.torch_native.actor import TorchNativeTrainRayActor
 from miles.backends.training_utils.torch_native.step_runner import LinearStepRunner, StepMetrics
-from miles.utils import train_dump_utils
 from miles.utils.context_utils import with_defer
 from miles.utils.distributed_utils import get_gloo_group
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
@@ -288,12 +288,7 @@ def _fsdp2_load_full_state_dict(self, model, full_state, device_mesh, cpu_offloa
 
         return model
 
-    def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
-        """Delegate checkpoint saving to the shared checkpoint utilities."""
-        if self.args.debug_rollout_only or self.args.save is None:
-            return
-
-        assert not self.args.async_save, "FSDPTrainRayActor does not support async_save yet."
+    def _save_checkpoint(self, rollout_id: int) -> None:
         checkpoint.save(self, rollout_id)
 
     @contextmanager
@@ -328,7 +323,7 @@ def _step_runner(self) -> LinearStepRunner:
 
     def _after_rollout(self, rollout_id: int, rollout_data) -> None:
         if self.args.save_debug_train_data is not None:
-            train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
+            train_dump.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
 
         if (
             self.args.ref_update_interval is not 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,

@@ -40,8 +40,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
@@ -132,6 +135,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
 
 

@@ -12,6 +12,8 @@
 from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
 from torch.distributed.checkpoint.stateful import Stateful
 
+from miles.backends.training_utils.checkpoint.tracker import CHECKPOINT_TRACKER_FILENAME
+
 logger = logging.getLogger(__name__)
 
 
@@ -92,7 +94,7 @@ def load(actor: Any) -> dict[str, Any] | None:
 
     target_step = getattr(actor.args, "ckpt_step", None)
     if target_step is None:
-        tracker_file = root_path / "latest_checkpointed_iteration.txt"
+        tracker_file = root_path / CHECKPOINT_TRACKER_FILENAME
         if not tracker_file.exists():
             logger.info(f"[FSDP] No tracker file at {tracker_file}; skipping load.")
             return None
@@ -243,7 +245,7 @@ def save(actor: Any, iteration: int) -> None:
         }
         _write_checkpoint_metadata(checkpoint_dir / "meta.json", metadata)
 
-        tracker_file = base_dir / "latest_checkpointed_iteration.txt"
+        tracker_file = base_dir / CHECKPOINT_TRACKER_FILENAME
         tracker_file.write_text(str(step_id))
         logger.info(f"[FSDP] Saved checkpoint to {checkpoint_dir}")
 

@@ -9,19 +9,35 @@
 from megatron.training.async_utils import maybe_finalize_async_save
 from torch_memory_saver import torch_memory_saver
 
-from miles.backends.megatron_utils.ft.types import TrainStepOutput
 from miles.backends.megatron_utils.hf_export import save_hf_model
 from miles.backends.megatron_utils.lora.utils import is_lora_enabled, lora_rollout_enabled
 from miles.backends.megatron_utils.rematerialize_utils import build_main_cast_context
 from miles.backends.megatron_utils.update_weight.hf_weight_iterator import get_hf_weight_iterator
+from miles.backends.training_utils.checkpoint.tracker import read_checkpoint_tracker_iteration
+from miles.backends.training_utils.data.rollout import (
+    DataIterator,
+    get_data_iterator,
+    get_num_rollouts,
+    get_rollout_data,
+)
+from miles.backends.training_utils.loss.objective import (
+    compute_advantages_and_returns,
+    get_log_probs_and_entropy,
+    get_values,
+    log_train_advantage_computation_event,
+)
+from miles.backends.training_utils.metrics import train_dump
+from miles.backends.training_utils.metrics.log_utils import log_cpu_memory, log_perf_data, log_rollout_data
+from miles.backends.training_utils.replay.data import fill_replay_data, register_replay_list_sequential
+from miles.backends.training_utils.types import TrainStepOutput
 from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement
 from miles.backends.training_utils.weight_update.snapshot_publisher import SnapshotPublisher
 from miles.backends.training_utils.weight_update.updater import WeightUpdater
 from miles.dashboard import hooks as dashboard_hooks
 from miles.ray.rollout.inference_controller import UpdatableEngines
 from miles.ray.specs.train import compute_trainer_pool_id
 from miles.ray.train_actor import TrainRayActor
-from miles.utils import object_store, train_dump_utils
+from miles.utils import object_store
 from miles.utils.argparse_utils import inplace_modify_args
 from miles.utils.audit_utils.event_logger.logger import event_logger_context
 from miles.utils.audit_utils.witness.allocator import WitnessInfo
@@ -43,18 +59,8 @@
 
 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 (
-    compute_advantages_and_returns,
-    get_log_probs_and_entropy,
-    get_values,
-    log_train_advantage_computation_event,
-)
 from ..training_utils.parallel import get_parallel_state
-from ..training_utils.replay_data import fill_replay_data, register_replay_list_sequential
 from .checkpoint import load_checkpoint
-from .checkpoint_tracker import read_checkpoint_tracker_iteration
 from .ft.checkpoint_transfer import recv_ckpt
 from .ft.checkpoint_transfer import send_ckpt as _send_ckpt
 from .ft.in_memory_checkpoint import InMemoryCheckpointManager
@@ -771,7 +777,7 @@ def _train_actor(
 
             self.prof.step(rollout_id=rollout_id)
 
-        train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
+        train_dump.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
 
         for m in all_replay_managers:
             if m.enabled:

@@ -6,12 +6,12 @@
 import torch.distributed as dist
 from megatron.core import mpu
 
+from miles.backends.training_utils.metrics.log_utils import aggregate_train_losses
 from miles.utils.distributed_utils import get_gloo_group
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
 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.parallel import ParallelState
 
 if TYPE_CHECKING:

@@ -10,7 +10,7 @@
 
 from miles.backends.megatron_utils.lora.utils import is_lora_model
 from miles.backends.megatron_utils.named_weights import named_params_and_buffers
-from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+from miles.backends.training_utils.checkpoint.io import write_checkpoint_dir
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.backends.training_utils.weight_update.snapshot_publisher import SnapshotPublisher
 from miles.utils.distributed_utils import get_gloo_group

@@ -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.object_store import StoreObjectRef
 from miles.utils.tracking_utils.structured_log import with_logs

@@ -13,7 +13,7 @@
 from megatron.core.utils import unwrap_model
 
 from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer
-from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+from miles.backends.training_utils.checkpoint.io import write_checkpoint_dir
 
 _WEIGHTS_KEY = "adapter_weights"
 _OPTIM_KEY = "adapter_optimizer"

@@ -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

@@ -10,7 +10,7 @@
 import torch
 import torch.distributed as dist
 
-from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+from miles.backends.training_utils.checkpoint.io import write_checkpoint_dir
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.backends.training_utils.weight_update.snapshot_publisher import SnapshotPublisher
 from miles.utils.lora.utils import (  # noqa: F401  (re-exported)

@@ -26,10 +26,18 @@
 from megatron.training.training import get_model
 
 from miles.backends.megatron_utils.ft.indep_dp import allreduce_grads_and_losses_across_replicas
-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.rollout import DataIterator, get_batch
+from miles.backends.training_utils.data.sampling_mask import get_rollout_sampling_masks
+from miles.backends.training_utils.loss.objective import loss_function
+from miles.backends.training_utils.metrics.checks import check_grad_norm, check_kl
+from miles.backends.training_utils.metrics.log_utils import (
+    aggregate_forward_results,
+    aggregate_train_losses,
+    log_train_step,
+)
+from miles.backends.training_utils.types import TrainStepOutcome
 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 +48,6 @@
 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.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
@@ -127,8 +127,5 @@ def _ref_context(self):
                 part.cpu()
             torch.cuda.empty_cache()
 
-    def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
-        if self.args.debug_rollout_only or self.args.save is None:
-            return
-        assert not self.args.async_save, "TorchtitanTrainRayActor does not support async_save yet."
+    def _save_checkpoint(self, rollout_id: int) -> None:
         self.trainer.checkpointer.save(self.trainer.step, last_step=True)