-
Notifications
You must be signed in to change notification settings - Fork 429
Expand file tree
/
Copy pathfully_async_trainer.py
More file actions
1207 lines (1054 loc) · 63.5 KB
/
Copy pathfully_async_trainer.py
File metadata and controls
1207 lines (1054 loc) · 63.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""
Implementation of fully async training in SkyRL.
For details, see https://docs.skyrl.ai/docs/tutorials/fully_async.
High-level notes:
- The global_step in each training loop iteration denotes the "current step being worked on", so
`global_step - 1` is the number of steps the model has finished training.
- We do not do any cross-epoch asynchrony here, so all the control logics like
generation workers and data buffer are initialized per-epoch. The async dataloader
and staleness manager are also reset / validated at the end of each epoch.
"""
import asyncio
import inspect
import os
import sys
import time
import traceback
from dataclasses import dataclass
from typing import Any, Iterable, List, Optional, Set, Tuple
import torch
from loguru import logger
from torchdata.stateful_dataloader import StatefulDataLoader
from tqdm import tqdm
from skyrl.backends.skyrl_train.inference_servers.engine_utils import (
get_sampling_params_for_backend,
)
from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch
from skyrl.backends.skyrl_train.utils.io import io
from skyrl.backends.skyrl_train.utils.ppo_utils import (
LOSSES_WITH_OLD_LOGPROBS,
PolicyLossType,
)
from skyrl.train.generators.base import GeneratorOutput
from skyrl.train.generators.utils import (
concatenate_generator_outputs,
get_metrics_from_generator_output,
prepare_generator_input,
)
from skyrl.train.trainer import RayPPOTrainer
from skyrl.train.utils import Timer
from skyrl.train.utils.metrics import ScalarGauges, TrainingPhaseGauge
from skyrl.train.utils.trainer_utils import (
ResumeMode,
build_dataloader,
get_group_completion_metrics,
get_intra_group_completion_time_std_cv,
zero_variance_filter,
)
@dataclass
class GeneratedOutputGroup:
"""
The GeneratorOutput for a single group of rollouts, along with the metadata.
Attributes:
generator_output (GeneratorOutput): The GeneratorOutput for a single group of rollouts.
That is, the output to the same prompt, but `n_samples_per_prompt` of it.
uid (str): The uid of the group. Underlyingly, it is the index of the data in train_dataloader.dataset.
global_step_when_scheduled (int): The global step when the group was scheduled for generation,
used for validating the staleness control.
group_completion_time_s (Optional[float]): Wall-clock time (seconds) for the whole group to
finish generation, i.e. the time for the slowest trajectory in the group to complete.
None if generation timing was not captured.
prompts (Optional[List[Any]]): The generator input prompts (in OpenAI message format) for this
group, one per trajectory (parallel to ``generator_output["response_ids"]``). Retained so
the trajectory logger can render prompt + response, mirroring the synchronous trainer which
logs ``generator_input["prompts"]``. None if not captured.
"""
generator_output: GeneratorOutput
uid: str
global_step_when_scheduled: int
group_completion_time_s: Optional[float] = None
prompts: Optional[List[Any]] = None
@dataclass
class _RolloutStat:
"""
Global statistics of the trajectories used for staleness control in `_AsyncStalenessManager`.
Note that these statistics are not per-epoch, but accumulates across all epochs.
Attributes:
submitted (int): The number of trajectories submitted to all generation workers, only used
for logging purposes (e.g. compute submitted / accepted ratio to see how many
trajectories failed). This is strictly increasing.
accepted (int): The number of trajectories finished generation that will be trained on (either
already consumed by, or about to be consumed by the training worker). This is strictly
increasing in the absence of filtering; ``on_rollout_filtered`` decrements it when an
already-generated group is dropped from training (see ``sample_full_batch``).
running (int): The number of trajectories currently being generated by the generation workers.
filtered (int): The number of finished trajectories that were dropped from training (e.g.
zero-variance groups under ``sample_full_batch``). This is strictly increasing. The
invariant ``submitted == accepted + filtered + running`` always holds.
For details, see https://docs.skyrl.ai/docs/tutorials/fully_async#async-staleness-manager
"""
submitted: int = 0
accepted: int = 0
running: int = 0
filtered: int = 0
class _AsyncStalenessManager:
"""
A controller that manages the capacity of the generation workers based on staleness control.
The goal is to never submit more trajectories to the generation workers than the training worker
can consume, so that the trajectories are not too stale (relative to max_staleness_steps).
This is enforced via a capacity rule, not a hard **per-group** staleness guarantee: we bound
the **aggregate** number of groups that can be ahead of training so that, in **steady state**,
staleness remains within the configured budget of `max_staleness_steps`.
In pathological cases (e.g., very long-running trajectories), an individual group may take
more than `max_staleness_steps` of training steps of time to finish generation. For such rare
cases, we still accept the trajectory and log the staleness metrics with a warning.
The key capacity formula is implemented in `_compute_capacity_unlocked`. For details and caveats,
see https://docs.skyrl.ai/docs/tutorials/fully_async#async-staleness-manager.
Reference:
- Modeled after AReal's StalenessManager: https://github.com/inclusionAI/AReaL/blob/b755c4447c2fff97889d8828293ee85f17a806f9/areal/core/staleness_manager.py
- The idea of this controller is from section 5.1 of AReal's paper: https://arxiv.org/pdf/2505.24298v3
"""
def __init__(self, max_concurrent_generation_groups: int, mini_batch_size: int, max_staleness_steps: int):
self.max_concurrent_generation_groups = max_concurrent_generation_groups
self.mini_batch_size = mini_batch_size
self.max_staleness_steps = max_staleness_steps
# Control logics.
self._stat = _RolloutStat()
self._cond = asyncio.Condition()
# The current version that is being worked on, i.e. FullyAsyncRayPPOTrainer.global_step.
# `self._current_global_step - 1` is the number of steps the model has finished training.
self._current_global_step = 1
def load_state_from_checkpoint(self, global_step: int) -> None:
"""
Load the state from a checkpoint.
"""
self._current_global_step = global_step
# trainer has already consumed (and hence submitted) this many trajectories.
self._stat.accepted = (global_step - 1) * self.mini_batch_size
self._stat.submitted = self._stat.accepted
async def validate_state_at_epoch_end(self, global_step: int) -> None:
"""
Check that the current version and accepted rollouts are consistent with the global step.
Args:
global_step: The global step we are about to train on (after incrementing).
"""
async with self._cond:
assert self._stat.running == 0, "We expect no rollouts are running at end of an epoch."
consumed = (global_step - 1) * self.mini_batch_size
assert (
self._stat.accepted == consumed
), f"Unexpected number of accepted rollouts. Got {self._stat.accepted} != {consumed}."
assert (
self._current_global_step == global_step
), f"Unexpected current version. Got {self._current_global_step} != {global_step}."
# All submitted rollouts must have finished as either accepted (trained) or filtered (dropped).
assert self._stat.submitted == self._stat.accepted + self._stat.filtered, (
"We expect all submitted rollouts to be accepted or filtered at end of an epoch. "
f"Got {self._stat.submitted} != {self._stat.accepted} + {self._stat.filtered}."
)
def _compute_capacity_unlocked(self) -> int:
# NOTE(Charlie): do not need a self._current_global_step + 1 here unlike AReal because our
# `_current_global_step` is "the version being worked on", not already finished steps.
consumer_capacity = (self.max_staleness_steps + self._current_global_step) * self.mini_batch_size
producer_staleness_capacity = consumer_capacity - (self._stat.accepted + self._stat.running)
producer_concurrency_capacity = self.max_concurrent_generation_groups - self._stat.running
return min(producer_concurrency_capacity, producer_staleness_capacity)
async def acquire_submission_slot(self) -> None:
"""Block until there is capacity, then reserve a slot (increments submitted and running).
This method always uses the latest current version tracked internally, which is
updated by `notify_capacity_change(new_global_step)`.
"""
async with self._cond:
while self._compute_capacity_unlocked() <= 0:
await self._cond.wait()
# Reserve slot
self._stat.submitted += 1
self._stat.running += 1
async def on_rollout_accepted(self) -> None:
async with self._cond:
self._stat.accepted += 1
self._stat.running -= 1
self._cond.notify_all()
async def on_rollout_filtered(self) -> None:
"""Reclassify an already-accepted group as filtered when it is dropped from training.
Without this, dropped groups keep counting toward ``accepted`` while ``current_global_step``
only advances on trained steps, shrinking producer capacity on every drop -> deadlock.
"""
async with self._cond:
self._stat.accepted -= 1
self._stat.filtered += 1
self._cond.notify_all()
async def on_rollout_rejected(self) -> None:
"""
Called when a generation is not accepted, or generation worker runs into error while generating a trajectory.
Currently, we do not call this method but instead raise errors. We might need to use this when we want to
filter out trajectories.
"""
async with self._cond:
self._stat.running -= 1
self._cond.notify_all()
async def notify_capacity_change(self, new_global_step: int) -> None:
# Called when current_global_step changes (e.g., after a training step)
async with self._cond:
self._current_global_step = int(new_global_step)
self._cond.notify_all()
class _AsyncDataloader:
"""
A train dataloader wrapper that accommodates the need of fully async training, including:
- Thread-safe dataloader iteration with a lock, as there are multiple parallel generation workers polling data.
- Records consumed data UIDs for checkpointing to avoid training on the same data upon resuming.
- Set the effective dataloader length to be divisible by mini-batch size, since we cannot rely on `drop_last`
because the batch size is 1 in fully async training.
"""
def __init__(self, train_dataloader: StatefulDataLoader, mini_batch_size: int):
self._train_dataloader = train_dataloader
self._train_dataloader_initial_state = train_dataloader.state_dict()
self._effective_dataloader_length = len(self._train_dataloader) // mini_batch_size * mini_batch_size
self._iter = enumerate(self._train_dataloader)
self._lock: asyncio.Lock = asyncio.Lock()
# `_consumed_data_uids` = do-not-redraw set (trained ∪ filtered); `_filtered_data_uids` =
# the dropped (not trained) subset. Both persisted so dropped prompts aren't regenerated on
# resume; trained count is `len(consumed) - len(filtered)`.
self._consumed_data_uids: Set[str] = set()
self._filtered_data_uids: Set[str] = set()
self._exhausted: bool = False # currently not used.
def load_state_from_checkpoint(
self, consumed_data_uids_set: Set[str], filtered_data_uids_set: Optional[Set[str]] = None
) -> None:
"""
Load the state from a checkpoint.
"""
self._consumed_data_uids = consumed_data_uids_set
self._filtered_data_uids = filtered_data_uids_set if filtered_data_uids_set is not None else set()
# Reset in case the dataloader loaded the state from the checkpoint, which we do not want.
self._train_dataloader.load_state_dict(self._train_dataloader_initial_state)
async def reset_at_epoch_end(self) -> None:
async with self._lock:
self._train_dataloader.load_state_dict(self._train_dataloader_initial_state) # reset to initial state
self._iter = enumerate(self._train_dataloader)
self._consumed_data_uids.clear()
self._filtered_data_uids.clear()
self._exhausted = False
async def get_next_non_consumed_data(self):
"""
Return the next batch of training data.
If we loaded from a checkpoint, it will skip the already-consumed data. Returns None if the dataloader is exhausted.
"""
assert self._iter is not None and self._lock is not None, "Dataloader not initialized; call reset() first"
async with self._lock:
try:
while True:
# Keep polling until we get a non-consumed data or the dataloader is exhausted.
iter_idx, rand_prompts = next(self._iter)
if iter_idx >= self._effective_dataloader_length:
raise StopIteration
uid = rand_prompts[0]["uid"]
if uid not in self._consumed_data_uids:
return rand_prompts
except StopIteration:
self._exhausted = True
return None
async def mark_consumed_uids(self, uids: Iterable[str]) -> None:
"""Mark UIDs as trained on (and hence consumed)."""
assert self._lock is not None, "Dataloader not initialized; call reset() first"
async with self._lock:
for uid in uids:
assert uid not in self._consumed_data_uids, "Duplicate UID found in mini-batch"
self._consumed_data_uids.add(uid)
async def mark_filtered_uids(self, uids: Iterable[str]) -> None:
"""Mark UIDs as dropped (not trained on): added to the do-not-redraw set and recorded as
filtered so they are skipped on resume and excluded from the trained-step count."""
assert self._lock is not None, "Dataloader not initialized; call reset() first"
async with self._lock:
for uid in uids:
assert uid not in self._consumed_data_uids, "Duplicate UID found in mini-batch"
self._consumed_data_uids.add(uid)
self._filtered_data_uids.add(uid)
def get_consumed_uids_list(self) -> List[str]:
return list(self._consumed_data_uids)
def get_filtered_uids_list(self) -> List[str]:
return list(self._filtered_data_uids)
def num_trained(self) -> int:
"""Number of UIDs consumed by training (excludes filtered/dropped UIDs)."""
return len(self._consumed_data_uids) - len(self._filtered_data_uids)
class FullyAsyncRayPPOTrainer(RayPPOTrainer):
def __init__(self, *args, **kwargs):
# Extract cfg before base init so we can initialize async-specific knobs used by our overrides.
cfg = kwargs.get("cfg", args[0] if len(args) > 0 else None)
assert cfg is not None, "cfg must be provided to FullyAsyncRayPPOTrainer"
# Initialize async-specific knobs
self.num_parallel_generation_workers = cfg.trainer.fully_async.num_parallel_generation_workers
self.mini_batch_size = cfg.trainer.policy_mini_batch_size
self.max_staleness_steps = cfg.trainer.fully_async.max_staleness_steps
self.sample_full_batch = cfg.trainer.fully_async.sample_full_batch
self._phase_gauge = TrainingPhaseGauge()
self._gen_buffer_maxsize = self.mini_batch_size * (self.max_staleness_steps + 1)
self._loop_gauges = ScalarGauges()
self._loop_gauges.set(
"skyrl_mini_batch_size", self.mini_batch_size, "Generation groups consumed per training step."
)
self._loop_gauges.set(
"skyrl_gen_buffer_maxsize", self._gen_buffer_maxsize, "Staleness-bounded generation-buffer capacity."
)
assert (
# otherwise wasted throughput
self.mini_batch_size <= self.num_parallel_generation_workers
and
# otherwise would never use all workers due to capacity constraint
self.num_parallel_generation_workers <= self.mini_batch_size * (self.max_staleness_steps + 1)
), (
"Invalid num_parallel_generation_workers, the following must hold: "
"mini_batch_size <= num_parallel_generation_workers <= mini_batch_size * (max_staleness_steps + 1). Got: "
f"{self.mini_batch_size=}, {self.num_parallel_generation_workers=}, {self.max_staleness_steps=}"
)
# Initialize base trainer
super().__init__(*args, **kwargs)
# Callbacks aren't wired into FullyAsyncRayPPOTrainer.train() yet — fail
# fast
if self._callback_handler.callbacks:
raise NotImplementedError("Callbacks are not yet supported by FullyAsyncRayPPOTrainer. ")
# Some async-specific validations
assert (
self.cfg.trainer.fully_async.enabled
), "trainer.fully_async.enabled must be True when using the fully async trainer."
assert (
self.cfg.trainer.train_batch_size == self.cfg.trainer.policy_mini_batch_size
), "train_batch_size must equal policy_mini_batch_size for fully async training"
assert (
self.cfg.trainer.algorithm.dynamic_sampling.type is None
), "dynamic sampling is not supported for fully async training yet."
if self.sample_full_batch:
assert self.cfg.trainer.algorithm.zero_variance_filter, (
"trainer.fully_async.sample_full_batch=True requires trainer.algorithm.zero_variance_filter=True "
"(it is the async-native equivalent of dynamic_sampling='filter')."
)
assert (
not self.cfg.generator.batched
), "batched is not supported for fully async training since a batched generate() call does not support pause/continue."
# TODO(Charlie): we can support it, just multi-turn partial rollout but synchronous.
assert not self.colocate_all, "colocate_all is not supported for async training yet."
# Ensure we're using a policy loss that doesn't depend on pi_old ~= pi_rollout
loss_type = self.cfg.trainer.algorithm.policy_loss_type
if loss_type == PolicyLossType.CISPO and self.cfg.trainer.algorithm.cispo.cispo_anchor == "rollout":
optimizes_against_rollout = True
else:
optimizes_against_rollout = loss_type not in LOSSES_WITH_OLD_LOGPROBS
assert optimizes_against_rollout, (
f"Found trainer.algorithm.policy_loss_type={loss_type} in "
f"{sorted([loss.value for loss in LOSSES_WITH_OLD_LOGPROBS])}. Fully async training should use "
"rollout logprobs (i.e. rollout_is, dppo, or cispo with cispo.cispo_anchor='rollout') instead "
"of recomputing logprobs, since stale policies are not kept for logprob computation."
)
# TODO(Charlie): need to assert we are doing TIS and returning logprobs
# Async-specific states
self.async_train_dataloader = _AsyncDataloader(self.train_dataloader, self.mini_batch_size)
self._staleness_manager = _AsyncStalenessManager(
max_concurrent_generation_groups=self.num_parallel_generation_workers,
mini_batch_size=self.mini_batch_size,
max_staleness_steps=self.max_staleness_steps,
)
def add_callback(self, callback):
raise NotImplementedError("Callbacks are not yet supported by FullyAsyncRayPPOTrainer. ")
def _build_train_dataloader_and_compute_training_steps(self):
"""
Overrides to build dataloader for fully async training. See `_AsyncDataloader` for more details.
"""
self.train_dataloader = build_dataloader(self.cfg, self.train_dataset, is_train=True, is_fully_async=True)
self.num_steps_per_epoch = len(self.train_dataloader) // self.mini_batch_size
self.total_training_steps = self.num_steps_per_epoch * self.cfg.trainer.epochs
if self.cfg.trainer.max_training_steps is not None:
self.total_training_steps = min(self.total_training_steps, self.cfg.trainer.max_training_steps)
logger.info(f"Length of train_dataloader: {len(self.train_dataloader)}")
logger.info(f"Number of steps per epoch: {self.num_steps_per_epoch}")
logger.info(f"Total training steps: {self.total_training_steps}")
async def train(self):
"""
Main fully async training loop for PPO
"""
self.global_step = 0
self.epoch = 0
resumed_start_epoch = None
# Load checkpoint state if resumption is enabled. Also load the data UIDs that are already trained on.
if self.resume_mode != ResumeMode.NONE:
with Timer("load_checkpoints"):
(
self.global_step,
_,
loaded_consumed_data_uids_set,
loaded_filtered_data_uids_set,
loaded_epoch,
) = self.load_checkpoints()
logger.info(f"Resumed training from global_step {self.global_step}")
if self.global_step > 0:
# Set async dataloader manager and staleness manager to the loaded state.
self.async_train_dataloader.load_state_from_checkpoint(
loaded_consumed_data_uids_set, loaded_filtered_data_uids_set
)
self._staleness_manager.load_state_from_checkpoint(
self.global_step + 1
) # +1 due to we haven't incremented yet
# Only trained UIDs map to completed steps (filtered UIDs are extra consumption),
# so validate the trained count is a whole number of steps.
num_trained_loaded = len(loaded_consumed_data_uids_set) - len(loaded_filtered_data_uids_set)
assert num_trained_loaded % self.mini_batch_size == 0, (
"Loaded trained (consumed minus filtered) data UIDs must be a multiple of "
f"mini_batch_size={self.mini_batch_size}. Got: {num_trained_loaded}"
)
# Use the persisted epoch; fall back to deriving it for pre-sample_full_batch
# checkpoints (where global_step stays aligned to epoch boundaries).
resumed_start_epoch = (
loaded_epoch if loaded_epoch is not None else self.global_step // self.num_steps_per_epoch
)
# Initialize weight sync state
with Timer("init_weight_sync_state"):
self.init_weight_sync_state()
# sync weights to inference engines
with Timer("sync_weights_to_inference_engines"):
await self.dispatch.save_weights_for_sampler()
# Per-step GPU utilization to the tracker. The base loop starts, flushes, and stops the
# monitor itself. The async loop overrides train() and must wire it here.
if self._ray_gpu_monitor is not None:
self._ray_gpu_monitor.start()
# Eval before training
if self.cfg.trainer.eval_interval > 0 and self.cfg.trainer.eval_before_train:
with self._phase_gauge.timed_phase("eval", self.all_timings):
eval_metrics = await self.eval()
self.tracker.log(eval_metrics, step=self.global_step, commit=True)
# main training loop
pbar = tqdm(total=self.total_training_steps, initial=self.global_step, desc="Training Step Progress")
start_epoch = resumed_start_epoch if resumed_start_epoch is not None else 0
self.global_step += 1 # start training at global_step 1
stop_training = False
self._profiler_start()
try:
for epoch in range(start_epoch, self.cfg.trainer.epochs):
self.epoch = epoch
# 0. Per-epoch prologue. Note that we do not do any cross-epoch asynchrony here.
# Buffer of completed generation, size bounded by capacity - consumed = B * (max_staleness_steps + 1)
generation_output_group_buffer = asyncio.Queue[GeneratedOutputGroup](maxsize=self._gen_buffer_maxsize)
# Maintain self.num_parallel_generation_workers concurrent group-generation workers
generator_tasks = [
asyncio.create_task(self._run_generate_for_a_group_loop(generation_output_group_buffer))
for _ in range(self.num_parallel_generation_workers)
]
# Lets the consumer detect epoch exhaustion (all generators done + buffer empty) instead of
# blocking forever on buffer.get() -- under sample_full_batch, drops can exhaust an epoch
# before num_steps_per_epoch steps complete.
all_generators_done = asyncio.Event()
generators_done_watcher = None
if self.sample_full_batch:
async def _watch_generators_done(tasks=generator_tasks, event=all_generators_done):
await asyncio.gather(*tasks, return_exceptions=True)
event.set()
generators_done_watcher = asyncio.create_task(_watch_generators_done())
# Steps trained in THIS epoch (not global_step % num_steps_per_epoch: sample_full_batch can
# end an epoch early, drifting global_step out of epoch alignment). On resume the dataloader
# already reflects this epoch's trained steps. The range below is just an upper bound.
trained_steps_this_epoch = self.async_train_dataloader.num_trained() // self.mini_batch_size
for _step_idx in range(self.global_step, (1 + epoch) * self.num_steps_per_epoch + 1):
with Timer("step", self.all_timings):
self._loop_gauges.set(
"skyrl_gen_buffer_qsize",
generation_output_group_buffer.qsize(),
"Completed generation groups buffered at step start.",
)
# 1. Wait until we have a full mini-batch buffered (dropping zero-variance groups if
# sample_full_batch).
(
cur_generation_group_mini_batch,
cur_dropped_groups,
epoch_exhausted,
) = await self._collect_generation_mini_batch(
generation_output_group_buffer, all_generators_done
)
if epoch_exhausted:
# Exhausted mid mini-batch: discard the partial batch (marked consumed so it
# isn't regenerated on resume) and end the epoch early.
if cur_generation_group_mini_batch:
for _ in cur_generation_group_mini_batch:
await self._staleness_manager.on_rollout_filtered()
await self.async_train_dataloader.mark_filtered_uids(
[g.uid for g in cur_generation_group_mini_batch]
)
logger.warning(
f"sample_full_batch: epoch {epoch} exhausted with a partial mini-batch of "
f"{len(cur_generation_group_mini_batch)} group(s); discarding and ending the epoch."
)
# Save the end-of-epoch checkpoint the normal is_epoch_end path would have, since
# we break before reaching it.
if self.cfg.trainer.ckpt_interval > 0:
with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
await asyncio.to_thread(self.save_checkpoints)
if self.cfg.trainer.hf_save_interval > 0:
with self._phase_gauge.timed_phase("save_hf_model", self.all_timings):
await asyncio.to_thread(self.save_models)
break
if self.sample_full_batch:
# The collect loop returns exactly mini_batch_size kept groups here.
num_dropped = len(cur_dropped_groups)
keep_rate = self.mini_batch_size / (self.mini_batch_size + num_dropped)
self.all_metrics["async/num_groups_dropped"] = num_dropped
self.all_metrics["async/keep_rate"] = keep_rate
self._loop_gauges.set(
"skyrl_gen_group_keep_rate",
keep_rate,
"Fraction of drained groups kept while collecting the last mini-batch.",
)
# 2. Post-process the generated groups, aggregating to a single GeneratorOutput, and convert to training format.
with self._phase_gauge.timed_phase("convert_to_training_input", self.all_timings):
training_input = await asyncio.to_thread(
self.convert_generation_group_mini_batch_to_training_input,
cur_generation_group_mini_batch,
cur_dropped_groups,
)
# 3. Run training and update consumed UIDs.
with self._phase_gauge.timed_phase("run_training", self.all_timings):
status = await self._run_training(training_input)
await self.async_train_dataloader.mark_consumed_uids(
[g.uid for g in cur_generation_group_mini_batch]
)
# 4. After training: pause generation, sync weights, resume.
with self._phase_gauge.timed_phase("sync_weights", self.all_timings):
await self.dispatch.save_weights_for_sampler()
# `sync_weights` above is the full bracket: it also pauses and
# resumes generation, which under vLLM DP costs seconds of
# coordinator quiesce that is not weight-sync work. The
# dispatch reports the transfer on its own alongside it.
self.all_timings.update(self.dispatch.get_timing_metrics())
# A training step completed: count it for this epoch's bookkeeping.
trained_steps_this_epoch += 1
# One profiler step per async global step.
self._profiler_step()
# 5. Set logs for this training step.
logger.info(status)
self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step})
pbar.update(1)
# 6. Eval. At interval and at the last step.
# NOTE(Charlie): eval does not overlap with training, but overlaps with generation.
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 self._phase_gauge.timed_phase("eval", self.all_timings):
eval_metrics = await self.eval()
self.all_metrics.update(eval_metrics)
# Log metrics for this step after evaluation
self.tracker.log(self.all_metrics, step=self.global_step, commit=False)
self.all_metrics = {}
# 7. Checkpointing. At interval and at the last step of each epoch.
is_epoch_end = trained_steps_this_epoch == self.num_steps_per_epoch
if self.cfg.trainer.ckpt_interval > 0:
if is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0:
with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
await asyncio.to_thread(self.save_checkpoints)
if self.cfg.trainer.hf_save_interval > 0:
if is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0:
with self._phase_gauge.timed_phase("save_hf_model", self.all_timings):
await asyncio.to_thread(self.save_models)
timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()}
if self._ray_gpu_monitor is not None:
timing_payload.update(self._ray_gpu_monitor.flush())
if self._vllm_metrics_scraper is not None:
timing_payload.update(await self._vllm_metrics_scraper.sample())
self.tracker.log(timing_payload, step=self.global_step, commit=True)
self.all_timings = {}
self.global_step += 1
if (
self.cfg.trainer.max_training_steps is not None
and self.global_step > self.cfg.trainer.max_training_steps
):
logger.info(
f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early."
)
for t in generator_tasks:
t.cancel()
await asyncio.gather(*generator_tasks, return_exceptions=True)
if generators_done_watcher is not None:
generators_done_watcher.cancel()
await asyncio.gather(generators_done_watcher, return_exceptions=True)
stop_training = True
break
# 8. Notify generation workers that the capacity has increased, unblocking them.
await self._staleness_manager.notify_capacity_change(self.global_step)
# Only trained UIDs map to completed steps; filtered/dropped UIDs are extra consumption.
expected_trained_in_epoch = self.mini_batch_size * trained_steps_this_epoch
actual_trained_in_epoch = self.async_train_dataloader.num_trained()
assert actual_trained_in_epoch == expected_trained_in_epoch, (
"Unexpected number of trained (consumed minus filtered) data UIDs. Got: "
f"{actual_trained_in_epoch} != {expected_trained_in_epoch}"
)
if stop_training:
break
# 9. Per-epoch epilogue.
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 generator tasks for this epoch
for t in generator_tasks:
t.cancel()
try:
await asyncio.gather(*generator_tasks, return_exceptions=True)
except Exception:
pass
if generators_done_watcher is not None:
generators_done_watcher.cancel()
await asyncio.gather(generators_done_watcher, return_exceptions=True)
# Per-epoch reset/validation for data loading and staleness management
assert all(
t.done() for t in generator_tasks
), "Generator tasks must be done before resetting the dataloader manager and validating the staleness manager."
assert (
generation_output_group_buffer.qsize() == 0
), f"We expect all generation output to be consumed by the training worker at end of an epoch, got {generation_output_group_buffer.qsize()}."
await self.async_train_dataloader.reset_at_epoch_end()
await self._staleness_manager.validate_state_at_epoch_end(self.global_step)
# End of an epoch.
finally:
self._profiler_stop()
if self._ray_gpu_monitor is not None:
self._ray_gpu_monitor.stop()
pbar.close()
if not stop_training:
# All epochs completed: advance past the last epoch so resuming from the final checkpoint
# does not redo the last epoch (whose consumed UIDs were cleared at its epoch end).
self.epoch = self.cfg.trainer.epochs
# safety net: always save final checkpoint at end of training.
if self.cfg.trainer.ckpt_interval > 0:
with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
await asyncio.to_thread(self.save_checkpoints)
logger.info("Saved final checkpoint.")
if self.cfg.trainer.hf_save_interval > 0:
with self._phase_gauge.timed_phase("save_hf_model", self.all_timings):
await asyncio.to_thread(self.save_models)
logger.info("Saved final model.")
# Drain any in-flight async checkpoint write before teardown. Unconditional:
# a save may have happened outside the periodic path. No-op when nothing is pending.
self.dispatch.finalize_pending_saves("policy")
if self.has_critic:
self.dispatch.finalize_pending_saves("critic")
if self._vllm_metrics_scraper is not None:
await self._vllm_metrics_scraper.aclose()
self.tracker.finish()
logger.info("Training done!")
async def _drain_next_group(
self, buffer: asyncio.Queue, all_generators_done: asyncio.Event
) -> Optional[GeneratedOutputGroup]:
"""Return the next generated group, or None if generation is exhausted (all workers finished
and the buffer is empty).
Only used under ``sample_full_batch``, where dropping groups can exhaust the epoch mid
mini-batch and a plain blocking ``buffer.get()`` would hang forever.
"""
while True:
if not buffer.empty():
return buffer.get_nowait()
if all_generators_done.is_set():
return None
get_task = asyncio.ensure_future(buffer.get())
done_task = asyncio.ensure_future(all_generators_done.wait())
done, pending = await asyncio.wait({get_task, done_task}, return_when=asyncio.FIRST_COMPLETED)
if get_task in done:
for t in pending:
t.cancel()
return get_task.result()
# all_generators_done fired first. Cancel the pending get and loop to re-check the buffer.
# If the get had already pulled an item (racing the cancel), return it rather than drop it
# (a successful get() pops from the queue). No put can actually race here since the event is
# only set after all producers stop, but this stays correct regardless.
get_task.cancel()
try:
return await get_task
except asyncio.CancelledError:
pass
def _should_keep_group(self, group: GeneratedOutputGroup) -> bool:
"""Whether a group has reward variance (train on it) vs. is zero-variance (drop it).
Token-level rewards (``List[List[float]]``) are collapsed to a per-trajectory sum, like the rest
of the trainer. Groups with <=1 live trajectory (singletons / mostly masked) are always kept.
"""
rewards = group.generator_output["rewards"]
if rewards and isinstance(rewards[0], list):
# Token-level rewards: collapse each trajectory to a scalar sequence reward.
seq_rewards = [float(sum(r)) for r in rewards]
else:
seq_rewards = rewards
kept_indices = zero_variance_filter(
seq_rewards,
[group.uid] * len(seq_rewards),
loss_masks=group.generator_output["loss_masks"],
tol=self.cfg.trainer.algorithm.zero_variance_filter_tol,
)
return len(kept_indices) > 0
async def _collect_generation_mini_batch(
self,
generation_output_group_buffer: asyncio.Queue,
all_generators_done: asyncio.Event,
) -> Tuple[List[GeneratedOutputGroup], List[GeneratedOutputGroup], bool]:
"""Pull a full mini-batch of generated groups from the buffer.
Without ``sample_full_batch``, blocks until ``mini_batch_size`` groups are available. With it,
drops zero-variance groups (freeing their capacity, marking UIDs consumed) and keeps pulling
until ``mini_batch_size`` non-zero-variance groups are collected, which can exhaust the epoch.
Returns ``(kept_groups, dropped_groups, epoch_exhausted)``. On exhaustion the kept groups are a
(possibly empty) partial batch to discard; dropped groups are kept for metrics only.
"""
kept_groups: List[GeneratedOutputGroup] = []
dropped_groups: List[GeneratedOutputGroup] = []
epoch_exhausted = False
# Buffer occupancy when this step starts waiting.
self.all_metrics["async/gen_buffer_qsize_at_wait_start"] = generation_output_group_buffer.qsize()
with self._phase_gauge.timed_phase("wait_for_generation_buffer", self.all_timings):
buffer_pbar = tqdm(total=self.mini_batch_size, initial=0, desc="Generation Buffer Progress", position=1)
while len(kept_groups) < self.mini_batch_size:
# We do finish-time FIFO here (not schedule-time FIFO).
if not self.sample_full_batch:
kept_groups.append(await generation_output_group_buffer.get())
buffer_pbar.update(1)
buffer_pbar.set_postfix({"buffer qsize": generation_output_group_buffer.qsize()})
continue
group = await self._drain_next_group(generation_output_group_buffer, all_generators_done)
if group is None:
epoch_exhausted = True
break
try:
if self._should_keep_group(group):
kept_groups.append(group)
buffer_pbar.update(1)
buffer_pbar.set_postfix({"buffer qsize": generation_output_group_buffer.qsize()})
else:
# Drop the zero-variance group: give its capacity back to the producers and mark
# its UID consumed (skipped, not regenerated, on resume).
dropped_groups.append(group)
await self._staleness_manager.on_rollout_filtered()
await self.async_train_dataloader.mark_filtered_uids([group.uid])
except Exception:
self._log_group_processing_error(group, len(kept_groups), len(dropped_groups))
raise
buffer_pbar.close()
return kept_groups, dropped_groups, epoch_exhausted
@staticmethod
def _log_group_processing_error(group: GeneratedOutputGroup, kept_so_far: int, dropped_so_far: int) -> None:
"""Log the offending group's reward / loss-mask shape before a drain-loop error propagates,
flushing stderr (generator ``os._exit`` on teardown can otherwise drop buffered output)."""
go = group.generator_output
rewards = go.get("rewards") if isinstance(go, dict) else None
loss_masks = go.get("loss_masks") if isinstance(go, dict) else None
logger.exception(
"sample_full_batch: error processing generated group "
f"(uid={group.uid}, kept_so_far={kept_so_far}, dropped_so_far={dropped_so_far}). "
f"rewards type={type(rewards).__name__} "
f"len={len(rewards) if hasattr(rewards, '__len__') else 'n/a'} "
f"value={rewards!r}; "
f"loss_masks type={type(loss_masks).__name__} "
f"len={len(loss_masks) if hasattr(loss_masks, '__len__') else 'n/a'}"
)
sys.stderr.flush()
async def _run_training(self, training_input: TrainingInputBatch):
# TODO(Charlie): share this code with the one-step-off async trainer.
# inference and calculate values, log probs, rewards, kl divergence
with Timer("fwd_logprobs_values_reward", self.all_timings):
training_input = await asyncio.to_thread(self.fwd_logprobs_values_reward, training_input)
# calculate kl divergence and create experiences
if self.cfg.trainer.algorithm.use_kl_in_reward:
with Timer("apply_reward_kl_penalty", self.all_timings):
training_input = self.apply_reward_kl_penalty(training_input)
# calculate advantages and returns / along with tensorboard logging
with Timer("compute_advantages_and_returns", self.all_timings):
training_input = self.compute_advantages_and_returns(training_input)
# remove some unwanted keys
for key in ["rewards"]:
training_input.pop(key)
training_input.metadata.pop("uids")
if self.cfg.trainer.dump_data_batch:
# dump data to file
with Timer("dump_data_batch"):
self.dump_data(training_input, file_name=f"global_step_{self.global_step}_training_input")
# train policy/critic model
with Timer("train_critic_and_policy", self.all_timings):
status = await asyncio.to_thread(self.train_critic_and_policy, training_input)
return status
async def _run_generate_for_a_group_loop(self, generation_output_group_buffer: asyncio.Queue):
"""
Generator worker: repeatedly pulls the next prompt (possibly blocked by staleness control),
generates one single generation group, respecting a pause/resume event, and enqueues the result.
"""
try:
while True:
# 0. Pull next batch from dataloader. If returns None, then dataloader is exhausted.
rand_prompts = await self.async_train_dataloader.get_next_non_consumed_data()
if rand_prompts is None:
return
# 1. Prepare generator input
assert len(rand_prompts) == 1
generator_input, uids = prepare_generator_input(
rand_prompts,
self.cfg.generator.n_samples_per_prompt,
get_sampling_params_for_backend(
self.cfg.generator.inference_engine.backend, self.cfg.generator.sampling_params
),
self.cfg.environment.env_class,
"train",
self.global_step,
)
assert all(uid == uids[0] for uid in uids), "Expect all uids to be the same"
# 2. Acquire capacity slot.
slot_acquired = False
await self._staleness_manager.acquire_submission_slot()
slot_acquired = True
# 3. Generate one rollout group
global_step_at_start = self.global_step # for staleness control
group_start_time = time.monotonic()
if "disable_tqdm" in inspect.signature(self.generator.generate).parameters:
# A workaround to disable tqdm for the SkyRLGymGenerator.generate method which will
# blast the console with each worker's progress bar.
cur_generator_output: GeneratorOutput = await self.generator.generate(
generator_input, disable_tqdm=True
)
else:
cur_generator_output: GeneratorOutput = await self.generator.generate(generator_input)
group_completion_time_s = time.monotonic() - group_start_time
# 4. Enqueue the completed group and mark accepted to free capacity slot.
try:
generation_output_group_buffer.put_nowait(
GeneratedOutputGroup(
generator_output=cur_generator_output,
uid=uids[0],
global_step_when_scheduled=global_step_at_start,
group_completion_time_s=group_completion_time_s,
prompts=generator_input["prompts"],
)
)
except asyncio.QueueFull:
raise AssertionError("Generation buffer should never be full given staleness control.")
await self._staleness_manager.on_rollout_accepted()
slot_acquired = False
except asyncio.CancelledError:
# Expected on epoch end / shutdown: release any held slot so staleness accounting stays
# consistent, then exit cleanly. (Previously os._exit(1) here, which crashed the process and
# masked the real traceback when the cancel was triggered by a training-loop error.)
if "slot_acquired" in locals() and slot_acquired:
try:
await self._staleness_manager.on_rollout_rejected()
except Exception:
pass
return
except Exception as e:
logger.error(f"Generator worker errored out with exception: {e}")
logger.error(f"Traceback: \n{traceback.format_exc()}")
sys.stderr.flush() # flush before os._exit, which otherwise drops buffered output
os._exit(1)
@staticmethod
def _reprefix_metrics(metrics: dict, suffix: str) -> dict:
"""Re-prefix metric keys into a separate view (``generate/X`` -> ``generate_<suffix>/X``),
keeping the leading namespace so views group together in trackers."""
out = {}
for k, v in metrics.items():
if "/" in k:
prefix, rest = k.split("/", 1)
out[f"{prefix}_{suffix}/{rest}"] = v
else:
out[f"{suffix}/{k}"] = v
return out
def convert_generation_group_mini_batch_to_training_input(
self,
cur_generation_group_mini_batch: List[GeneratedOutputGroup],
dropped_groups: Optional[List[GeneratedOutputGroup]] = None,
) -> TrainingInputBatch:
"""Concatenate the mini-batch of generated groups and convert to a TrainingInputBatch.
``dropped_groups`` (zero-variance groups dropped this step under ``sample_full_batch``) are not
trained on, but are folded into the reward/generation metrics to keep them comparable across runs.
"""
dropped_groups = dropped_groups or []
generator_outputs = []
# Prompts for the kept groups, kept parallel to the concatenated `generator_output` so the
# trajectory logger can render prompt + response (mirrors the sync trainer's use of
# `generator_input["prompts"]`).
prompts: List[Any] = []
uids = []
stalenesses = []
staleness_violation_count = 0
# Per-group completion times (seconds) and the intra-group spread (std and coefficient of
# variation) of per-trajectory completion times. Helpful to understand tail latency
group_completion_times: List[float] = []
intra_group_stds: List[float] = []
intra_group_cvs: List[float] = []