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

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

Add optional sequence_length parameter to all four trainers that enables
customers to specify their desired context length for serverless training
jobs. The parameter is passed in ServerlessJobConfig for recipe filtering.

During trainer initialization, _get_fine_tuning_options_and_model_arn
filters recipes by SequenceLength field, picking the smallest recipe
with context length >= the requested value. Raises ValueError if no
sufficient recipe exists or if recipes lack SequenceLength metadata.

Changes:
- ServerlessJobConfig: add sequence_length field
- _parse_context_length: parse values like '8K' to integers
- _get_fine_tuning_options_and_model_arn: filter by SequenceLength
- _create_serverless_config: conditionally include sequence_length
- SFTTrainer, DPOTrainer, RLVRTrainer, RLAIFTrainer: accept and
  thread sequence_length through init and train methods
- Unit tests for all new functionality

* fix: use codegen for SequenceLength shape instead of manual shapes.py edit

Add SequenceLength to service-2.json and regenerate shapes.py via codegen
(python -m sagemaker.core.tools.codegen) instead of editing shapes.py manually.

* fix: correct test imports and mock setup for sequence_length tests

* address PR review: sequence_length as recipe pre-filter and simplify config

- Move sequence_length filtering above recipe selection to reduce
  recipes_with_template before existing logic runs
- Always pass sequence_length to ServerlessJobConfig (no None guard)

* test: add integration test for SFT trainer with sequence_length

* Regenerate code via codegen and remove hardcoded path in data_extractor.py

* comment out gpu testing flag

* fix: reject bare integer inputs in _parse_context_length

Bare integers like "128" were silently parsed as token counts (128),
causing every recipe to pass the >= filter and silently selecting the
shortest recipe instead of failing. Now raises ValueError with an
actionable message for any input not ending in 'K'.

* Validate customer sequence lengths against the recipe's SequenceLength

Enforce that customer-selected lengths fit the recipe's supported sequence length:
- RL: max_prompt_length + max_response_length <= SequenceLength
- SFT/DPO: max_length (dataset_max_len) <= SequenceLength

FineTuningOptions.validate_length_constraints() checks this and is wired into the
RL/base trainers. The ceiling comes from the recipe's SequenceLength metadata
('<n>K'), parsed to an int via _parse_context_length; no-op when absent.

(Recipes side unified on the SequenceLength metadata field; the former
ContextLength field was dropped, so read SequenceLength here.)

* Rename context_length to sequence_length; exact-match recipe filter

Recipes vend only SequenceLength, so use that name throughout the fine-tuning
options and validation path (context_length -> sequence_length,
_parse_context_length -> _parse_sequence_length).

Filter recipes by exact SequenceLength match and keep all candidates at that
length, so _select_recipe_by_training_type picks the right recipe when LORA and
FULL variants share the same sequence length. The client sends only descriptors
in ServerlessJobConfig (including SequenceLength); the selectable-length cap is
derived from the recipe catalog rather than a hardcoded enum.

* Enforce SFT/DPO sequence-length ceiling on dataset_max_len

validate_length_constraints() checked only the RL prompt+response sum, so the
single-example ceiling for SFT/DPO recipes was never enforced. verl and llmft
recipes vend that ceiling as dataset_max_len; check it against the recipe's
SequenceLength so an over-limit value is rejected client-side.

Verified against the recipe catalog: dataset_max_len is the param verl/llmft
SFT/DPO recipes actually render (max_context_length appears in Nova overrides
but is never rendered into a recipe body; max_length is Nova-RFT only).

* test: point sequence_length integ tests at qwen3-5-9b verl recipes

The SFT sequence_length test targeted llama-3-2-1b-instruct, which only
offers a 4K recipe, so requesting 8K failed recipe selection. Point it at
huggingface-vlm-qwen3-5-9b and use 16K, which has a matching verl SFT recipe.

Add a matching RLVR sequence_length test on the same model at 8K.

---------

Co-authored-by: Mohamed Zeidan <81834882+mohamedzeidan2021@users.noreply.github.com>

Author: guanweim

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