Skip to content

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

Closed
guanweim wants to merge 0 commit into
aws:masterfrom
guanweim:context_length
Closed

feat(train): Add SequenceLength support for SFT, DPO, RLVR, RLAIF trainers#5898
guanweim wants to merge 0 commit into
aws:masterfrom
guanweim:context_length

Conversation

@guanweim

Copy link
Copy Markdown

Description of changes:

Summary

  • Add optional sequence_length parameter to SFTTrainer, DPOTrainer, RLVRTrainer, and RLAIFTrainer for specifying context length in serverless training jobs
  • Recipes are filtered by SequenceLength field, picking the smallest recipe with context length >= the requested value
  • Add SequenceLength to service-2.json and regenerate shapes via codegen

Test plan

  • Unit tests for _parse_context_length helper (K suffix, lowercase, integer, None, empty)
  • Unit tests for recipe filtering by sequence length
  • Unit tests for ValueError when no sufficient context length exists
  • Unit tests for each trainer passing sequence_length through
  • Unit tests for _create_serverless_config with/without sequence_length
    By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

elif (isinstance(training_type, TrainingType) and training_type == TrainingType.FULL) or training_type == "FULL":
recipe = next((r for r in recipes_with_template if not r.get("Peft") and not r.get("IsSubscriptionModel")), None)

# Override recipe selection when sequence_length is explicitly requested

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is cleaner to think about as a recipe filter. So move the sequence_length filter above the recipe selection logic and reduce the list of available recipes. Afterwards the existing logic can stay the same

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated

accept_eula=accept_eula,
)
if sequence_length is not None:
config_kwargs["sequence_length"] = sequence_length

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can probably pass on the None value to ServerlessJobConfig

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated

@mollyheamazon mollyheamazon left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add integ test as sanity check as this is working e2e.

training_type,
sagemaker_session,
sequence_length=None,
hub_name: str = "SageMakerPublicHub"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be set to hub_name: Optional[str] = None for original pattern and consistency.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated

@guanweim

Copy link
Copy Markdown
Author

Please add integ test as sanity check as this is working e2e.

added

@guanweim guanweim requested review from mollyheamazon and zachgk May 29, 2026 00:41
@guanweim guanweim temporarily deployed to manual-approval June 16, 2026 23:44 — with GitHub Actions Inactive
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants