[train] Fix issue with unset pad_token_id (#1232)

# What does this PR do?

Fixes #1231 
<!-- devin-review-badge-begin -->

---

<a href="https://app.devin.ai/review/novasky-ai/skyrl/pull/1232"
target="_blank">
  <picture>
<source media="(prefers-color-scheme: dark)"
srcset="https://static.devin.ai/assets/gh-open-in-devin-review-dark.svg?v=1">
<img
src="https://static.devin.ai/assets/gh-open-in-devin-review-light.svg?v=1"
alt="Open with Devin">
  </picture>
</a>
<!-- devin-review-badge-end -->

---------

Signed-off-by: SumanthRH <sumanthrh99@gmail.com>

Author: SumanthRH

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.distributed
 import ray
-from transformers import AutoTokenizer, AutoConfig
+from transformers import AutoConfig
 from huggingface_hub import snapshot_download
 
 import os
@@ -42,7 +42,7 @@
 from skyrl.backends.skyrl_train.workers.megatron.megatron_model_wrapper import MegatronModelWrapper
 from skyrl.backends.skyrl_train.utils.profiler import Profiler
 from skyrl.backends.skyrl_train.weight_sync import WeightExtractor, WeightChunk
-
+from skyrl.utils.tok import get_tokenizer
 
 if TYPE_CHECKING:
     from skyrl.backends.skyrl_train.inference_engines.base import InferenceEngineInterface
@@ -205,7 +205,7 @@ def init_configs(
         """
         Initialize the Megatron-Bridge bridge and provider objects + hf_config and tokenizer
         """
-        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer = get_tokenizer(model_path, trust_remote_code=True)
         hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
 
         override_config_kwargs = {

skyrl/backends/skyrl_train_backend.py

@@ -30,6 +30,7 @@
     create_ray_wrapped_inference_engines,
 )
 from skyrl.backends.skyrl_train.inference_engines.inference_engine_client import InferenceEngineClient
+from skyrl.utils.tok import get_tokenizer
 
 
 class SkyRLTrainBackendOverrides(BaseModel, extra="allow"):
@@ -112,7 +113,7 @@ def __init__(self, base_model: str, config: SkyRLTrainBackendOverrides):
         self._model_metadata: types.ModelMetadata | None = None
         self._cfg = None
         self._dispatch: WorkerDispatch | None = None
-        self._tokenizer = AutoTokenizer.from_pretrained(self.base_model)
+        self._tokenizer: AutoTokenizer = get_tokenizer(self.base_model)
         self._inference_engine_client = None
         self._inference_engines_initialized = False
 

skyrl/train/entrypoints/main_base.py

@@ -4,7 +4,7 @@
 
 from ray.util.placement_group import placement_group, PlacementGroup
 
-from transformers import AutoTokenizer, PreTrainedTokenizerBase
+from transformers import PreTrainedTokenizerBase
 from skyrl.train.dataset import PromptDataset
 from skyrl.train.utils import validate_cfg
 
@@ -24,6 +24,7 @@
 import os
 from loguru import logger
 from skyrl.train.utils.tracking import Tracking
+from skyrl.utils.tok import get_tokenizer
 import multiprocessing as mp
 import asyncio
 
@@ -122,7 +123,12 @@ def __init__(self, cfg: SkyRLTrainConfig):
             cfg: The fully resolved SkyRLTrainConfig instance.
         """
         self.cfg = cfg
-        self.tokenizer = self.get_tokenizer()
+        self.tokenizer = get_tokenizer(
+            self.cfg.trainer.policy.model.path,
+            trust_remote_code=True,
+            use_fast=not self.cfg.trainer.disable_fast_tokenizer,
+            padding_side="left",
+        )
         self.train_dataset = self.get_train_dataset()
         self.eval_dataset = self.get_eval_dataset()
         self.colocate_pg = self.get_colocate_pg()
@@ -135,19 +141,6 @@ def __init__(self, cfg: SkyRLTrainConfig):
     def get_cfg_as_str(cfg: SkyRLTrainConfig) -> str:
         return get_config_as_yaml_str(cfg)
 
-    def get_tokenizer(self, padding_side="left"):
-        """Initializes a tokenizer for the given model."""
-        tokenizer = AutoTokenizer.from_pretrained(
-            self.cfg.trainer.policy.model.path,
-            trust_remote_code=True,
-            use_fast=not self.cfg.trainer.disable_fast_tokenizer,
-        )
-        tokenizer.padding_side = padding_side
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-            tokenizer.pad_token_id = tokenizer.eos_token_id
-        return tokenizer
-
     def get_train_dataset(self):
         """Initializes the training dataset.
 

skyrl/utils/tok.py

@@ -0,0 +1,14 @@
+"""Tokenization related utilities"""
+
+from transformers import AutoTokenizer
+
+
+def get_tokenizer(model_name_or_path, **tokenizer_kwargs) -> AutoTokenizer:
+    """Gets tokenizer for the given base model with the given parameters
+
+    Sets the pad token ID to EOS token ID if `None`"""
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, **tokenizer_kwargs)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer

tests/backends/skyrl_train/gpu/utils.py

@@ -35,6 +35,7 @@
 from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import RemoteInferenceClient
 from skyrl.backends.skyrl_train.inference_servers.server_group import ServerGroup
 from skyrl.backends.skyrl_train.inference_servers.router import InferenceRouter
+from skyrl.utils.tok import get_tokenizer
 
 TEST_DATA_PATH = os.path.expanduser("~/data/gsm8k/validation.parquet")
 
@@ -460,7 +461,7 @@ def create(
         # Extract served_model_name from config if set
         served_model_name = ie_cfg.served_model_name
 
-        tokenizer = AutoTokenizer.from_pretrained(cfg.trainer.policy.model.path)
+        tokenizer = get_tokenizer(cfg.trainer.policy.model.path)
 
         # Return both router and server group if created to keep references alive
         router = None