-
Notifications
You must be signed in to change notification settings - Fork 429
Expand file tree
/
Copy pathsft_trainer.py
More file actions
2259 lines (1936 loc) · 98.1 KB
/
Copy pathsft_trainer.py
File metadata and controls
2259 lines (1936 loc) · 98.1 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
"""
SFT (Supervised Fine-Tuning) trainer for SkyRL.
Supports both FSDP and Megatron backends via a single ``SFTTrainer`` class.
The backend is selected dynamically based on ``SFTConfig.strategy``.
Usage::
from skyrl.train.config.sft_config import SFTConfig, SFTPlacementConfig
from skyrl.train.sft_trainer import SFTTrainer
cfg = SFTConfig(strategy="megatron")
trainer = SFTTrainer(cfg)
trainer.setup()
trainer.train()
trainer.shutdown()
Or as a CLI entrypoint::
python -m skyrl.train.main_sft strategy=megatron model.path=Qwen/Qwen3-0.6B
"""
import functools
import json
import multiprocessing as mp
import os
import tempfile
from dataclasses import asdict
from typing import Any, Optional
import numpy as np
import ray
import torch
from datasets import Dataset, load_dataset
from loguru import logger
from ray.util.placement_group import placement_group
from torchdata.stateful_dataloader import StatefulDataLoader
from transformers import AutoTokenizer
from skyrl.backends.skyrl_train.training_batch import (
TensorList,
TrainingInputBatch,
pad_training_input_batch,
)
from skyrl.backends.skyrl_train.utils.io import io
from skyrl.backends.skyrl_train.workers.worker import PPORayActorGroup
from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch
from skyrl.env_vars import SKYRL_RAY_PG_TIMEOUT_IN_S
from skyrl.train.config import SkyRLTrainConfig
from skyrl.train.config.sft_config import (
SFTConfig,
TrainOnWhat,
_normalize_dataset_cfg,
build_skyrl_config_for_sft,
)
from skyrl.train.dataset.pretokenized import load_from_pretokenized
from skyrl.train.dataset.sft_dataset import ConcatSFTDataset, SFTDataset, TextDataset
from skyrl.train.generators.utils import (
get_response_ids_and_loss_mask_from_messages,
)
from skyrl.train.utils import get_ray_pg_ready_with_timeout
from skyrl.train.utils.async_batch_collator import AsyncBatchCollator
from skyrl.train.utils.callbacks import (
CallbackHandler,
CallbackInput,
TrainingCallback,
TrainingControl,
)
from skyrl.train.utils.ray_gpu_monitor import RayGpuMonitor
from skyrl.train.utils.tracking import Tracking
from skyrl.train.utils.trainer_utils import (
GLOBAL_STEP_PREFIX,
cleanup_old_checkpoints,
extract_step_from_path,
validate_consistency_for_latest_checkpoint,
)
from skyrl.train.utils.utils import ResolvedPlacementGroup, Timer
from skyrl.utils.tok import (
check_is_vlm,
get_processor,
get_tokenizer,
)
# ---------------------------------------------------------------------------
# Tokenization helpers
# ---------------------------------------------------------------------------
def _tokenize_chat_slice_worker(args):
"""Worker function for parallel chat-format tokenization with slice-based loading.
Each worker loads the full dataset (HF caches it locally after the parent's
first call) and tokenizes only its assigned index range.
Must be top-level for pickling with spawn.
"""
(
dataset_name,
dataset_split,
start_idx,
end_idx,
tokenizer_path,
max_length,
messages_key,
train_on_what_str,
tools_key,
system_key,
) = args
# Worker loads tokenizer from cached path
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_path,
trust_remote_code=True,
use_fast=True,
local_files_only=True,
)
train_on_what = TrainOnWhat(train_on_what_str)
# Reload the dataset using the original split string and slice by index.
# The parent has already loaded once so this hits the HF cache.
dataset = load_dataset(dataset_name, split=dataset_split)
dataset_slice = dataset.select(range(start_idx, end_idx))
# Tokenize and filter inline
results = []
for example in dataset_slice:
tokenized = tokenize_chat_example(
example,
tokenizer,
max_length=max_length,
messages_key=messages_key,
train_on_what=train_on_what,
tools_key=tools_key,
system_key=system_key,
)
if tokenized is not None:
results.append(tokenized)
return results
def _tokenize_alpaca_slice_worker(args):
"""Worker function for parallel Alpaca-format tokenization with slice-based loading.
Each worker loads the full dataset (HF caches it locally after the parent's
first call) and tokenizes only its assigned index range.
Must be top-level for pickling with spawn.
"""
dataset_name, dataset_split, start_idx, end_idx, tokenizer_path, max_length = args
# Worker loads tokenizer from cached path
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_path,
trust_remote_code=True,
use_fast=True,
local_files_only=True,
)
# Reload the dataset using the original split string and slice by index.
dataset = load_dataset(dataset_name, split=dataset_split)
dataset_slice = dataset.select(range(start_idx, end_idx))
# Tokenize and filter inline
results = []
for example in dataset_slice:
tokenized = tokenize_sft_example(example, tokenizer, max_length)
if tokenized is not None:
results.append(tokenized)
return results
def _compute_cache_key(
dataset_name: str,
dataset_split: str,
model_path: str,
max_length: Optional[int],
messages_key: str,
train_on_what: str,
tools_key: Optional[str],
system_key: Optional[str],
) -> str:
"""Compute a cache key (hash) for a tokenized dataset.
The hash uniquely identifies the dataset and tokenization parameters so
that cached results can be safely reused when parameters match.
Args:
dataset_name: HuggingFace dataset name
dataset_split: Dataset split string (e.g., "train[:100000]")
model_path: Model name/path (tokenizer identity)
max_length: Maximum sequence length for truncation
messages_key: Column name for messages
train_on_what: Training target (last_assistant_message or all_assistant_messages)
tools_key: Column name for tools (if applicable)
system_key: Column name for system prompt (if applicable)
Returns:
A hex string hash (e.g., "a3f2c1...")
"""
import hashlib
# Build a deterministic string from all relevant parameters
cache_params = json.dumps(
{
"dataset_name": dataset_name,
"dataset_split": dataset_split,
"model_path": model_path,
"max_length": max_length,
"messages_key": messages_key,
"train_on_what": train_on_what,
"tools_key": tools_key,
"system_key": system_key,
},
sort_keys=True,
)
return hashlib.sha256(cache_params.encode()).hexdigest()[:16]
def _get_cache_path(cache_dir: str, cache_key: str) -> str:
"""Get the full path to a cached tokenized dataset.
The cache is stored as an arrow-backed HuggingFace dataset on disk
(``Dataset.save_to_disk``), so this path is a directory, not a file.
Args:
cache_dir: Base cache directory
cache_key: Cache key (hash) for this dataset
Returns:
Path to the cache directory (e.g., /path/to/cache/a3f2c1).
"""
return os.path.join(cache_dir, cache_key)
def _load_from_cache(cache_path: str) -> Optional[list]:
"""Load tokenized dataset from cache.
Reads an arrow-backed HF ``Dataset`` directory written by
:func:`_save_to_cache` and materializes it back to the ``list[dict]``
representation expected by the trainer (which slices, shuffles, and
concatenates the result during the training loop).
Args:
cache_path: Path to cached dataset directory.
Returns:
List of tokenized examples, or ``None`` if the cache directory
does not exist or fails to load.
"""
if not os.path.isdir(cache_path):
return None
try:
logger.info(f"Loading tokenized dataset from cache: {cache_path}")
dataset = Dataset.load_from_disk(cache_path)
tokenized = dataset.to_list()
logger.info(f"Loaded {len(tokenized)} examples from cache")
return tokenized
except Exception as e:
logger.warning(f"Failed to load cache from {cache_path}: {e}")
return None
def _save_to_cache(cache_path: str, tokenized: list) -> None:
"""Save tokenized dataset to cache.
Materializes the in-memory ``list[dict]`` as a HuggingFace ``Dataset``
and writes it via ``save_to_disk``. At 1M-row scale, the arrow-backed,
memory-mapped format reads and writes dramatically faster than pickle
while also being portable across Python versions. The write goes to a
sibling ``<cache_path>.tmp`` directory which is then atomically renamed
onto ``cache_path`` for NFS safety.
Args:
cache_path: Path to the cache directory to create.
tokenized: List of tokenized examples.
"""
try:
import shutil
parent_dir = os.path.dirname(cache_path)
os.makedirs(parent_dir, exist_ok=True)
logger.info(f"Saving {len(tokenized)} examples to cache: {cache_path}")
# Build the HF Dataset from rows and write to a sibling temp dir.
# An atomic rename onto cache_path makes concurrent readers see only
# a fully-written cache (NFS-safe; matches the previous pickle path).
dataset = Dataset.from_list(tokenized)
temp_path = cache_path + ".tmp"
# Clean up any stale temp dir from an interrupted prior run.
if os.path.isdir(temp_path):
shutil.rmtree(temp_path)
dataset.save_to_disk(temp_path)
# If a previous cache exists at the final path, drop it before
# rename so the swap is the only visible state change.
if os.path.isdir(cache_path):
shutil.rmtree(cache_path)
os.rename(temp_path, cache_path)
logger.info("Cache saved successfully")
except Exception as e:
logger.warning(f"Failed to save cache to {cache_path}: {e}")
@functools.lru_cache(maxsize=512)
def _parse_tools_str(tools: str) -> Optional[tuple]:
"""Parse a JSON-encoded tools string. Cached because tool-calling datasets
typically share one schema across thousands of rows (e.g. APIGen's airline
domain), and `apply_chat_template` re-tokenizes the schema on every row."""
tools = tools.strip()
if not tools:
return None
parsed = json.loads(tools)
if not parsed:
return None
# Return a tuple so the cache stores an immutable value; caller re-lists it.
return tuple(parsed)
def _coerce_tools(tools: Any) -> Optional[list]:
"""Coerce a dataset's ``tools`` field into a list[dict] for ``apply_chat_template``.
Tool-calling datasets ship the schema list as ``list[dict]`` (parquet-typed),
JSON-encoded ``str``, or absent. HF chat templates expect ``list[dict]``;
returns ``None`` when there are no tools so the caller can omit the kwarg.
"""
if tools is None:
return None
if isinstance(tools, str):
cached = _parse_tools_str(tools)
return list(cached) if cached else None
if isinstance(tools, list):
return tools or None
raise TypeError(f"Unsupported `tools` type: {type(tools).__name__}")
def _normalize_tool_call_payload(tc: Any) -> Optional[list]:
"""Normalize an assistant message's ``tool_calls`` into the OpenAI-style list.
Datasets in the wild use:
* ``[]`` / ``""`` / ``None`` — no tool call;
* a single JSON-encoded ``{"name": ..., "arguments": ...}`` dict (APIGen-MT);
* a JSON-encoded list of such dicts (xLAM, ToolACE);
* an already-parsed list of OpenAI-style ``{"type": "function", "function": {...}}``.
HF chat templates expect ``[{"type": "function", "function": {"name", "arguments"}}]``,
so we coerce to that shape. Returns ``None`` when the message has no tool call.
"""
if tc is None:
return None
if isinstance(tc, str):
tc = tc.strip()
if not tc or tc == "[]":
return None
tc = json.loads(tc)
if isinstance(tc, dict):
tc = [tc]
if not isinstance(tc, list) or not tc:
return None
out = []
for call in tc:
if not isinstance(call, dict):
raise TypeError(f"tool call entry must be a dict, got {type(call).__name__}")
fn = call["function"] if isinstance(call.get("function"), dict) else call
arguments = fn.get("arguments", {})
if isinstance(arguments, str):
try:
arguments = json.loads(arguments)
except json.JSONDecodeError:
pass
out.append({"type": "function", "function": {"name": fn.get("name"), "arguments": arguments}})
return out
# Fields we explicitly normalize; everything else on the message is forwarded
# so model-specific templates can read e.g. ``reasoning_content`` (Qwen3,
# Nemotron-3 thinking models) or ``tool_call_id``.
_NORMALIZED_KEYS = frozenset({"role", "content", "tool_calls"})
def _normalize_chat_messages(messages: list[dict]) -> list[dict]:
"""Return messages in a shape that HF chat templates accept.
Normalizes ``content`` (``None`` → ``""``), promotes string-encoded
``tool_calls`` on assistant turns into the OpenAI list-of-dicts form,
drops empty/placeholder ``tool_calls`` on every role, and preserves any
other fields on the message verbatim.
"""
out = []
for msg in messages:
role = msg["role"]
new_msg = {k: v for k, v in msg.items() if k not in _NORMALIZED_KEYS}
new_msg["role"] = role
new_msg["content"] = msg.get("content", "") or ""
if role == "assistant":
tool_calls = _normalize_tool_call_payload(msg.get("tool_calls"))
if tool_calls:
new_msg["tool_calls"] = tool_calls
out.append(new_msg)
return out
def tokenize_sft_example(example: dict, tokenizer, max_length: int = 512, **tokenizer_kwargs) -> dict | None:
"""Tokenize an Alpaca-format SFT example via ``apply_chat_template``.
Converts the instruction/input/output fields into a two-message chat
(user + assistant) and delegates to :func:`tokenize_chat_example`.
This ensures tokenization matches the HF / TRL convention (proper
special tokens, chat template formatting).
Returns dict with input_ids, attention_mask, num_actions (response length),
or None if the example was fully truncated.
"""
instruction = example.get("instruction", "")
input_text = example.get("input", "")
output = example.get("output", "")
# Build user content: instruction + optional input
user_content = instruction
if input_text:
user_content = f"{instruction}\n\n{input_text}"
user_content = user_content.strip()
messages = [
{"role": "user", "content": user_content},
{"role": "assistant", "content": output},
]
return tokenize_chat_example(
{"messages": messages},
tokenizer,
max_length=max_length,
messages_key="messages",
**tokenizer_kwargs,
)
def tokenize_chat_example(
example: dict,
tokenizer,
max_length: Optional[int] = None,
messages_key: str = "messages",
train_on_what: TrainOnWhat = TrainOnWhat.LAST_ASSISTANT_MESSAGE,
tools_key: Optional[str] = "tools",
system_key: Optional[str] = "system",
processor=None,
**tokenizer_kwargs,
) -> dict | None:
"""Tokenize a chat-format example with configurable loss targets.
Uses ``apply_chat_template`` to tokenize the conversation and determine
which tokens to train on based on ``train_on_what``.
For tool-calling datasets (e.g. APIGen-MT, xLAM, ToolACE), if the example
has a ``tools_key`` column, its parsed value is forwarded as ``tools=`` to
``apply_chat_template`` so the model sees the function schemas. Likewise,
a ``system_key`` column is prepended as a leading system message when no
system message is already present.
Args:
example: Dict containing a ``messages_key`` column with chat messages.
tokenizer: HuggingFace tokenizer with ``apply_chat_template``.
max_length: Maximum sequence length (truncation boundary).
messages_key: Key in *example* that holds the messages list.
train_on_what: Which tokens to compute loss on.
tools_key: Key in *example* whose value is the per-row tool schema list
(or JSON-encoded string thereof). ``None`` disables the lookup.
system_key: Key in *example* whose value is a system prompt string.
``None`` disables the lookup.
**tokenizer_kwargs: Extra kwargs forwarded to ``apply_chat_template``
(e.g. ``enable_thinking``).
Returns:
Dict with ``input_ids``, ``attention_mask``, ``num_actions``, and
optionally ``loss_mask`` (a per-token list of 0/1 within the action
window). Returns ``None`` when the example should be skipped.
"""
# Validate supported modes
_SUPPORTED = {TrainOnWhat.LAST_ASSISTANT_MESSAGE, TrainOnWhat.ALL_ASSISTANT_MESSAGES}
if train_on_what not in _SUPPORTED:
raise NotImplementedError(
f"train_on_what={train_on_what!r} is not yet supported. "
f"Supported values: {sorted(v.value for v in _SUPPORTED)}"
)
messages = list(example[messages_key])
# Trim trailing tool observations with no follow-up assistant response
# (common in APIGen-MT). Trailing user turns still drop the row.
while messages and messages[-1]["role"] == "tool":
messages.pop()
if not messages or messages[-1]["role"] != "assistant":
return None
messages = _normalize_chat_messages(messages)
system_prompt = example.get(system_key) if system_key else None
if system_prompt and messages[0].get("role") != "system":
messages = [{"role": "system", "content": system_prompt}] + messages
# Per-row schemas yield to an explicit caller-provided ``tools`` kwarg.
tools = _coerce_tools(example.get(tools_key)) if tools_key else None
if tools is not None:
tokenizer_kwargs = {"tools": tools, **tokenizer_kwargs}
# Detect image content. VLM tokenization runs through the HF processor and
# only supports last-assistant training (image token positions are tied to
# a single forward over the full conversation).
has_images = any(
isinstance(m.get("content"), list)
and any(isinstance(d, dict) and d.get("type") == "image" for d in m["content"])
for m in messages
)
if has_images and train_on_what == TrainOnWhat.ALL_ASSISTANT_MESSAGES:
raise NotImplementedError("Training on all assistant messages with vision inputs is not yet supported")
if train_on_what == TrainOnWhat.LAST_ASSISTANT_MESSAGE:
return _tokenize_chat_last_assistant(
messages, tokenizer, max_length, processor if has_images else None, **tokenizer_kwargs
)
else:
# ALL_ASSISTANT_MESSAGES
return _tokenize_chat_all_assistants(messages, tokenizer, max_length, **tokenizer_kwargs)
def _unbatch(proc_tok_output):
"""Strip the outer batch list from a processor output.
``processor.apply_chat_template`` returns batched ids (``list[list[int]]``)
while a plain tokenizer returns a flat ``list[int]``. Normalize to flat.
"""
if proc_tok_output and isinstance(proc_tok_output[0], list):
return proc_tok_output[0]
return proc_tok_output
def _tokenize_chat_last_assistant(
messages: list[dict],
tokenizer,
max_length: Optional[int] = None,
processor=None,
**tokenizer_kwargs,
) -> dict | None:
"""Tokenize a conversation and compute loss only on the last assistant message.
Args:
messages: Full conversation (must end with an assistant message).
tokenizer: HuggingFace tokenizer with ``apply_chat_template``.
max_length: Optional sequence length cap; truncates both prompt and full
conversation to this limit.
processor: Optional HF processor used for VLM examples. When provided,
image tensors (``pixel_values``, ``image_grid_thw``) are produced and
returned alongside the token ids.
**tokenizer_kwargs: Extra kwargs forwarded to ``apply_chat_template``.
Returns:
Dict with ``input_ids``, ``attention_mask``, ``num_actions`` (number of
last-assistant tokens), and, for VLM examples, ``pixel_values`` and
``image_grid_thw``. Returns ``None`` if truncation left no response
tokens or dropped an image's placeholder tokens.
"""
# The processor (when present) tokenizes text and images together; otherwise
# the plain tokenizer handles text. ``return_dict=True`` so we can read both
# token ids and any image tensors back out.
processing_class = processor or tokenizer
length_kwargs = {}
if max_length is not None and processor is None:
length_kwargs = dict(truncation=True, max_length=max_length)
# Tokenize prompt (everything except last assistant message)
prompt_ids = processing_class.apply_chat_template(
messages[:-1],
add_generation_prompt=True,
tokenize=True,
return_dict=True,
**length_kwargs,
**tokenizer_kwargs,
)
# Tokenize full conversation
full_ids = processing_class.apply_chat_template(
messages,
add_generation_prompt=False,
tokenize=True,
return_dict=True,
**length_kwargs,
**tokenizer_kwargs,
)
full_input_ids = _unbatch(full_ids["input_ids"])
full_prompt_ids = _unbatch(prompt_ids["input_ids"])
# VLM samples can't be safely truncated (it would drop image placeholder tokens
# and break image/text alignment), so drop anything that exceeds the limit.
if processor is not None and max_length is not None and len(full_input_ids) > max_length:
logger.warning(
f"Dropping VLM sample longer than max_length={max_length}, consider increasing max_length if you see this warning too much"
)
return None
vlm_kwargs = {} # We only support Qwen-style image kwargs at the moment
if "pixel_values" in full_ids and "image_grid_thw" in full_ids:
vlm_kwargs = dict(
pixel_values=full_ids["pixel_values"],
image_grid_thw=full_ids["image_grid_thw"],
)
num_actions = len(full_input_ids) - len(full_prompt_ids)
if num_actions <= 0:
return None
return {
"input_ids": full_input_ids,
"attention_mask": [1] * len(full_input_ids),
"num_actions": num_actions,
"loss_mask": [1] * num_actions,
**vlm_kwargs,
}
def _tokenize_chat_all_assistants(
messages: list[dict],
tokenizer,
max_length: Optional[int] = None,
**tokenizer_kwargs,
) -> dict | None:
"""Tokenize a conversation and compute loss on all assistant messages.
Builds a per-token loss mask covering every assistant turn. ``num_actions``
spans from the first assistant token to the end of the conversation, with
interior 0s masking out user/system tokens between assistant turns.
Args:
messages: Full conversation. May start with system/user messages;
must contain at least one assistant message.
tokenizer: HuggingFace tokenizer with ``apply_chat_template``.
max_length: Optional sequence length cap; truncates to this limit.
**tokenizer_kwargs: Extra kwargs forwarded to ``apply_chat_template``.
Returns:
Dict with ``input_ids``, ``attention_mask``, ``num_actions``, and
``loss_mask`` (per-token 0/1 list within the action window), or
``None`` if no assistant tokens survived after truncation.
"""
# Find the index of the first assistant message.
i = 0
while i < len(messages) and messages[i]["role"] != "assistant":
i += 1
# Encode leading non-assistant messages separately because
# `get_response_ids_and_loss_mask_from_messages` does not accept system messages.
initial_token_ids = tokenizer.apply_chat_template(
messages[:i],
add_generation_prompt=False,
tokenize=True,
return_dict=False,
**tokenizer_kwargs,
)
# no assistant message
if i >= len(messages):
return None
later_token_ids, loss_mask, _ = get_response_ids_and_loss_mask_from_messages(
messages[i:], tokenizer, tokenizer_kwargs=tokenizer_kwargs
)
input_ids = initial_token_ids + later_token_ids
# truncate
if max_length is not None:
input_ids = input_ids[:max_length]
max_assistant_length = max(max_length - len(initial_token_ids), 0)
loss_mask = loss_mask[:max_assistant_length]
if sum(loss_mask) == 0:
return None # No assistant tokens survived truncation
num_actions = len(loss_mask)
return {
"input_ids": input_ids,
"attention_mask": [1] * len(input_ids),
"num_actions": num_actions,
"loss_mask": loss_mask,
}
# ---------------------------------------------------------------------------
# Collation
# ---------------------------------------------------------------------------
def collate_sft_batch(examples: list, tokenizer) -> TrainingInputBatch:
"""Collate tokenized examples into a TrainingInputBatch.
Creates the batch format expected by forward_backward with cross_entropy loss:
- sequences: [batch_size, seq_len] - token IDs (left-padded)
- attention_mask: [batch_size, seq_len] - 1 for real tokens, 0 for padding
- loss_mask: [batch_size, num_actions] - 1 for tokens to compute loss on
All examples are expected to carry a ``loss_mask`` key (guaranteed by both
``_tokenize_chat_last_assistant`` and ``_tokenize_chat_all_assistants``).
"""
max_len = max(len(ex["input_ids"]) for ex in examples)
max_num_actions = max(ex["num_actions"] for ex in examples)
num_examples = len(examples)
# Fill NumPy buffers by slice, then convert once.
sequences_np = np.full((num_examples, max_len), tokenizer.pad_token_id, dtype=np.int64)
attention_mask_np = np.zeros((num_examples, max_len), dtype=np.int64)
loss_mask_np = np.zeros((num_examples, max_num_actions), dtype=np.int64)
# VLM image tensors travel as a TensorList (one variable-shape tensor per
# sample). Mixed text+image batches are not supported; every sample in a VLM
# batch must carry images. Check homogeneity up front so a mixed batch fails
# here with a clear message rather than a KeyError deep in the pad loop.
num_with_images = sum("pixel_values" in ex for ex in examples)
if num_with_images not in (0, len(examples)):
raise ValueError(
f"Mixed text+image batches are not supported: {num_with_images}/{len(examples)} "
"samples carry 'pixel_values'. Every sample in a VLM batch must carry images."
)
batch_has_images = num_with_images > 0
pixel_values = []
image_grid_thw = []
for i, ex in enumerate(examples):
# Left-pad sequences; right-align response loss masks.
pad_len = max_len - len(ex["input_ids"])
sequences_np[i, pad_len:] = ex["input_ids"]
attention_mask_np[i, pad_len:] = ex["attention_mask"]
action_pad = max_num_actions - ex["num_actions"]
loss_mask_np[i, action_pad:] = ex["loss_mask"]
if batch_has_images:
pixel_values.append(torch.as_tensor(ex["pixel_values"]))
image_grid_thw.append(torch.as_tensor(ex["image_grid_thw"]))
batch = TrainingInputBatch(
{
"sequences": torch.from_numpy(sequences_np),
"attention_mask": torch.from_numpy(attention_mask_np),
"loss_mask": torch.from_numpy(loss_mask_np),
"pixel_values": TensorList(pixel_values) if batch_has_images else None,
"image_grid_thw": TensorList(image_grid_thw) if batch_has_images else None,
}
)
batch.metadata = {"response_length": max_num_actions}
return batch
def collate_sft_examples(
examples: list, collator, batch_size: int, pad_to_batch_size: bool = False
) -> TrainingInputBatch:
"""Top-level collate function for the SFT ``StatefulDataLoader``.
Defined at module scope (not a lambda/closure) so it is picklable when the
dataloader uses worker processes with the ``spawn`` start method. Delegates
to the trainer's configured ``collator`` (``DefaultCollator`` or
``PackedDataCollator``).
When ``pad_to_batch_size`` is set (the train path, non-packed), a final
short batch is padded up to ``batch_size`` rows via
:func:`pad_training_input_batch`, which zeros ``loss_mask`` on the padded
rows so they contribute no gradient. This lets every example in an epoch be
trained on (instead of dropping the tail) while still dispatching a full,
evenly-shardable ``batch_size`` batch. Packed batches are never row-padded
(their rows are FFD bins, already rounded to a multiple of ``dp_size``).
"""
batch = collator(examples, batch_size=batch_size)
if pad_to_batch_size:
pad_rows = batch_size - len(examples)
if pad_rows > 0:
batch = pad_training_input_batch(batch, pad_rows)
return batch
# ---------------------------------------------------------------------------
# SFTTrainer
# ---------------------------------------------------------------------------
def _format_eval_metrics(eval_metrics: dict) -> str:
"""Render per-dataset eval metrics (``{name}/loss``) for stdout logging."""
return ", ".join(f"{k}={v:.4f}" for k, v in eval_metrics.items())
class SFTTrainer:
"""SFT trainer supporting FSDP and Megatron backends.
Unlike RayPPOTrainer, this does NOT subclass it. SFT's concerns are
fundamentally different: no generation, no critic, no advantages, no
KL penalty. Sharing a base class would create confusing dead code paths.
Usage::
trainer = SFTTrainer(SFTConfig(strategy="megatron"))
trainer.setup()
trainer.train()
trainer.shutdown()
"""
def __init__(
self,
cfg: SFTConfig,
skyrl_cfg: SkyRLTrainConfig | None = None,
callbacks: Optional[list[TrainingCallback]] = None,
):
self.sft_cfg = cfg
_normalize_dataset_cfg(cfg)
# Accept a pre-built bridge config to avoid redundant rebuilds.
# When not provided (e.g. standalone usage), build it here.
self.cfg = skyrl_cfg if skyrl_cfg is not None else build_skyrl_config_for_sft(cfg)
self.tokenizer = None
self.processor = None # set in setup() for VLM models
self.is_vlm = False
self.dispatch: WorkerDispatch | None = None
self.tracker: Tracking | None = None
# Stateful dataloaders, built in train() once data is tokenized.
self.train_dataloader: StatefulDataLoader | None = None
# One ``(name, dataloader)`` pair per configured eval dataset; ``None``
# when eval is disabled. Names are unique (enforced in config validation)
# and namespace the eval metrics as ``eval/{name}/...``.
self.eval_dataloaders: list[tuple[str, StatefulDataLoader]] | None = None
self._checkpoint_dataloader_state: dict | None = None
self.global_step = 0
# running count of total non-padding tokens trained on
self._total_tokens_processed = 0
self.collator = None # built in setup() once the tokenizer is available
self._num_training_gpus: int = cfg.placement.num_nodes * cfg.placement.num_gpus_per_node
self._ray_gpu_monitor = RayGpuMonitor() if cfg.enable_ray_gpu_monitor else None
self._callback_handler = CallbackHandler(callbacks)
self._training_control = TrainingControl()
# Loop metadata used to build CallbackInput. Populated in train().
self._total_steps: int = 0
self._steps_per_epoch: int = 0
self._current_epoch: int = 0
@property
def _torch_profiler_enabled(self) -> bool:
"""Whether to dispatch policy profiler RPCs."""
return self.cfg.trainer.policy.torch_profiler_config.enable
def _build_collator(self, tokenizer):
"""Select the batch collator from the configured packing mode.
``PackedDataCollator`` performs controller-level FFD bin-packing
(Megatron-only, ``use_sequence_packing=True``); ``DefaultCollator``
left-pads each example. The choice is fixed by static config; the
``tokenizer`` is passed in by :meth:`setup` once it is available. The
packed config is validated here.
"""
# Imported lazily to avoid a circular import: ``collators`` imports
# ``collate_sft_batch`` from this module.
from skyrl.train.dataset.collators import DefaultCollator, PackedDataCollator
if self.sft_cfg.use_sequence_packing:
from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import (
is_fp8_enabled,
)
self._validate_packing_cfg()
transformer_config_kwargs = self.sft_cfg.megatron_config.transformer_config_kwargs or {}
return PackedDataCollator(
tokenizer=tokenizer,
max_tokens_per_microbatch=self.sft_cfg.resolved_bin_capacity(),
tp_size=self.sft_cfg.megatron_config.tensor_model_parallel_size,
pp_size=self.sft_cfg.megatron_config.pipeline_model_parallel_size,
cp_size=self.sft_cfg.megatron_config.context_parallel_size,
dp_size=self._dp_size(),
batch_size=self.sft_cfg.batch_size,
micro_train_batch_size_per_gpu=self.sft_cfg.micro_train_batch_size_per_gpu,
fp8_enabled=is_fp8_enabled(transformer_config_kwargs.get("fp8")),
fp8_recipe=transformer_config_kwargs.get("fp8_recipe"),
)
return DefaultCollator(
tokenizer=tokenizer,
micro_train_batch_size_per_gpu=self.sft_cfg.micro_train_batch_size_per_gpu,
)
def _dp_size(self) -> int:
"""Number of DP ranks under the configured Megatron parallelism."""
total_gpus = self.sft_cfg.placement.num_nodes * self.sft_cfg.placement.num_gpus_per_node
tp = self.sft_cfg.megatron_config.tensor_model_parallel_size
pp = self.sft_cfg.megatron_config.pipeline_model_parallel_size
cp = self.sft_cfg.megatron_config.context_parallel_size
return total_gpus // (tp * pp * cp)
def _validate_packing_cfg(self):
"""Validate the config when ``use_sequence_packing=True``."""
if self.sft_cfg.strategy != "megatron":
raise ValueError(
f"use_sequence_packing=True only supports strategy='megatron'; got "
f"{self.sft_cfg.strategy!r}. Use the FSDP packing path instead."
)
# Sequence packing needs the THD layout, so it implies
# remove_microbatch_padding=True. Auto-enable it (warning if the user
# explicitly set it False) instead of erroring on the contradiction.
if not self.sft_cfg.remove_microbatch_padding:
logger.warning(
"use_sequence_packing=True requires the THD layout; "
"setting remove_microbatch_padding=True (was False)."
)
self.sft_cfg.remove_microbatch_padding = True
# ------------------------------------------------------------------ #
# Setup
# ------------------------------------------------------------------ #
def setup(self):
"""Initialize tokenizer, workers, dispatch, and tracker.
Ray must already be initialized before calling this (either via
``initialize_ray`` on the head node or inside a Ray task).
"""
tokenizer_kwargs = {
"trust_remote_code": True,
"use_fast": not self.cfg.trainer.disable_fast_tokenizer,
"padding_side": "left",
}
self.is_vlm = check_is_vlm(self.cfg.trainer.policy.model.path)
if self.is_vlm:
self.processor = get_processor(self.cfg.trainer.policy.model.path, **tokenizer_kwargs)
# Sequence packing / microbatch padding removal are unsupported for
# VLMs (3D RoPE + image token positions). ``remove_microbatch_padding``
# defaults to True, so disable both unconditionally and mirror the
# change onto the already-built trainer config the workers receive.
if self.sft_cfg.use_sequence_packing or self.sft_cfg.remove_microbatch_padding:
logger.warning("VLM detected: disabling sequence packing / microbatch padding removal.")
self.sft_cfg.use_sequence_packing = False
self.sft_cfg.remove_microbatch_padding = False
self.cfg.trainer.remove_microbatch_padding = False
self.tokenizer = get_tokenizer(self.cfg.trainer.policy.model.path, **tokenizer_kwargs)
self.collator = self._build_collator(self.tokenizer)
self._init_tracker()
self._init_workers()
def _init_workers(self):
"""Create PPORayActorGroup and WorkerDispatch.
Selects the correct PolicyWorker based on strategy.
"""
if self.sft_cfg.strategy == "megatron":
from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import (
PolicyWorker,
)
else:
from skyrl.backends.skyrl_train.workers.fsdp.fsdp_worker import PolicyWorker
num_gpus = self.sft_cfg.placement.num_gpus_per_node
raw_pg = placement_group(
[{"GPU": num_gpus, "CPU": num_gpus}] * self.sft_cfg.placement.num_nodes,
strategy="PACK",
)
get_ray_pg_ready_with_timeout(raw_pg, timeout=SKYRL_RAY_PG_TIMEOUT_IN_S)
pg = ResolvedPlacementGroup(raw_pg)
actor_group = PPORayActorGroup(
self.cfg.trainer,
num_nodes=self.sft_cfg.placement.num_nodes,
num_gpus_per_node=num_gpus,
ray_actor_type=PolicyWorker,
pg=pg,
num_gpus_per_actor=1,
colocate_all=False,
sequence_parallel_size=self.cfg.trainer.policy.sequence_parallel_size,
record_memory=self.cfg.trainer.policy.record_memory,
)
num_training_steps = (
self.sft_cfg.dummy_run_max_steps if self.sft_cfg.dummy_run_full_ctx else self.sft_cfg.num_steps
)
if self.sft_cfg.max_training_steps is not None:
num_training_steps = (
self.sft_cfg.max_training_steps
if num_training_steps is None
else min(num_training_steps, self.sft_cfg.max_training_steps)
)
# num_steps may be None when num_epochs is used; without an explicit cap,
# the worker will use its default large value for the LR scheduler.
ray.get(
actor_group.async_init_model(
self.sft_cfg.model.path,
num_training_steps=num_training_steps,
)
)
ray.get(actor_group.async_run_ray_method("pass_through", "_set_pad_token_id", self.tokenizer.pad_token_id))
self.dispatch = WorkerDispatch(self.cfg, policy_actor_group=actor_group)
def _init_tracker(self):