diff --git a/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json b/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json
index ceb4f316dc..6b551c5fd6 100644
--- a/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json
+++ b/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json
@@ -44132,6 +44132,10 @@
"EvaluatorArn":{
"shape":"EvaluatorArn",
"documentation":" The evaluator Amazon Resource Name (ARN) used as reward function or reward prompt.
"
+ },
+ "SequenceLength":{
+ "shape":"SequenceLength",
+ "documentation":" The sequence length for the training job.
"
}
},
"documentation":" The configuration for the serverless training job.
"
@@ -44143,6 +44147,19 @@
"Evaluation"
]
},
+ "SequenceLength":{
+ "type":"string",
+ "enum":[
+ "1K",
+ "2K",
+ "4K",
+ "8K",
+ "16K",
+ "32K",
+ "64K",
+ "128K"
+ ]
+ },
"ServerlessMaxConcurrency":{
"type":"integer",
"box":true,
diff --git a/sagemaker-core/src/sagemaker/core/resources.py b/sagemaker-core/src/sagemaker/core/resources.py
index ff462080bb..322ff2f356 100644
--- a/sagemaker-core/src/sagemaker/core/resources.py
+++ b/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__)
diff --git a/sagemaker-core/src/sagemaker/core/shapes/shapes.py b/sagemaker-core/src/sagemaker/core/shapes/shapes.py
index 3eb547ec1c..e0e79a5751 100644
--- a/sagemaker-core/src/sagemaker/core/shapes/shapes.py
+++ b/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):
diff --git a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py
index 61d7a1b223..439fef722d 100644
--- a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py
+++ b/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",
},
diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py
index c17199a611..19a266ccfb 100644
--- a/sagemaker-train/src/sagemaker/train/base_trainer.py
+++ b/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)
diff --git a/sagemaker-train/src/sagemaker/train/common.py b/sagemaker-train/src/sagemaker/train/common.py
index cd6f2c5f6e..be0c301e51 100644
--- a/sagemaker-train/src/sagemaker/train/common.py
+++ b/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."""
diff --git a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py
index dc40bb981e..980d5dd27d 100644
--- a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py
+++ b/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 ("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
diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py
index a83d1ed2fe..f187bde8cf 100644
--- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py
+++ b/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,
diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py
index 5077c288d5..95020c1518 100644
--- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py
+++ b/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
diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py
index f20b54e7b1..8d77aca4e9 100644
--- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py
@@ -132,6 +132,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.
@@ -159,6 +163,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,
@@ -197,6 +202,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
@@ -211,6 +217,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
@@ -460,6 +467,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(
@@ -469,6 +477,9 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
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()
# Apply recipe/overrides if provided (overrides > recipe > Hub defaults)
diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py
index 3f14c348f0..f182835f16 100644
--- a/sagemaker-train/src/sagemaker/train/sft_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/sft_trainer.py
@@ -116,6 +116,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.
@@ -157,6 +161,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,
@@ -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,8 +215,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()
@@ -346,12 +353,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,
diff --git a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py
index 68446991c4..61d9ff3122 100644
--- a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py
+++ b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py
@@ -135,3 +135,39 @@ def test_sft_trainer_nova_workflow(sagemaker_session_us_east_1):
assert training_job.training_job_status == "Completed"
assert hasattr(training_job, 'output_model_package_arn')
assert training_job.output_model_package_arn is not None
+
+
+# @pytest.mark.gpu_intensive
+def test_sft_trainer_lora_with_sequence_length(sagemaker_session):
+ """Test SFT training workflow with LORA and sequence_length specified."""
+ unique_id = f"{int(time.time())}-{random.randint(1000, 9999)}"
+
+ sft_trainer = SFTTrainer(
+ model="meta-textgeneration-llama-3-2-1b-instruct",
+ training_type=TrainingType.LORA,
+ model_package_group="arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models",
+ training_dataset="s3://mc-flows-sdk-testing/input_data/sft/sample_data_256_final.jsonl",
+ s3_output_path="s3://mc-flows-sdk-testing/output/",
+ accept_eula=True,
+ sequence_length="8K",
+ base_job_name=f"sft-seqlen-integ-{unique_id}",
+ )
+
+ training_job = sft_trainer.train(wait=False)
+
+ max_wait_time = 3600
+ poll_interval = 30
+ start_time = time.time()
+
+ while time.time() - start_time < max_wait_time:
+ training_job.refresh()
+ status = training_job.training_job_status
+
+ if status in ["Completed", "Failed", "Stopped"]:
+ break
+
+ time.sleep(poll_interval)
+
+ assert training_job.training_job_status == "Completed"
+ assert hasattr(training_job, 'output_model_package_arn')
+ assert training_job.output_model_package_arn is not None
diff --git a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py
index 1c5e495f63..02e5e1d22f 100644
--- a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py
+++ b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py
@@ -31,9 +31,11 @@
_create_mlflow_config,
_validate_eula_for_gated_model,
_validate_model_region_availability,
- _validate_s3_path_exists
+ _validate_s3_path_exists,
+ _parse_sequence_length
)
from sagemaker.core.resources import ModelPackage, ModelPackageGroup
+from sagemaker.core.utils.utils import Unassigned
from sagemaker.ai_registry.dataset import DataSet
from sagemaker.train.common import TrainingType
from sagemaker.train.configs import InputData
@@ -632,7 +634,6 @@ def test__convert_input_data_to_channels(self):
def test__validate_eula_for_gated_model_with_model_package(self):
"""Test EULA validation returns True for ModelPackage input"""
- from sagemaker.core.resources import ModelPackage
model_package = Mock(spec=ModelPackage)
result = _validate_eula_for_gated_model(model_package, False, True)
@@ -1033,6 +1034,144 @@ def test__get_fine_tuning_options_subscription_enabled_but_not_subscribed(self,
assert "max_steps" in options._specs
assert "customer_data_percent" not in options._specs
+ def test__create_serverless_config_with_sequence_length(self):
+ config = _create_serverless_config("model-arn", "SFT", TrainingType.LORA, accept_eula=True, sequence_length="8K")
+
+ assert config.sequence_length == "8K"
+ assert config.base_model_arn == "model-arn"
+
+ def test__create_serverless_config_without_sequence_length(self):
+ config = _create_serverless_config("model-arn", "SFT", TrainingType.LORA, accept_eula=True)
+
+ assert config.sequence_length is None
+
+ def test__parse_sequence_length_with_k_suffix(self):
+ assert _parse_sequence_length("8K") == 8192
+ assert _parse_sequence_length("32K") == 32768
+ assert _parse_sequence_length("128K") == 131072
+
+ def test__parse_sequence_length_with_lowercase(self):
+ assert _parse_sequence_length("8k") == 8192
+
+ def test__parse_sequence_length_with_integer(self):
+ with pytest.raises(ValueError, match="Invalid sequence_length '4096'"):
+ _parse_sequence_length("4096")
+
+ def test__parse_sequence_length_with_none(self):
+ assert _parse_sequence_length(None) == 0
+
+ def test__parse_sequence_length_with_empty(self):
+ assert _parse_sequence_length("") == 0
+
+ @patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
+ def test__get_fine_tuning_options_filters_by_exact_sequence_length(self, mock_get_hub_content):
+ mock_session = Mock()
+ mock_session.boto_session.region_name = "us-east-1"
+ mock_s3 = Mock()
+ mock_s3.get_object.return_value = {
+ "Body": Mock(read=Mock(return_value=b'{"max_length": {"default": 32768}}'))
+ }
+ mock_session.boto_session.client.return_value = mock_s3
+
+ mock_get_hub_content.return_value = {
+ 'hub_content_arn': "arn:aws:sagemaker:us-east-1:123456789012:model/test-model",
+ 'hub_content_document': {
+ "GatedBucket": False,
+ "RecipeCollection": [
+ {
+ "CustomizationTechnique": "SFT",
+ "SmtjRecipeTemplateS3Uri": "s3://bucket/template-4k.json",
+ "SmtjOverrideParamsS3Uri": "s3://bucket/params-4k.json",
+ "Peft": True,
+ "SequenceLength": "4K"
+ },
+ {
+ "CustomizationTechnique": "SFT",
+ "SmtjRecipeTemplateS3Uri": "s3://bucket/template-32k.json",
+ "SmtjOverrideParamsS3Uri": "s3://bucket/params-32k.json",
+ "Peft": True,
+ "SequenceLength": "32K"
+ }
+ ]
+ }
+ }
+
+ result = _get_fine_tuning_options_and_model_arn("test-model", "SFT", "LORA", mock_session, sequence_length="32K")
+
+ assert result is not None
+ options, model_arn, is_gated_model = result
+ # Should pick the recipe whose SequenceLength exactly matches the request.
+ mock_s3.get_object.assert_called_once()
+ call_args = mock_s3.get_object.call_args[1]
+ assert "params-32k" in call_args["Key"]
+
+ @patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
+ def test__get_fine_tuning_options_keeps_all_recipes_at_same_sequence_length(self, mock_get_hub_content):
+ # Multiple recipes share the same SequenceLength (LORA + FULL). Selection
+ # by training_type must resolve to the LORA one, not an arbitrary match.
+ mock_session = Mock()
+ mock_session.boto_session.region_name = "us-east-1"
+ mock_s3 = Mock()
+ mock_s3.get_object.return_value = {
+ "Body": Mock(read=Mock(return_value=b'{"max_length": {"default": 32768}}'))
+ }
+ mock_session.boto_session.client.return_value = mock_s3
+
+ mock_get_hub_content.return_value = {
+ 'hub_content_arn': "arn:aws:sagemaker:us-east-1:123456789012:model/test-model",
+ 'hub_content_document': {
+ "GatedBucket": False,
+ "RecipeCollection": [
+ {
+ "CustomizationTechnique": "SFT",
+ "SmtjRecipeTemplateS3Uri": "s3://bucket/template-32k-full.json",
+ "SmtjOverrideParamsS3Uri": "s3://bucket/params-32k-full.json",
+ "Peft": False,
+ "SequenceLength": "32K"
+ },
+ {
+ "CustomizationTechnique": "SFT",
+ "SmtjRecipeTemplateS3Uri": "s3://bucket/template-32k-lora.json",
+ "SmtjOverrideParamsS3Uri": "s3://bucket/params-32k-lora.json",
+ "Peft": True,
+ "SequenceLength": "32K"
+ }
+ ]
+ }
+ }
+
+ result = _get_fine_tuning_options_and_model_arn("test-model", "SFT", "LORA", mock_session, sequence_length="32K")
+
+ assert result is not None
+ mock_s3.get_object.assert_called_once()
+ call_args = mock_s3.get_object.call_args[1]
+ assert "params-32k-lora" in call_args["Key"]
+
+ @patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
+ def test__get_fine_tuning_options_raises_when_no_exact_sequence_length(self, mock_get_hub_content):
+ mock_session = Mock()
+ mock_session.boto_session.region_name = "us-east-1"
+
+ mock_get_hub_content.return_value = {
+ 'hub_content_arn': "arn:aws:sagemaker:us-east-1:123456789012:model/test-model",
+ 'hub_content_document': {
+ "GatedBucket": False,
+ "RecipeCollection": [
+ {
+ "CustomizationTechnique": "SFT",
+ "SmtjRecipeTemplateS3Uri": "s3://bucket/template-4k.json",
+ "SmtjOverrideParamsS3Uri": "s3://bucket/params-4k.json",
+ "Peft": True,
+ "SequenceLength": "4K"
+ }
+ ]
+ }
+ }
+
+ # Requesting 128K but only 4K available — no exact match, should raise.
+ with pytest.raises(ValueError, match="No recipes found with SequenceLength == 128K"):
+ _get_fine_tuning_options_and_model_arn("test-model", "SFT", "LORA", mock_session, sequence_length="128K")
+
# ===========================================================================
# Hub recipe/image resolution helpers
diff --git a/sagemaker-train/tests/unit/train/test_common.py b/sagemaker-train/tests/unit/train/test_common.py
index 25d1decbc1..74230cad60 100644
--- a/sagemaker-train/tests/unit/train/test_common.py
+++ b/sagemaker-train/tests/unit/train/test_common.py
@@ -54,3 +54,78 @@ def test_to_dict_all_none_returns_empty(self):
})
result = options.to_dict()
assert result == {}
+
+
+import pytest
+
+
+class TestValidateLengthConstraints:
+ """FineTuningOptions.validate_length_constraints() — sum vs sequence_length."""
+
+ def _rl_options(self, prompt, response, sequence_length):
+ # Per-field max defaults to sequence_length; when unknown, omit it so the
+ # per-field range check is a no-op and only the sum-guard is exercised.
+ prompt_spec = {"type": "integer", "default": 1024, "min": 512}
+ response_spec = {"type": "integer", "default": 2048, "min": 100}
+ if sequence_length is not None:
+ prompt_spec["max"] = sequence_length
+ response_spec["max"] = sequence_length
+ opts = FineTuningOptions(
+ {"max_prompt_length": prompt_spec, "max_response_length": response_spec},
+ sequence_length=sequence_length,
+ )
+ object.__setattr__(opts, "max_prompt_length", prompt)
+ object.__setattr__(opts, "max_response_length", response)
+ return opts
+
+ def test_rl_sum_within_sequence_length_passes(self):
+ opts = self._rl_options(prompt=4096, response=8192, sequence_length=262144)
+ opts.validate_length_constraints() # no raise
+
+ def test_rl_sum_equal_to_sequence_length_passes(self):
+ opts = self._rl_options(prompt=100000, response=162144, sequence_length=262144)
+ opts.validate_length_constraints() # sum == ceiling, allowed
+
+ def test_rl_sum_exceeds_sequence_length_raises(self):
+ opts = self._rl_options(prompt=200000, response=200000, sequence_length=262144)
+ with pytest.raises(ValueError) as exc:
+ opts.validate_length_constraints()
+ msg = str(exc.value)
+ assert "400000" in msg and "262144" in msg
+
+ def test_no_sequence_length_is_noop(self):
+ opts = self._rl_options(prompt=200000, response=200000, sequence_length=None)
+ opts.validate_length_constraints() # unknown ceiling -> no raise
+
+ # verl/llmft SFT/DPO recipes vend the single-example ceiling as
+ # "dataset_max_len". Validation must enforce it so the ceiling is not
+ # silently unenforced for these frameworks.
+ def test_sft_dataset_max_len_within_sequence_length_passes(self):
+ opts = FineTuningOptions(
+ {"dataset_max_len": {"type": "integer", "default": 4096, "min": 4096, "max": 131072}},
+ sequence_length=131072,
+ )
+ opts.dataset_max_len = 65536
+ opts.validate_length_constraints() # no raise
+
+ def test_sft_dataset_max_len_exceeds_sequence_length_raises(self):
+ opts = FineTuningOptions(
+ {"dataset_max_len": {"type": "integer", "default": 4096, "min": 4096, "max": 131072}},
+ sequence_length=131072,
+ )
+ # Bypass the per-field max so the sequence-length guard is what raises.
+ object.__setattr__(opts, "dataset_max_len", 200000)
+ with pytest.raises(ValueError) as exc:
+ opts.validate_length_constraints()
+ msg = str(exc.value)
+ assert "dataset_max_len" in msg and "131072" in msg
+
+ def test_framework_agnostic_gated_on_sequence_length_metadata_only(self):
+ # No sequence_length metadata (e.g. a recipe that does not vend it) -> the
+ # ceiling is unknown and validation is a no-op regardless of param name.
+ opts = FineTuningOptions(
+ {"dataset_max_len": {"type": "integer", "default": 4096, "min": 4096}},
+ sequence_length=None,
+ )
+ object.__setattr__(opts, "dataset_max_len", 999999)
+ opts.validate_length_constraints() # unknown ceiling -> no raise
diff --git a/sagemaker-train/tests/unit/train/test_dpo_trainer.py b/sagemaker-train/tests/unit/train/test_dpo_trainer.py
index e960451196..95a13a2c1c 100644
--- a/sagemaker-train/tests/unit/train/test_dpo_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_dpo_trainer.py
@@ -509,6 +509,70 @@ def test_train_wait_false_skips_wait(self, mock_training_job_create, mock_model_
mock_wait.assert_not_called()
+ @patch('sagemaker.train.dpo_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.dpo_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_sequence_length_default_none(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = DPOTrainer(model="test-model", model_package_group="test-group")
+ assert trainer.sequence_length is None
+
+ @patch('sagemaker.train.dpo_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.dpo_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_with_sequence_length(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = DPOTrainer(model="test-model", model_package_group="test-group", sequence_length="8K")
+ assert trainer.sequence_length == "8K"
+
+ @patch('sagemaker.train.dpo_trainer._resolve_model_and_name')
+ @patch('sagemaker.train.dpo_trainer._get_fine_tuning_options_and_model_arn')
+ @patch('sagemaker.train.dpo_trainer.TrainDefaults.get_role')
+ @patch('sagemaker.train.dpo_trainer.TrainDefaults.get_sagemaker_session')
+ @patch('sagemaker.train.dpo_trainer._get_unique_name')
+ @patch('sagemaker.train.dpo_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.dpo_trainer._create_input_data_config')
+ @patch('sagemaker.train.dpo_trainer._convert_input_data_to_channels')
+ @patch('sagemaker.train.dpo_trainer._create_output_config')
+ @patch('sagemaker.train.dpo_trainer._create_serverless_config')
+ @patch('sagemaker.train.dpo_trainer._create_mlflow_config')
+ @patch('sagemaker.train.dpo_trainer._create_model_package_config')
+ @patch('sagemaker.core.resources.TrainingJob.create')
+ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_job_create,
+ mock_model_package_config, mock_mlflow_config, mock_serverless_config,
+ mock_output_config, mock_convert_channels, mock_input_config,
+ mock_validate_group, mock_unique_name, mock_get_sagemaker_session,
+ mock_get_role, mock_get_options, mock_resolve_model):
+ mock_validate_group.return_value = "test-group"
+ mock_resolve_model.return_value = ("test-model", "test-model")
+ mock_get_sagemaker_session.return_value = Mock()
+ mock_fine_tuning_options = Mock()
+ mock_fine_tuning_options.to_dict.return_value = {}
+ mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)
+ mock_get_role.return_value = "test-role"
+ mock_unique_name.return_value = "test-job-name"
+ mock_input_config.return_value = [Mock()]
+ mock_convert_channels.return_value = [Mock()]
+ mock_output_config.return_value = Mock()
+ mock_serverless_config.return_value = Mock()
+ mock_mlflow_config.return_value = Mock()
+ mock_model_package_config.return_value = Mock()
+ mock_training_job = Mock()
+ mock_training_job_create.return_value = mock_training_job
+
+ trainer = DPOTrainer(model="test-model", model_package_group="test-group",
+ training_dataset="s3://bucket/train", sequence_length="16K")
+ trainer.train(wait=False)
+
+ mock_serverless_config.assert_called_once()
+ call_kwargs = mock_serverless_config.call_args[1]
+ assert call_kwargs["sequence_length"] == "16K"
+
+
class TestDPOTrainerComputeDispatch:
"""Tests for compute dispatch in DPOTrainer."""
diff --git a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py b/sagemaker-train/tests/unit/train/test_rlaif_trainer.py
index a4f16fee42..34edeac298 100644
--- a/sagemaker-train/tests/unit/train/test_rlaif_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_rlaif_trainer.py
@@ -682,4 +682,70 @@ def test_train_wait_false_skips_wait(self, mock_training_job_create, mock_model_
trainer = RLAIFTrainer(model="test-model", model_package_group="test-group", training_dataset="s3://bucket/train")
trainer.train(wait=False, wait_timeout=600)
- mock_wait.assert_not_called()
\ No newline at end of file
+ mock_wait.assert_not_called()
+
+ @patch('sagemaker.train.rlaif_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.rlaif_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_sequence_length_default_none(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_hyperparams._specs = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = RLAIFTrainer(model="test-model", model_package_group="test-group")
+ assert trainer.sequence_length is None
+
+ @patch('sagemaker.train.rlaif_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.rlaif_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_with_sequence_length(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_hyperparams._specs = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = RLAIFTrainer(model="test-model", model_package_group="test-group", sequence_length="128K")
+ assert trainer.sequence_length == "128K"
+
+ @patch('sagemaker.train.rlaif_trainer._resolve_model_and_name')
+ @patch('sagemaker.train.rlaif_trainer._get_fine_tuning_options_and_model_arn')
+ @patch('sagemaker.train.rlaif_trainer.TrainDefaults.get_role')
+ @patch('sagemaker.train.rlaif_trainer.TrainDefaults.get_sagemaker_session')
+ @patch('sagemaker.train.rlaif_trainer._get_unique_name')
+ @patch('sagemaker.train.rlaif_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.rlaif_trainer._create_input_data_config')
+ @patch('sagemaker.train.rlaif_trainer._convert_input_data_to_channels')
+ @patch('sagemaker.train.rlaif_trainer._create_output_config')
+ @patch('sagemaker.train.rlaif_trainer._create_serverless_config')
+ @patch('sagemaker.train.rlaif_trainer._create_mlflow_config')
+ @patch('sagemaker.train.rlaif_trainer._create_model_package_config')
+ @patch('sagemaker.core.resources.TrainingJob.create')
+ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_job_create,
+ mock_model_package_config, mock_mlflow_config, mock_serverless_config,
+ mock_output_config, mock_convert_channels, mock_input_config,
+ mock_validate_group, mock_unique_name, mock_get_sagemaker_session,
+ mock_get_role, mock_get_options, mock_resolve_model):
+ mock_validate_group.return_value = "test-group"
+ mock_resolve_model.return_value = ("test-model", "test-model")
+ mock_get_sagemaker_session.return_value = Mock()
+ mock_fine_tuning_options = Mock()
+ mock_fine_tuning_options.to_dict.return_value = {}
+ mock_fine_tuning_options._specs = {}
+ mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)
+ mock_get_role.return_value = "test-role"
+ mock_unique_name.return_value = "test-job-name"
+ mock_input_config.return_value = [Mock()]
+ mock_convert_channels.return_value = [Mock()]
+ mock_output_config.return_value = Mock()
+ mock_serverless_config.return_value = Mock()
+ mock_mlflow_config.return_value = Mock()
+ mock_model_package_config.return_value = Mock()
+ mock_training_job = Mock()
+ mock_training_job_create.return_value = mock_training_job
+
+ trainer = RLAIFTrainer(model="test-model", model_package_group="test-group",
+ training_dataset="s3://bucket/train", sequence_length="64K")
+ trainer.train(wait=False)
+
+ mock_serverless_config.assert_called_once()
+ call_kwargs = mock_serverless_config.call_args[1]
+ assert call_kwargs["sequence_length"] == "64K"
diff --git a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py
index 5a6178392b..198b8a50ff 100644
--- a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py
@@ -512,6 +512,70 @@ def test_train_wait_false_skips_wait(self, mock_training_job_create, mock_model_
mock_wait.assert_not_called()
+ @patch('sagemaker.train.rlvr_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.rlvr_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_sequence_length_default_none(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = RLVRTrainer(model="test-model", model_package_group="test-group")
+ assert trainer.sequence_length is None
+
+ @patch('sagemaker.train.rlvr_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.rlvr_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_with_sequence_length(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = RLVRTrainer(model="test-model", model_package_group="test-group", sequence_length="32K")
+ assert trainer.sequence_length == "32K"
+
+ @patch('sagemaker.train.rlvr_trainer._resolve_model_and_name')
+ @patch('sagemaker.train.rlvr_trainer._get_fine_tuning_options_and_model_arn')
+ @patch('sagemaker.train.rlvr_trainer.TrainDefaults.get_role')
+ @patch('sagemaker.train.rlvr_trainer.TrainDefaults.get_sagemaker_session')
+ @patch('sagemaker.train.rlvr_trainer._get_unique_name')
+ @patch('sagemaker.train.rlvr_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.rlvr_trainer._create_input_data_config')
+ @patch('sagemaker.train.rlvr_trainer._convert_input_data_to_channels')
+ @patch('sagemaker.train.rlvr_trainer._create_output_config')
+ @patch('sagemaker.train.rlvr_trainer._create_serverless_config')
+ @patch('sagemaker.train.rlvr_trainer._create_mlflow_config')
+ @patch('sagemaker.train.rlvr_trainer._create_model_package_config')
+ @patch('sagemaker.core.resources.TrainingJob.create')
+ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_job_create,
+ mock_model_package_config, mock_mlflow_config, mock_serverless_config,
+ mock_output_config, mock_convert_channels, mock_input_config,
+ mock_validate_group, mock_unique_name, mock_get_sagemaker_session,
+ mock_get_role, mock_get_options, mock_resolve_model):
+ mock_validate_group.return_value = "test-group"
+ mock_resolve_model.return_value = ("test-model", "test-model")
+ mock_get_sagemaker_session.return_value = Mock()
+ mock_fine_tuning_options = Mock()
+ mock_fine_tuning_options.to_dict.return_value = {}
+ mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)
+ mock_get_role.return_value = "test-role"
+ mock_unique_name.return_value = "test-job-name"
+ mock_input_config.return_value = [Mock()]
+ mock_convert_channels.return_value = [Mock()]
+ mock_output_config.return_value = Mock()
+ mock_serverless_config.return_value = Mock()
+ mock_mlflow_config.return_value = Mock()
+ mock_model_package_config.return_value = Mock()
+ mock_training_job = Mock()
+ mock_training_job_create.return_value = mock_training_job
+
+ trainer = RLVRTrainer(model="test-model", model_package_group="test-group",
+ training_dataset="s3://bucket/train", sequence_length="4K")
+ trainer.train(wait=False)
+
+ mock_serverless_config.assert_called_once()
+ call_kwargs = mock_serverless_config.call_args[1]
+ assert call_kwargs["sequence_length"] == "4K"
+
+
class TestRLVRTrainerComputeDispatch:
"""Tests for compute dispatch in RLVRTrainer."""
diff --git a/sagemaker-train/tests/unit/train/test_sft_trainer.py b/sagemaker-train/tests/unit/train/test_sft_trainer.py
index 7586a8d7df..cd2010ceda 100644
--- a/sagemaker-train/tests/unit/train/test_sft_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_sft_trainer.py
@@ -523,6 +523,70 @@ def test_train_wait_false_skips_wait(self, mock_training_job_create, mock_model_
mock_wait.assert_not_called()
+ @patch('sagemaker.train.sft_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_sequence_length_default_none(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = SFTTrainer(model="test-model", model_package_group="test-group")
+ assert trainer.sequence_length is None
+
+ @patch('sagemaker.train.sft_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn')
+ def test_init_with_sequence_length(self, mock_finetuning_options, mock_validate_group):
+ mock_validate_group.return_value = "test-group"
+ mock_hyperparams = Mock()
+ mock_hyperparams.to_dict.return_value = {}
+ mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
+ trainer = SFTTrainer(model="test-model", model_package_group="test-group", sequence_length="8K")
+ assert trainer.sequence_length == "8K"
+
+ @patch('sagemaker.train.sft_trainer._resolve_model_and_name')
+ @patch('sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn')
+ @patch('sagemaker.train.sft_trainer.TrainDefaults.get_role')
+ @patch('sagemaker.train.sft_trainer.TrainDefaults.get_sagemaker_session')
+ @patch('sagemaker.train.sft_trainer._get_unique_name')
+ @patch('sagemaker.train.sft_trainer._validate_and_resolve_model_package_group')
+ @patch('sagemaker.train.sft_trainer._create_input_data_config')
+ @patch('sagemaker.train.sft_trainer._convert_input_data_to_channels')
+ @patch('sagemaker.train.sft_trainer._create_output_config')
+ @patch('sagemaker.train.sft_trainer._create_serverless_config')
+ @patch('sagemaker.train.sft_trainer._create_mlflow_config')
+ @patch('sagemaker.train.sft_trainer._create_model_package_config')
+ @patch('sagemaker.core.resources.TrainingJob.create')
+ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_job_create,
+ mock_model_package_config, mock_mlflow_config, mock_serverless_config,
+ mock_output_config, mock_convert_channels, mock_input_config,
+ mock_validate_group, mock_unique_name, mock_get_sagemaker_session,
+ mock_get_role, mock_get_options, mock_resolve_model):
+ mock_validate_group.return_value = "test-group"
+ mock_resolve_model.return_value = ("test-model", "test-model")
+ mock_get_sagemaker_session.return_value = Mock()
+ mock_fine_tuning_options = Mock()
+ mock_fine_tuning_options.to_dict.return_value = {}
+ mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)
+ mock_get_role.return_value = "test-role"
+ mock_unique_name.return_value = "test-job-name"
+ mock_input_config.return_value = [Mock()]
+ mock_convert_channels.return_value = [Mock()]
+ mock_output_config.return_value = Mock()
+ mock_serverless_config.return_value = Mock()
+ mock_mlflow_config.return_value = Mock()
+ mock_model_package_config.return_value = Mock()
+ mock_training_job = Mock()
+ mock_training_job_create.return_value = mock_training_job
+
+ trainer = SFTTrainer(model="test-model", model_package_group="test-group",
+ training_dataset="s3://bucket/train", sequence_length="16K")
+ trainer.train(wait=False)
+
+ mock_serverless_config.assert_called_once()
+ call_kwargs = mock_serverless_config.call_args[1]
+ assert call_kwargs["sequence_length"] == "16K"
+
+
class TestSFTTrainerComputeDispatch:
"""Tests for compute dispatch in SFTTrainer."""