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

@@ -9717,6 +9717,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
@@ -9726,6 +9727,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

@@ -16206,6 +16206,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/common_utils/finetune_utils.py

@@ -568,8 +568,28 @@ def _resolve_model_package_arn(model_package) -> Optional[str]:
         return None
 
 
+def _parse_context_length(value) -> int:
+    """Parse a context 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,27 @@ 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
+        if sequence_length:
+            requested = _parse_context_length(sequence_length)
+            candidates_with_context = [r for r in recipes_with_template if r.get("SequenceLength")]
+            if candidates_with_context:
+                filtered = [r for r in candidates_with_context if _parse_context_length(r.get("SequenceLength")) >= requested]
+                if filtered:
+                    filtered.sort(key=lambda r: _parse_context_length(r.get("SequenceLength")))
+                    recipes_with_template = filtered
+                else:
+                    available = sorted(set(r.get("SequenceLength") for r in candidates_with_context))
+                    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)
 
@@ -881,7 +922,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 +932,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 +948,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

@@ -104,6 +104,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.
@@ -127,6 +131,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,
@@ -163,6 +168,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
@@ -175,8 +181,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
@@ -285,12 +291,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

@@ -115,6 +115,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.
@@ -139,6 +143,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,
     ):
@@ -161,15 +166,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)
@@ -248,13 +255,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,

sagemaker-train/src/sagemaker/train/rlvr_trainer.py

@@ -130,6 +130,10 @@ class RLVRTrainer(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.
@@ -157,6 +161,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,
@@ -195,6 +200,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
@@ -209,6 +215,7 @@ def __init__(
                                                                      self.sagemaker_session or TrainDefaults.get_sagemaker_session(
                                                                      sagemaker_session=self.sagemaker_session
                                                                     ),
+                                                                     sequence_length=self.sequence_length,
                                                                      compute=self.compute)
 
         # Remove constructor-handled hyperparameters
@@ -450,6 +457,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
                                                      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-train/src/sagemaker/train/sft_trainer.py

@@ -115,6 +115,10 @@ class SFTTrainer(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.
         recipe (Optional[str]):
             Path to a user recipe YAML file (local path or S3 URI). When provided,
             enables 3-level recipe resolution: Hub defaults < recipe file < overrides dict.
@@ -156,6 +160,7 @@ def __init__(
         networking: Optional[VpcConfig] = None,
         accept_eula: Optional[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,
@@ -194,6 +199,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
@@ -208,8 +214,9 @@ def __init__(
                                                                      self.sagemaker_session or TrainDefaults.get_sagemaker_session(
                                                                      sagemaker_session=self.sagemaker_session
                                                                      ),
+                                                                     sequence_length=self.sequence_length,
                                                                      compute=self.compute)
-        
+
         # Process hyperparameters
         self._process_hyperparameters()
 
@@ -339,12 +346,14 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
             disable_output_compression=getattr(self, 'disable_output_compression', False),
         )
 
-        serverless_config = _create_serverless_config(model_arn=self._model_arn,
-                                                     customization_technique=CustomizationTechnique.SFT.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.SFT.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,
             mlflow_resource_arn=self.mlflow_resource_arn,