feat(profiler): drive torch.profiler around the training loop

docs/api-pages.yaml

@@ -147,7 +147,7 @@
             - skyrl.train.config.config.MegatronConfig
             - skyrl.train.config.config.MegatronDDPConfig
             - skyrl.train.config.config.MegatronLoraConfig
-            - skyrl.train.config.config.MegatronTorchProfilerConfig
+            - skyrl.train.config.config.TorchProfilerConfig
         - heading: Placement
           description: ""
           objects:

docs/content/docs/configuration/config.mdx

@@ -273,6 +273,23 @@ policy:
   use_torch_compile: false  # Enable torch compile for the entropy calculation
   record_memory: false  # Dump memory snapshot for debugging
 
+  torch_profiler_config:    # torch.profiler for policy training
+    enable: false
+    ranks: [0]
+    save_path: null         # required when enable=true; absolute local path
+    skip_first: 10
+    wait: 0
+    warmup: 1
+    active: 1
+    repeat: 1                # 0 = repeat forever
+    activities: ["cpu", "cuda"]
+    record_shapes: true
+    profile_memory: false
+    with_stack: true
+    with_flops: false
+    with_modules: false
+    export_type: "chrome_trace"  # chrome_trace | stacks
+
   model_config_kwargs: {}     # pass through kwargs to the HuggingFace model config for FSDP training backends (i.e. for overriding vocab size, etc) - for megatron, use policy.megatron_config.transformer_config_kwargs instead
 
 ```
@@ -283,6 +300,22 @@ policy:
 - `policy.use_torch_compile`: Whether to enable torch compile for entropy calculation
 - `policy.record_memory`: Whether to record memory usage. If `True`, this will use PyTorch's [memory snapshotting utility](https://docs.pytorch.org/docs/stable/torch_cuda_memory.html) to record memory usage and dump memory snapshots after each policy model training step.
 
+### Torch Profiler Configuration
+
+`policy.torch_profiler_config` profiles policy training steps for FSDP and Megatron in RL and SFT. It writes one Kineto/[HTA](https://github.com/facebookresearch/HolisticTraceAnalysis)-friendly `*.pt.trace.json` per active window and profiled rank.
+
+Scope: policy forward/backward/optimizer only. RL critic/ref models and generation are not profiled.
+
+FSDP restriction: profiling is rejected when manual CPU offload and colocation are both active (`policy.fsdp_config.cpu_offload: false` with `placement.colocate_all` or `placement.colocate_policy_ref`). Use FSDP native CPU offload, disable colocation, or use Megatron.
+
+- `policy.torch_profiler_config.enable`: Enables profiling.
+- `policy.torch_profiler_config.ranks`: Global ranks to profile, for example `[0]`.
+- `policy.torch_profiler_config.save_path`: Required absolute local trace directory. Relative paths land under Ray's `/tmp/ray/.../working_dir_files`; cloud URIs are rejected.
+- `policy.torch_profiler_config.{skip_first,wait,warmup,active,repeat}`: Passed to [`torch.profiler.schedule`](https://docs.pytorch.org/docs/stable/profiler.html#torch.profiler.schedule).
+- `policy.torch_profiler_config.activities`: Subset of `["cpu", "cuda"]`.
+- `policy.torch_profiler_config.{record_shapes,profile_memory,with_stack,with_flops,with_modules}`: Passed to `torch.profiler.profile`.
+- `policy.torch_profiler_config.export_type`: `chrome_trace` for `*.pt.trace.json`; `stacks` for self-CUDA-time stacks (`with_stack: true` required).
+
 ### LoRA Configuration
 
 LoRA (Low-Rank Adaptation) enables parameter-efficient fine-tuning by training only a small number of additional low-rank matrices instead of the full model weights:

examples/train/async/async_trainer.py

@@ -51,63 +51,73 @@ async def train(self):
         start_epoch = self.global_step // len(self.train_dataloader)
         # Start from step 1
         self.global_step += 1
-        for epoch in range(start_epoch, self.cfg.trainer.epochs):
-            # while this is just off by one, you can image a more general queue based approach
-            # where the generation buffer holds a list of objects that the trainer can read from
-            # bit by bit.
-            generation_buffer = asyncio.Queue(maxsize=1)
-            self.sync_finished = asyncio.Event()
-            self.generation_ack = asyncio.Event()
-
-            # start generator task
-            generator_task = asyncio.create_task(self._run_generate_loop(generation_buffer))
-
-            for idx in range(len(self.train_dataloader)):
-                with Timer("step", self.all_timings):
-                    status = await self._run_training(generation_buffer)
-
-                    # request the generation loop that we should sync sometime soon.
-                    if idx != len(self.train_dataloader) - 1:
-                        await self.generation_ack.wait()
-
-                    # sync weights
-                    async with Timer("sync_weights", self.all_timings):
-                        await self.dispatch.save_weights_for_sampler()
-
-                    self.sync_finished.set()
-                    self.generation_ack.clear()
-
-                # 5. set logs
-                logger.info(status)
-                # log epoch info
-                self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step})
-                self.tracker.log(self.all_metrics, step=self.global_step)
-                self.all_metrics = {}
-                pbar.update(1)
-
-                if self.cfg.trainer.eval_interval > 0 and (
-                    self.global_step % self.cfg.trainer.eval_interval == 0
-                    or self.global_step == self.total_training_steps
-                ):
-                    with Timer("eval", self.all_timings):
-                        eval_metrics = await self.eval()
-                        self.all_metrics.update(eval_metrics)
-                if self.cfg.trainer.ckpt_interval > 0 and self.global_step % self.cfg.trainer.ckpt_interval == 0:
-                    with Timer("save_checkpoints", self.all_timings):
-                        self.save_checkpoints()
-                if self.cfg.trainer.hf_save_interval > 0 and self.global_step % self.cfg.trainer.hf_save_interval == 0:
-                    with Timer("save_hf_model", self.all_timings):
-                        self.save_models()
-                self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step)
-                self.all_timings = {}
-                self.global_step += 1
-
-            if self.cfg.trainer.update_ref_every_epoch and self.ref_model is not None:
-                with Timer("update_ref_with_policy", self.all_timings):
-                    await asyncio.to_thread(self.update_ref_with_policy)
-
-            # cancel generation task for this epoch
-            generator_task.cancel()
+        self._profiler_start()
+        try:
+            for epoch in range(start_epoch, self.cfg.trainer.epochs):
+                # while this is just off by one, you can image a more general queue based approach
+                # where the generation buffer holds a list of objects that the trainer can read from
+                # bit by bit.
+                generation_buffer = asyncio.Queue(maxsize=1)
+                self.sync_finished = asyncio.Event()
+                self.generation_ack = asyncio.Event()
+
+                # start generator task
+                generator_task = asyncio.create_task(self._run_generate_loop(generation_buffer))
+
+                for idx in range(len(self.train_dataloader)):
+                    with Timer("step", self.all_timings):
+                        status = await self._run_training(generation_buffer)
+
+                        # request the generation loop that we should sync sometime soon.
+                        if idx != len(self.train_dataloader) - 1:
+                            await self.generation_ack.wait()
+
+                        # sync weights
+                        async with Timer("sync_weights", self.all_timings):
+                            await self.dispatch.save_weights_for_sampler()
+
+                        self.sync_finished.set()
+                        self.generation_ack.clear()
+
+                    # One profiler step per global step.
+                    self._profiler_step()
+
+                    # 5. set logs
+                    logger.info(status)
+                    # log epoch info
+                    self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step})
+                    self.tracker.log(self.all_metrics, step=self.global_step)
+                    self.all_metrics = {}
+                    pbar.update(1)
+
+                    if self.cfg.trainer.eval_interval > 0 and (
+                        self.global_step % self.cfg.trainer.eval_interval == 0
+                        or self.global_step == self.total_training_steps
+                    ):
+                        with Timer("eval", self.all_timings):
+                            eval_metrics = await self.eval()
+                            self.all_metrics.update(eval_metrics)
+                    if self.cfg.trainer.ckpt_interval > 0 and self.global_step % self.cfg.trainer.ckpt_interval == 0:
+                        with Timer("save_checkpoints", self.all_timings):
+                            self.save_checkpoints()
+                    if (
+                        self.cfg.trainer.hf_save_interval > 0
+                        and self.global_step % self.cfg.trainer.hf_save_interval == 0
+                    ):
+                        with Timer("save_hf_model", self.all_timings):
+                            self.save_models()
+                    self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step)
+                    self.all_timings = {}
+                    self.global_step += 1
+
+                if self.cfg.trainer.update_ref_every_epoch and self.ref_model is not None:
+                    with Timer("update_ref_with_policy", self.all_timings):
+                        await asyncio.to_thread(self.update_ref_with_policy)
+
+                # cancel generation task for this epoch
+                generator_task.cancel()
+        finally:
+            self._profiler_stop()
 
         pbar.close()
         if self.cfg.trainer.ckpt_interval > 0:

examples/train/megatron/run_megatron.sh

@@ -17,7 +17,7 @@ MEGATRON_TP=2
 MEGATRON_PP=2
 MEGATRON_CP=1
 
-# torch profiler config
+# torch.profiler for policy training steps.
 ENABLE_TORCH_PROFILER=false
 RANKS_TO_PROFILE="[0]"
 SAVE_PATH="$HOME/megatron_prof/tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}"
@@ -33,9 +33,9 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \
   trainer.placement.ref_num_gpus_per_node=$NUM_GPUS \
   generator.inference_engine.num_engines=$NUM_GPUS \
   generator.inference_engine.tensor_parallel_size=1 \
-  trainer.policy.megatron_config.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \
-  trainer.policy.megatron_config.torch_profiler_config.ranks=$RANKS_TO_PROFILE \
-  trainer.policy.megatron_config.torch_profiler_config.save_path=$SAVE_PATH \
+  trainer.policy.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \
+  trainer.policy.torch_profiler_config.ranks=$RANKS_TO_PROFILE \
+  trainer.policy.torch_profiler_config.save_path=$SAVE_PATH \
   trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \
   trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \
   trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \
@@ -69,4 +69,4 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \
   trainer.run_name="gsm8k_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}" \
   trainer.resume_mode=null \
   trainer.ckpt_path="$HOME/ckpts/gsm8k_megatron_ckpt" \
-  $@
+  $@

examples/train/megatron/run_megatron_nemotron_mini_4b.sh

@@ -37,7 +37,7 @@ MEGATRON_CP=1
 NUM_INFERENCE_ENGINES=8
 INFERENCE_TP=1
 
-# torch profiler config
+# torch.profiler config
 ENABLE_TORCH_PROFILER=false
 RANKS_TO_PROFILE="[0]"
 SAVE_PATH="$HOME/megatron_prof/tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}"
@@ -53,9 +53,9 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \
   trainer.placement.ref_num_gpus_per_node=$NUM_GPUS \
   generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \
   generator.inference_engine.tensor_parallel_size=$INFERENCE_TP \
-  trainer.policy.megatron_config.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \
-  trainer.policy.megatron_config.torch_profiler_config.ranks=$RANKS_TO_PROFILE \
-  trainer.policy.megatron_config.torch_profiler_config.save_path=$SAVE_PATH \
+  trainer.policy.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \
+  trainer.policy.torch_profiler_config.ranks=$RANKS_TO_PROFILE \
+  trainer.policy.torch_profiler_config.save_path=$SAVE_PATH \
   trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \
   trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \
   trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \
@@ -89,4 +89,4 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \
   trainer.run_name="gsm8k_megatron_nemotron_mini_4b_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \
   trainer.resume_mode=null \
   trainer.ckpt_path="$HOME/ckpts/gsm8k_megatron_nemotron_ckpt" \
-  $@
+  $@