feat(train): Add SequenceLength support for SFT, DPO, RLVR, RLAIF trainers

sagemaker-core/sample/sagemaker/2017-07-24/service-2.json

@@ -44132,6 +44132,10 @@
         "EvaluatorArn":{
           "shape":"EvaluatorArn",
           "documentation":"<p> The evaluator Amazon Resource Name (ARN) used as reward function or reward prompt. </p>"
+        },
+        "SequenceLength":{
+          "shape":"SequenceLength",
+          "documentation":"<p> The sequence length for the training job. </p>"
         }
       },
       "documentation":"<p> The configuration for the serverless training job. </p>"
@@ -44143,6 +44147,19 @@
         "Evaluation"
       ]
     },
+    "SequenceLength":{
+      "type":"string",
+      "enum":[
+        "1K",
+        "2K",
+        "4K",
+        "8K",
+        "16K",
+        "32K",
+        "64K",
+        "128K"
+      ]
+    },
     "ServerlessMaxConcurrency":{
       "type":"integer",
       "box":true,

sagemaker-core/src/sagemaker/core/resources.py

@@ -46,7 +46,6 @@
 from sagemaker.core.serializers.base import BaseSerializer
 from sagemaker.core.deserializers.base import BaseDeserializer
 
-
 logger = get_textual_rich_logger(__name__)
 
 

sagemaker-core/src/sagemaker/core/shapes/shapes.py

@@ -9781,6 +9781,7 @@ class ServerlessJobConfig(Base):
     peft:  The parameter-efficient fine-tuning configuration.
     evaluation_type:  The evaluation job type. Required when serverless job type is Evaluation.
     evaluator_arn:  The evaluator Amazon Resource Name (ARN) used as reward function or reward prompt.
+    sequence_length:  The sequence length for the training job.
     """
 
     base_model_arn: StrPipeVar
@@ -9790,6 +9791,7 @@ class ServerlessJobConfig(Base):
     peft: Optional[StrPipeVar] = Unassigned()
     evaluation_type: Optional[StrPipeVar] = Unassigned()
     evaluator_arn: Optional[StrPipeVar] = Unassigned()
+    sequence_length: Optional[StrPipeVar] = Unassigned()
 
 
 class MlflowConfig(Base):

sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py

@@ -16272,6 +16272,7 @@
             {"name": "Peft", "shape": "Peft", "type": "string"},
             {"name": "EvaluationType", "shape": "EvaluationType", "type": "string"},
             {"name": "EvaluatorArn", "shape": "EvaluatorArn", "type": "string"},
+            {"name": "SequenceLength", "shape": "SequenceLength", "type": "string"},
         ],
         "type": "structure",
     },

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -568,6 +568,10 @@ def _yaml_safe_default(value):
                 if hp_value is not None:
                     _set_spec_default(override_spec, hp_key, _yaml_safe_default(hp_value))
 
+        # Enforce selected lengths fit the recipe's supported sequence length.
+        if hasattr(self.hyperparameters, "validate_length_constraints"):
+            self.hyperparameters.validate_length_constraints()
+
         # Build hyperparameters early to inject into recipe template before runtime.
         final_hyperparameters = self.hyperparameters.to_dict()
         _validate_hyperparameter_values(final_hyperparameters)

sagemaker-train/src/sagemaker/train/common.py

@@ -22,16 +22,60 @@ class CustomizationTechnique(Enum):
 
 class FineTuningOptions:
     """Dynamic class for fine-tuning options with validation."""
-    
-    def __init__(self, options_dict: Dict[str, Any]):
+
+    def __init__(self, options_dict: Dict[str, Any], sequence_length: int = None):
         self._specs = options_dict.copy()
         self._user_set = set()
         self._initialized = False
+        # Recipe's supported sequence-length ceiling (prompt + response for RL,
+        # or max_length for SFT/DPO). Used by validate_length_constraints().
+        self._sequence_length = sequence_length
         # Extract default values and set as attributes (no validation during init)
         for key, spec in options_dict.items():
             default_value = spec.get('default') if isinstance(spec, dict) else spec
             super().__setattr__(key, default_value)
         self._initialized = True
+
+    # SFT/DPO recipes express the single-example length ceiling as
+    # "dataset_max_len" (verl and llmft templates). It maps to the recipe's
+    # sequence-length ceiling, so validation checks it.
+    _MAX_LENGTH_PARAMS = ("dataset_max_len",)
+
+    def validate_length_constraints(self):
+        """Enforce that selected lengths fit the recipe's sequence_length.
+
+        For RL recipes: max_prompt_length + max_response_length must not exceed
+        sequence_length. For SFT/DPO: the single-example length (dataset_max_len)
+        must not exceed it. Per-field min/max are already enforced on assignment;
+        this adds the cross-field sum check that a per-field max cannot express.
+
+        Framework-agnostic: gated only on the recipe's sequence_length metadata,
+        not on the model family. No-op if sequence_length is unknown or the
+        relevant params are absent.
+        """
+        if self._sequence_length is None:
+            return
+
+        prompt = getattr(self, "max_prompt_length", None)
+        response = getattr(self, "max_response_length", None)
+        if prompt is not None and response is not None:
+            total = prompt + response
+            if total > self._sequence_length:
+                raise ValueError(
+                    f"max_prompt_length ({prompt}) + max_response_length ({response}) "
+                    f"= {total} exceeds the recipe's supported sequence length "
+                    f"({self._sequence_length}). Lower max_prompt_length and/or "
+                    f"max_response_length so their sum is within {self._sequence_length}."
+                )
+
+        for param in self._MAX_LENGTH_PARAMS:
+            max_length = getattr(self, param, None)
+            if max_length is not None and max_length > self._sequence_length:
+                raise ValueError(
+                    f"{param} ({max_length}) exceeds the recipe's supported sequence "
+                    f"length ({self._sequence_length}). Set {param} to "
+                    f"{self._sequence_length} or lower."
+                )
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert back to dictionary for hyperparameters with string values."""

sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py

@@ -568,8 +568,28 @@ def _resolve_model_package_arn(model_package) -> Optional[str]:
         return None
 
 
+def _parse_sequence_length(value) -> int:
+    """Parse a sequence length value like '8K', '32K', '128K' into an integer (e.g., 8192)."""
+    if not value:
+        return 0
+    value = str(value).strip().upper()
+    if not value.endswith("K"):
+        raise ValueError(
+            f"Invalid sequence_length '{value}'. "
+            f"Expected a value ending in 'K', e.g. '8K' or '128K'."
+        )
+    try:
+        return int(value[:-1]) * 1024
+    except ValueError:
+        raise ValueError(
+            f"Invalid sequence_length '{value}'. "
+            f"Expected a numeric value followed by 'K', e.g. '8K' or '128K'."
+        )
+
+
 def _get_fine_tuning_options_and_model_arn(model_name: str, customization_technique: str, training_type, sagemaker_session,
-                                         hub_name: Optional[str] = None, compute: Optional[Union[HyperPodCompute, TrainingJobCompute]] = None) -> tuple:
+                                         sequence_length=None, hub_name: Optional[str] = None,
+                                         compute: Optional[Union[HyperPodCompute, TrainingJobCompute]] = None) -> tuple:
     """Get fine-tuning options and model ARN for given customization technique.
     Returns:
         tuple: (FineTuningOptions, model_arn, is_gated_model)
@@ -621,6 +641,29 @@ def _get_fine_tuning_options_and_model_arn(model_name: str, customization_techni
         if not recipes_with_template:
             raise ValueError(f"No recipes found with {platform_label} for technique: {customization_technique}")
 
+        # Filter by SequenceLength before recipe selection if sequence_length is requested.
+        # Multiple recipes may share the same SequenceLength (e.g. LORA and FULL
+        # variants); keep every exact match so _select_recipe_by_training_type can
+        # pick the right one for the requested training type.
+        if sequence_length:
+            requested = _parse_sequence_length(sequence_length)
+            candidates_with_sequence = [r for r in recipes_with_template if r.get("SequenceLength")]
+            if candidates_with_sequence:
+                filtered = [r for r in candidates_with_sequence if _parse_sequence_length(r.get("SequenceLength")) == requested]
+                if filtered:
+                    recipes_with_template = filtered
+                else:
+                    available = sorted(set(r.get("SequenceLength") for r in candidates_with_sequence))
+                    raise ValueError(
+                        f"No recipes found with SequenceLength == {sequence_length}. "
+                        f"Available sequence lengths: {available}"
+                    )
+            else:
+                raise ValueError(
+                    f"No recipes found with {platform_label} for technique: {customization_technique},training_type:{training_type}, "
+                    f"and sequence length:{sequence_length}"
+                )
+
         # Select recipe based on training type
         recipe = _select_recipe_by_training_type(recipes_with_template, training_type)
 
@@ -679,10 +722,16 @@ def _get_fine_tuning_options_and_model_arn(model_name: str, customization_techni
             except Exception as e:
                 logger.debug(f"Could not fetch subscription recipe override_params: {type(e).__name__}: {e}")
 
+        # Supported sequence-length ceiling: the recipe's SequenceLength ("<n>K")
+        # is the single source of truth. Parse it to an int so we can validate that
+        # max_prompt_length + max_response_length (or max_length) stays within
+        # what the recipe's hardware can serve. 0 if absent/unparseable (no-op).
+        sequence_length_ceiling = _parse_sequence_length(recipe.get("SequenceLength")) or None
+
         if options_dict:
-            return FineTuningOptions(options_dict), model_arn, is_gated_model
+            return FineTuningOptions(options_dict, sequence_length=sequence_length_ceiling), model_arn, is_gated_model
         else:
-            return FineTuningOptions({}), model_arn, is_gated_model
+            return FineTuningOptions({}, sequence_length=sequence_length_ceiling), model_arn, is_gated_model
 
     except Exception as e:
         logger.debug("Exception getting fine-tuning options: %s", e)
@@ -881,7 +930,8 @@ def _resolve_model_and_name(model, sagemaker_session=None):
 
 
 def _create_serverless_config(model_arn, customization_technique,
-                           training_type, accept_eula, evaluator_arn=None, job_type=JOB_TYPE) -> Optional['ServerlessJobConfig']:
+                           training_type, accept_eula, evaluator_arn=None,
+                           sequence_length=None, job_type=JOB_TYPE) -> Optional['ServerlessJobConfig']:
     """Create serverless job configuration for fine-tuning.
 
     Args:
