feat(train): Add inherited list_supported_models to BaseTrainer - #6187
Conversation
Expose a list_supported_models capability on all fine-tuning trainers (SFT, RLVR, RLAIF, DPO, CPT) mirroring MultiTurnRLTrainer. Rather than duplicating the method per trainer, add a single inherited classmethod on BaseTrainer that resolves cls._customization_technique and delegates to the existing recipe_utils._list_hub_models_by_recipe primitive. RLAIFTrainer previously used the technique inline; add its _customization_technique class attribute so the inherited method resolves. MultiTurnRLTrainer keeps its own override (CreateJob flow, non-enum technique). Add unit tests covering per-trainer technique resolution, the base delegation path, and the missing-technique guard. --- X-AI-Prompt: Add list_supported_models to SFT/RLVR/RLAIF/DPO trainers like MTRL; refactored to a shared BaseTrainer method X-AI-Tool: Kiro
Add an integration test that queries SageMakerPublicHub for each fine-tuning trainer (SFT, RLVR, RLAIF, DPO, CPT) and asserts the inherited list_supported_models returns a non-empty, sorted list of model names. This closes the contract gap unit tests cannot cover: that each trainer's _customization_technique string matches the live hub recipe keywords (@recipe:finetuning_{technique}_...). Mirrors the existing MTRL integ test. --- X-AI-Prompt: Add a parametrized integ test for list_supported_models across all fine-tuning trainers X-AI-Tool: Kiro
|
Claude finished @jam-jee's task in 1m 28s —— View job PR Review —
|
_list_hub_models_by_recipe built the search keyword as "@recipe:{type}_{technique}_" (trailing underscore) and matched by prefix, assuming every recipe keyword carries a "_{strategy}" component. Techniques with no strategy — CPT is tagged as the bare "@recipe:finetuning_cpt" — never matched, so CPTTrainer.list_supported_models() always returned an empty list despite the hub having CPT-tagged models. Match the bare base keyword OR the "{base}_{strategy}" form. The "{base}_" guard prevents a shorter technique from matching a longer one that shares its prefix (e.g. "rl" must not match "rlvr"). Rework the integ test into a hub-content-agnostic oracle: independently scan the active hub and assert list_supported_models returns exactly the tagged set per technique (0 or many), instead of a brittle non-empty assertion that broke under the private test hub which carries no CPT models. Add unit coverage for the bare-keyword match and the prefix-collision guard. Verified against SageMakerPublicHub: SFT=36, RLVR=33, RLAIF=31, DPO=32, CPT=5 models, each matching the independent oracle. --- X-AI-Prompt: Integ test surfaced CPT list returning empty; fixed suffix-less recipe keyword matching and made the integ test hub-agnostic X-AI-Tool: Kiro
|
Claude finished @jam-jee's task in 1m 38s —— View job PR Review —
|
|
Verified fortress reported issues are LOW severity. |
|
As this is public discoverable method, worth adding docs for the same. |
| Returns: | ||
| Sorted list of hub content model names supporting the technique. | ||
| """ | ||
| from sagemaker.train.common_utils.recipe_utils import _list_hub_models_by_recipe |
There was a problem hiding this comment.
Nit[non-blocking]: We can move the import to above
| @classmethod | ||
| @_telemetry_emitter( | ||
| feature=Feature.MODEL_CUSTOMIZATION, | ||
| func_name="BaseTrainer.list_supported_models", |
There was a problem hiding this comment.
non blocking but a callout: we might not be able to trace which Trainer is used to list supported models.
There was a problem hiding this comment.
good point. But I dont think we need granular list_supported_models metrics ATM.
2 participants
Problem
MultiTurnRLTrainerexposeslist_supported_models()so users can discover whichhub models have a matching recipe, but the other fine-tuning trainers
(
SFTTrainer,RLVRTrainer,RLAIFTrainer,DPOTrainer,CPTTrainer) have noequivalent — users have no programmatic way to find which models support a given
technique.
Why it matters
Without this, users must consult external docs or trial-and-error a model name into
train()to find out whether a technique is supported for it. A discovery methodper trainer makes the supported-model set self-describing from the SDK.
Fix (symptom → root cause → change)
The listing logic is already generic:
MultiTurnRLTrainer.list_supported_modelsjust delegates to
recipe_utils._list_hub_models_by_recipe(recipe_type, technique, session), and the hub keyword filter is keyed off the trainer's_customization_techniquevalue (@recipe:finetuning_{technique}_...). Rather thancopy the method into each trainer, add a single inherited
classmethodonBaseTrainerthat resolvescls._customization_techniqueand delegates to thatsame primitive. Every fine-tuning trainer inherits it for free:
SFTTrainer.list_supported_models()→ techniqueSFTRLVRTrainer.list_supported_models()→ techniqueRLVRRLAIFTrainer.list_supported_models()→ techniqueRLAIFDPOTrainer.list_supported_models()→ techniqueDPOCPTTrainer.list_supported_models()→ techniqueCPTRLAIFTrainerpreviously referenced its technique inline and lacked the_customization_techniqueclass attribute, so it is added there for the inheritedmethod to resolve.
MultiTurnRLTrainerkeeps its own override (it is a CreateJobflow whose technique is not a
CustomizationTechniqueenum value and delegates to_list_all_mtrl_models). If a subclass defines no technique, the method raisesNotImplementedErrorwith a clear message.Tests
test_base_trainer_compute.py: base classmethod delegates with the class'stechnique; raises
NotImplementedErrorwhen the technique is missing.test_sft_trainer.py/test_rlvr_trainer.py/test_dpo_trainer.py/test_rlaif_trainer.py: each subclass resolves and passes its own technique(
SFT/RLVR/DPO/RLAIF) through the inherited method, including sessionpassthrough.
All 250 unit tests across the touched trainer suites (including the existing MTRL
suite, which confirms the base method does not clash with MTRL's override) pass.
Manual verification
N/A — unit coverage is sufficient; the change is a thin delegation to the already
integration-tested
_list_hub_models_by_recipehelper.X-AI-Prompt: Add list_supported_models to SFT/RLVR/RLAIF/DPO trainers like MTRL; refactored to a shared BaseTrainer method
X-AI-Tool: Kiro