Skip to content

Commit 8459e50

Browse files
dyurk-lilaclaude
dyurk-lila
and
claude
committed
feat(profiler): drive torch.profiler in the full-context trainer
Bracket FullCtxTrainer.train() with _profiler_start / _profiler_step / _profiler_stop, matching the wiring in RayPPOTrainer.train(). One profile_step per dummy global step; stop runs in a finally so the open kineto trace window isn't leaked if a step raises. No-op unless torch_profiler_config.enable is set, so non-profiling runs pay nothing. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 277c67e commit 8459e50

1 file changed

Lines changed: 79 additions & 62 deletions

File tree

‎examples/train_scripts/full_context/trainer_full_ctx.py‎

Lines changed: 79 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -26,68 +26,85 @@ async def train(self):
2626

2727
# Run a few training steps
2828
self.global_step += 1 # start from 1
29-
for step in range(self.cfg.trainer.num_dummy_steps):
30-
logger.info(f"Running dummy training step {step + 1}/{self.cfg.trainer.num_dummy_steps}")
31-
32-
# Run a single training step
33-
with Timer("step", self.all_timings):
34-
# Create training input directly with max length sequences
35-
num_samples = self.cfg.trainer.train_batch_size * self.cfg.generator.n_samples_per_prompt
36-
uids = [str(i) for i in range(self.cfg.trainer.train_batch_size)]
37-
prompt_token_ids = [
38-
[random.randint(0, self.tokenizer.vocab_size - 1)] * self.cfg.generator.max_input_length
39-
] * self.cfg.trainer.train_batch_size
40-
prompt_token_ids = sum(
41-
[
42-
[prompt_token_id] * self.cfg.generator.n_samples_per_prompt
43-
for prompt_token_id in prompt_token_ids
44-
],
45-
[],
46-
)
47-
response_ids = [
48-
[random.randint(0, self.tokenizer.vocab_size - 1)]
49-
* self.cfg.generator.sampling_params.max_generate_length
50-
] * num_samples
51-
uids = sum([[uid] * self.cfg.generator.n_samples_per_prompt for uid in uids], [])
52-
53-
dummy_generator_output = {
54-
"prompt_token_ids": prompt_token_ids,
55-
"response_ids": response_ids,
56-
"rewards": [
57-
[0] * (self.cfg.generator.sampling_params.max_generate_length - 1) + [random.randint(0, 1)]
58-
]
59-
* num_samples,
60-
"loss_masks": [[1] * self.cfg.generator.sampling_params.max_generate_length] * num_samples,
61-
}
62-
training_input = self.convert_to_training_input(dummy_generator_output, uids)
63-
64-
with Timer("fwd_logprobs_values_reward", self.all_timings):
65-
training_input = self.fwd_logprobs_values_reward(training_input)
66-
67-
# 1.5 apply kl divergence penalty to rewards
68-
if self.cfg.trainer.algorithm.use_kl_in_reward:
69-
with Timer("apply_reward_kl_penalty", self.all_timings):
70-
training_input = self.apply_reward_kl_penalty(training_input)
71-
72-
# 3. calculate advantages and returns
73-
with Timer("compute_advantages_and_returns", self.all_timings):
74-
training_input = self.compute_advantages_and_returns(training_input)
75-
# remove some unwanted keys
76-
for key in ["rewards"]:
77-
training_input.pop(key)
78-
training_input.metadata.pop("uids")
79-
80-
# 4. train policy/critic model
81-
with Timer("train_critic_and_policy", self.all_timings):
82-
status = self.train_critic_and_policy(training_input)
83-
84-
self.tracker.log(self.all_metrics, step=self.global_step)
85-
self.all_metrics = {}
86-
self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step)
87-
self.all_timings = {}
88-
self.global_step += 1
89-
90-
logger.info(f"Step {step + 1} completed. Status: {status}")
29+
self._profiler_start()
30+
try:
31+
for step in range(self.cfg.trainer.num_dummy_steps):
32+
logger.info(f"Running dummy training step {step + 1}/{self.cfg.trainer.num_dummy_steps}")
33+
34+
# Run a single training step
35+
with Timer("step", self.all_timings):
36+
# Create training input directly with max length sequences
37+
num_samples = self.cfg.trainer.train_batch_size * self.cfg.generator.n_samples_per_prompt
38+
uids = [str(i) for i in range(self.cfg.trainer.train_batch_size)]
39+
prompt_token_ids = [
40+
[random.randint(0, self.tokenizer.vocab_size - 1)] * self.cfg.generator.max_input_length
41+
] * self.cfg.trainer.train_batch_size
42+
prompt_token_ids = sum(
43+
[
44+
[prompt_token_id] * self.cfg.generator.n_samples_per_prompt
45+
for prompt_token_id in prompt_token_ids
46+
],
47+
[],
48+
)
49+
response_ids = [
50+
[random.randint(0, self.tokenizer.vocab_size - 1)]
51+
* self.cfg.generator.sampling_params.max_generate_length
52+
] * num_samples
53+
uids = sum(
54+
[[uid] * self.cfg.generator.n_samples_per_prompt for uid in uids],
55+
[],
56+
)
57+
58+
dummy_generator_output = {
59+
"prompt_token_ids": prompt_token_ids,
60+
"response_ids": response_ids,
61+
"rewards": [
62+
[0] * (self.cfg.generator.sampling_params.max_generate_length - 1) + [random.randint(0, 1)]
63+
]
64+
* num_samples,
65+
"loss_masks": [[1] * self.cfg.generator.sampling_params.max_generate_length] * num_samples,
66+
}
67+
training_input = self.convert_to_training_input(dummy_generator_output, uids)
68+
69+
with Timer("fwd_logprobs_values_reward", self.all_timings):
70+
training_input = self.fwd_logprobs_values_reward(training_input)
71+
72+
# 1.5 apply kl divergence penalty to rewards
73+
if self.cfg.trainer.algorithm.use_kl_in_reward:
74+
with Timer("apply_reward_kl_penalty", self.all_timings):
75+
training_input = self.apply_reward_kl_penalty(training_input)
76+
77+
# 3. calculate advantages and returns
78+
with Timer("compute_advantages_and_returns", self.all_timings):
79+
training_input = self.compute_advantages_and_returns(training_input)
80+
# remove some unwanted keys
81+
for key in ["rewards"]:
82+
training_input.pop(key)
83+
training_input.metadata.pop("uids")
84+
85+
# 4. train policy/critic model
86+
with Timer("train_critic_and_policy", self.all_timings):
87+
status = self.train_critic_and_policy(training_input)
88+
89+
# Advance the torch profiler schedule once per global step
90+
# (no-op unless profiling is enabled).
91+
self._profiler_step()
92+
93+
self.tracker.log(self.all_metrics, step=self.global_step)
94+
self.all_metrics = {}
95+
self.tracker.log(
96+
{"timing/" + k: v for k, v in self.all_timings.items()},
97+
step=self.global_step,
98+
)
99+
self.all_timings = {}
100+
self.global_step += 1
101+
102+
logger.info(f"Step {step + 1} completed. Status: {status}")
103+
finally:
104+
# Always stop/flush the profiler when the loop exits -- including via
105+
# an exception -- so the open kineto trace window isn't leaked. No-op
106+
# when profiling is disabled.
107+
self._profiler_stop()
91108

92109
self.tracker.finish()
93110
logger.info("Dummy training completed successfully!")

0 commit comments

Comments
 (0)