@@ -890,6 +940,7 @@ def _create_serverless_config(model_arn, customization_technique,
         training_type: Training type (TrainingType enum or string)
         accept_eula: Boolean indicating if EULA is accepted
         evaluator_arn: Optional evaluator ARN for RLVR/RLAIF
+        sequence_length: Optional sequence length enum value (e.g., "1K", "2K", "4K", "8K", "16K", "32K", "64K", "128K")
         job_type: Type of job (default: "FineTuning")
 
     Returns:
@@ -905,7 +956,8 @@ def _create_serverless_config(model_arn, customization_technique,
         customization_technique=customization_technique,
         peft=peft,
         evaluator_arn=evaluator_arn,
-        accept_eula=accept_eula
+        accept_eula=accept_eula,
+        sequence_length=sequence_length,
     )
 
     return serverless_config

sagemaker-train/src/sagemaker/train/dpo_trainer.py

@@ -105,6 +105,10 @@ class DPOTrainer(BaseTrainer):
         stopping_condition (Optional[StoppingCondition]):
             The stopping condition to override training runtime limit.
             If not specified, uses SageMaker service default (24 hours for serverless training).
+        sequence_length (Optional[str]):
+            The sequence length for the training job. Valid values are
+            "1K", "2K", "4K", "8K", "16K", "32K", "64K", "128K".
+            If not specified, the service will use default recipe selection behavior.
         is_multimodal (Optional[bool]):
             Whether the training dataset contains multimodal data. If None (default),
             auto-detected from the training dataset at train time.
@@ -128,6 +132,7 @@ def __init__(
             networking: Optional[VpcConfig] = None,
             accept_eula: bool = False,
             stopping_condition: Optional[StoppingCondition] = None,
+            sequence_length: Optional[str] = None,
             recipe: Optional[str] = None,
             overrides: Optional[dict] = None,
             is_multimodal: Optional[bool] = None,
@@ -164,6 +169,7 @@ def __init__(
         self.kms_key_id = kms_key_id
         self.networking = networking
         self.stopping_condition = stopping_condition
+        self.sequence_length = sequence_length
         self._recipe_path = recipe
         self._overrides = overrides
         self._recipe_resolver = None
@@ -176,8 +182,8 @@ def __init__(
                                                                                       self.training_type,
                                                                                       self.sagemaker_session or TrainDefaults.get_sagemaker_session(
                                                                                       sagemaker_session=self.sagemaker_session
-
                                                                                     ),
+                                                                                      sequence_length=self.sequence_length,
                                                                                       compute=self.compute)
 
         # Process hyperparameters
@@ -292,12 +298,14 @@ def train(self,
             disable_output_compression=getattr(self, 'disable_output_compression', False),
         )
 
-        serverless_config = _create_serverless_config(model_arn=self._model_arn,
-                                                     customization_technique=CustomizationTechnique.DPO.value,
-                                                     training_type=self.training_type,
-                                                     accept_eula=self.accept_eula,
-                                                     job_type=JOB_TYPE
-                                                     )
+        serverless_config = _create_serverless_config(
+            model_arn=self._model_arn,
+            customization_technique=CustomizationTechnique.DPO.value,
+            training_type=self.training_type,
+            accept_eula=self.accept_eula,
+            sequence_length=self.sequence_length,
+            job_type=JOB_TYPE
+        )
 
         mlflow_config = _create_mlflow_config(
             sagemaker_session,

sagemaker-train/src/sagemaker/train/rlaif_trainer.py

@@ -117,6 +117,10 @@ class RLAIFTrainer(BaseTrainer):
         stopping_condition (Optional[StoppingCondition]):
             The stopping condition to override training runtime limit.
             If not specified, uses SageMaker service default (24 hours for serverless training).
+        sequence_length (Optional[str]):
+            The sequence length for the training job. Valid values are
+            "1K", "2K", "4K", "8K", "16K", "32K", "64K", "128K".
+            If not specified, the service will use default recipe selection behavior.
         is_multimodal (Optional[bool]):
             Whether the training dataset contains multimodal data. If None (default),
             auto-detected from the training dataset at train time.
@@ -141,6 +145,7 @@ def __init__(
         networking: Optional[VpcConfig] = None,
         accept_eula: bool = False,
         stopping_condition: Optional[StoppingCondition] = None,
+        sequence_length: Optional[str] = None,
         is_multimodal: Optional[bool] = None,
         **kwargs,
     ):
@@ -163,15 +168,17 @@ def __init__(
         self.kms_key_id = kms_key_id
         self.networking = networking
         self.stopping_condition = stopping_condition
+        self.sequence_length = sequence_length
         self.is_multimodal = is_multimodal
 
         # Initialize fine-tuning options with beta session fallback
-        self.hyperparameters, self._model_arn, is_gated_model = _get_fine_tuning_options_and_model_arn(self._model_name,
-                                                                     CustomizationTechnique.RLAIF.value,
-                                                                     self.training_type,
-                                                                     self.sagemaker_session or TrainDefaults.get_sagemaker_session(
-                                                                     sagemaker_session=self.sagemaker_session
-                                                                    ))
+        self.hyperparameters, self._model_arn, is_gated_model = _get_fine_tuning_options_and_model_arn(
+            self._model_name,
+            CustomizationTechnique.RLAIF.value,
+            self.training_type,
+            self.sagemaker_session or TrainDefaults.get_sagemaker_session(sagemaker_session=self.sagemaker_session),
+            sequence_length=self.sequence_length
+        )
 
         # Validate and set EULA acceptance
         self.accept_eula = _validate_eula_for_gated_model(model, accept_eula, is_gated_model)
@@ -256,13 +263,15 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
         )
 
         evaluator_arn = getattr(self, '_evaluator_arn', None)
-        serverless_config = _create_serverless_config(model_arn=self._model_arn,
-                                                     customization_technique=CustomizationTechnique.RLAIF.value,
-                                                     training_type=self.training_type,
-                                                     accept_eula=self.accept_eula,
-                                                     evaluator_arn=evaluator_arn,
-                                                     job_type=JOB_TYPE
-                                                     )
+        serverless_config = _create_serverless_config(
+            model_arn=self._model_arn,
+            customization_technique=CustomizationTechnique.RLAIF.value,
+            training_type=self.training_type,
+            accept_eula=self.accept_eula,
+            evaluator_arn=evaluator_arn,
+            sequence_length=self.sequence_length,
+            job_type=JOB_TYPE
+        )
 
         mlflow_config = _create_mlflow_config(
             sagemaker_session,
@@ -271,6 +280,9 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
             mlflow_run_name=self.mlflow_run_name,
         )
 
+        # Enforce prompt + response fit the recipe's supported sequence length.
+        self.hyperparameters.validate_length_constraints()
+
         final_hyperparameters = self.hyperparameters.to_dict()
 
         # Resolve is_multimodal: auto-detect from training dataset if not explicitly set