diff --git a/packages/nmp_common/src/nmp/common/api/utils.py b/packages/nmp_common/src/nmp/common/api/utils.py index d03d32aeac..518d66d9a8 100644 --- a/packages/nmp_common/src/nmp/common/api/utils.py +++ b/packages/nmp_common/src/nmp/common/api/utils.py @@ -6,6 +6,7 @@ import json import logging import os +from collections import defaultdict from copy import deepcopy from datetime import datetime from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union @@ -147,35 +148,58 @@ def _walk_spec(d: Dict, visitor: Callable[[str, Any, Dict], None]): _walk_spec(item, visitor) -def _normalize_refs_and_schema_keys(spec: Dict) -> Dict: +def _normalize_refs_and_schema_keys(spec: Dict, *, strict_collisions: bool = False) -> Dict: """Normalize all ``$ref`` values and schema dictionary keys. Schema keys are renamed first so that ``$ref`` values can be rewritten consistently. When two raw keys normalize to the same target and have - identical content the duplicate is silently dropped. When the content - differs a warning is logged and the *existing* schema wins (the - duplicate is still dropped). + identical content the duplicate is silently dropped. + + When the content *differs*, two distinct Pydantic models are fighting over + one schema name, and keeping one silently makes the other's ``$ref``\\ s point + at the wrong contract. With ``strict_collisions`` this raises ``ValueError`` + so the build fails loudly — used for self-contained plugin specs (e.g. the + merged ``/apis/customization`` app, where each backend must namespace its + own models). Without it (the default, used for the platform/service specs) + it logs a warning and keeps the first-seen schema, preserving legacy + behaviour: the platform spec carries pre-existing such collisions that + predate this gate and are tracked separately. """ schemas = spec["components"]["schemas"] - rename_map: Dict[str, str] = {} - - for old_key, value in list(schemas.items()): - new_key = normalize_schema_name(old_key) - if new_key == old_key: - continue - - if new_key in schemas and schemas[new_key] != value: - logger.warning( - "Schema key collision: '%s' renamed to '%s' which already exists " - "with different content. Keeping the existing schema. This likely " - "means two distinct Pydantic models from different modules share " - "the same class name.", - old_key, - new_key, + # Normalize each raw key once, then group by the name it maps to. Two + # distinct models from different modules (e.g. each customization backend's + # own ``TrainingSpec``) produce module-qualified raw keys that normalize to + # the same bare name; collapsing them would make one silently steal the + # other's ``$ref``\\ s, shipping a wrong contract. + key_to_target: Dict[str, str] = {old_key: normalize_schema_name(old_key) for old_key in schemas} + + by_target: Dict[str, List[str]] = defaultdict(list) + for old_key, target in key_to_target.items(): + by_target[target].append(old_key) + + for target, old_keys in by_target.items(): + # Identical-content duplicates are harmless — they dedup to one schema. + # Keep one representative per distinct content shape; more than one means + # genuinely different models are colliding on a single name. + reps: List[str] = [] + for key in old_keys: + if not any(schemas[key] == schemas[rep] for rep in reps): + reps.append(key) + if len(reps) > 1: + message = ( + f"OpenAPI schema name collision: {sorted(old_keys)} all normalize to " + f"'{target}' with differing content. Two distinct Pydantic models share " + f"a class name across modules — namespace them (e.g. via a per-backend " + f"NamespacedModel base) so they emit distinct schema names." ) + if strict_collisions: + raise ValueError(message) + logger.warning("%s Keeping the first-seen schema.", message) - rename_map[old_key] = new_key + # Iterating ``key_to_target`` preserves ``schemas`` order, so the first raw + # key that maps to a given target still wins the collapse below. + rename_map: Dict[str, str] = {old_key: target for old_key, target in key_to_target.items() if target != old_key} for old_key, new_key in rename_map.items(): if new_key not in schemas: @@ -331,9 +355,9 @@ def _sort_schemas(spec: Dict) -> Dict: return spec -def tweak_spec(spec: Dict) -> Dict: +def tweak_spec(spec: Dict, *, strict_collisions: bool = False) -> Dict: _walk_spec(spec, _anyof_null_visitor) - spec = _normalize_refs_and_schema_keys(spec) + spec = _normalize_refs_and_schema_keys(spec, strict_collisions=strict_collisions) spec = _split_input_output_schemas(spec) spec = _annotate_string_references(spec) spec = _sync_schema_titles(spec) diff --git a/packages/nmp_common/tests/api/test_utils_openapi_spec.py b/packages/nmp_common/tests/api/test_utils_openapi_spec.py index 6a6bc9e0af..77c15e43f4 100644 --- a/packages/nmp_common/tests/api/test_utils_openapi_spec.py +++ b/packages/nmp_common/tests/api/test_utils_openapi_spec.py @@ -328,6 +328,93 @@ def test_tweak_spec_full_pipeline(): assert result == expected +def _collision_spec(): + """Two distinct models sharing a class name across modules that normalize to + the same bare name with *differing* content.""" + return { + "components": { + "schemas": { + "automodel__schema__TrainingSpec": { + "type": "object", + "properties": {"finetuning_type": {"type": "string"}}, + }, + "unsloth__schemas__TrainingSpec": { + "type": "object", + "properties": {"use_gradient_checkpointing": {"type": "string"}}, + }, + } + }, + "paths": { + "/a": { + "post": { + "requestBody": { + "content": {"application/json": {"schema": {"$ref": REF + "automodel__schema__TrainingSpec"}}} + } + } + } + }, + } + + +def test_tweak_spec_raises_on_collision_when_strict(): + """With ``strict_collisions`` (plugin specs, e.g. the customization app) a + differing-content collision must fail the build loudly rather than silently + keeping one and mis-pointing the other's ``$ref``s.""" + with pytest.raises(ValueError, match="schema name collision"): + tweak_spec(_collision_spec(), strict_collisions=True) + + +def test_tweak_spec_warns_and_collapses_on_collision_by_default(caplog): + """Non-strict (platform/service specs) preserves legacy behaviour: warn and + keep the first-seen schema, so pre-existing platform collisions don't newly + break the build.""" + import logging + + with caplog.at_level(logging.WARNING, logger="nmp.common.api.utils"): + result = tweak_spec(_collision_spec()) + + assert "schema name collision" in caplog.text + schemas = result["components"]["schemas"] + assert set(schemas) == {"TrainingSpec"} + # First-seen (automodel) wins the collapse. + assert "finetuning_type" in schemas["TrainingSpec"]["properties"] + ref = result["paths"]["/a"]["post"]["requestBody"]["content"]["application/json"]["schema"]["$ref"] + assert ref == REF + "TrainingSpec" + + +def test_tweak_spec_dedups_identical_content_module_qualified_collision(): + """Two module-qualified keys with *identical* content dedup to one schema — + no raise (the fix must not over-trigger on genuinely-shared shapes).""" + spec = { + "components": { + "schemas": { + "pkg_a__schema__Shared": { + "type": "object", + "properties": {"value": {"type": "string"}}, + }, + "pkg_b__schema__Shared": { + "type": "object", + "properties": {"value": {"type": "string"}}, + }, + } + }, + "paths": { + "/a": { + "post": { + "requestBody": { + "content": {"application/json": {"schema": {"$ref": REF + "pkg_a__schema__Shared"}}} + } + } + } + }, + } + + result = tweak_spec(spec) + assert set(result["components"]["schemas"]) == {"Shared"} + ref = result["paths"]["/a"]["post"]["requestBody"]["content"]["application/json"]["schema"]["$ref"] + assert ref == REF + "Shared" + + def test_anyof_null_collapse_preserves_format_and_write_only(): """Collapsing ``anyOf: [SecretStr, null]`` must keep ``format`` / ``writeOnly`` so SDK + docs treat the field as sensitive.""" diff --git a/packages/nmp_customization_common/src/nmp/customization_common/schema.py b/packages/nmp_customization_common/src/nmp/customization_common/schema.py new file mode 100644 index 0000000000..3490b2460a --- /dev/null +++ b/packages/nmp_customization_common/src/nmp/customization_common/schema.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Namespaced Pydantic base for customization-backend schemas. + +Every backend (automodel, unsloth, …) is merged into one FastAPI app under +``/apis/customization``, so two backends that each define e.g. ``TrainingSpec`` +collide in the generated OpenAPI. Pydantic freezes a model's JSON-schema name at +class-creation time, so the prefix has to be applied there — this metaclass does +it. ``class TrainingSpec(AutomodelSchema)`` is emitted as ``AutomodelTrainingSpec``. + +Usage: declare a per-backend subclass that sets ``__schema_namespace__`` and +inherit *that* from every model the backend owns:: + + class AutomodelSchema(NamespacedModel): + __schema_namespace__ = "Automodel" + + class TrainingSpec(AutomodelSchema): # emitted as ``AutomodelTrainingSpec`` + ... + +A model whose class name already starts with the prefix (e.g. +``AutomodelJobInput``) is left unchanged, so top-level request/response names stay +stable. +""" + +from __future__ import annotations + +from typing import ClassVar + +from pydantic import BaseModel, ConfigDict + +# ``type(BaseModel)`` is pydantic's ``ModelMetaclass``; deriving from it this way +# avoids importing pydantic internals. +_ModelMeta = type(BaseModel) + + +class _NamespacedMeta(_ModelMeta): + def __new__(mcs, name, bases, ns, **kw): + prefix = ns.get("__schema_namespace__") or next( + (p for b in bases if (p := getattr(b, "__schema_namespace__", None))), None + ) + if prefix and not name.startswith(prefix): # don't double-prefix e.g. AutomodelJobInput + name = f"{prefix}{name}" + ns["__qualname__"] = name + return super().__new__(mcs, name, bases, ns, **kw) + + +class NamespacedModel(BaseModel, metaclass=_NamespacedMeta): + """Base for backend schemas. + + Declare a per-backend subclass that sets ``__schema_namespace__`` and inherit + *that* from every model, so each backend's schemas emit distinct component + names in the merged ``/apis/customization`` OpenAPI spec. + """ + + # Dunder ClassVar → pydantic ignores it as a field; the metaclass reads it at + # class-creation time to compute the emitted schema name. + __schema_namespace__: ClassVar[str | None] = None + model_config = ConfigDict(extra="forbid") # the config every backend model already sets diff --git a/packages/nmp_customization_common/tests/test_schema.py b/packages/nmp_customization_common/tests/test_schema.py new file mode 100644 index 0000000000..71e9e40cac --- /dev/null +++ b/packages/nmp_customization_common/tests/test_schema.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the ``NamespacedModel`` backend-schema base. + +These pin the behaviour the customization OpenAPI fix relies on: each backend's +models emit a namespace-prefixed schema name at *class-creation time*, so two +backends that each define e.g. ``TrainingSpec`` don't collapse into one component +when merged into the single ``/apis/customization`` FastAPI app. +""" + +import pytest +from nmp.customization_common.schema import NamespacedModel +from pydantic import Field, ValidationError + + +class AutomodelSchema(NamespacedModel): + __schema_namespace__ = "Automodel" + + +class UnslothSchema(NamespacedModel): + __schema_namespace__ = "Unsloth" + + +# Two backends each define a leaf literally named ``TrainingSpec`` — exactly the +# collision the metaclass exists to prevent. Capture each binding under an alias +# before the next ``class TrainingSpec`` statement shadows the module name. +class TrainingSpec(AutomodelSchema): + lr: float = 1e-4 + + +AutomodelTrainingSpec = TrainingSpec + + +class TrainingSpec(UnslothSchema): # noqa: F811 — intentional same-name redefinition + steps: int = 100 + + +UnslothTrainingSpec = TrainingSpec + + +class AutomodelJobInput(AutomodelSchema): + training: AutomodelTrainingSpec = Field(default_factory=AutomodelTrainingSpec) + + +def test_leaf_models_are_prefixed_at_class_creation(): + assert AutomodelTrainingSpec.__name__ == "AutomodelTrainingSpec" + assert UnslothTrainingSpec.__name__ == "UnslothTrainingSpec" + + +def test_same_leaf_name_across_backends_emits_distinct_component_names(): + a_title = AutomodelTrainingSpec.model_json_schema()["title"] + u_title = UnslothTrainingSpec.model_json_schema()["title"] + assert a_title == "AutomodelTrainingSpec" + assert u_title == "UnslothTrainingSpec" + # The two normalize to different bare names, so the platform normalizer won't + # collapse them into one component. + assert a_title != u_title + + +def test_nested_ref_points_at_the_owning_backends_schema(): + schema = AutomodelJobInput.model_json_schema() + assert "AutomodelTrainingSpec" in schema["$defs"] + assert schema["properties"]["training"]["$ref"].endswith("/AutomodelTrainingSpec") + + +def test_top_level_name_already_prefixed_is_not_double_prefixed(): + assert AutomodelJobInput.__name__ == "AutomodelJobInput" + assert AutomodelJobInput.model_json_schema()["title"] == "AutomodelJobInput" + + +def test_backend_base_class_itself_is_not_prefixed(): + # The intermediate per-backend base keeps its own name (it is never emitted + # as a request/response schema, but must not be mangled either). + assert AutomodelSchema.__name__ == "AutomodelSchema" + assert UnslothSchema.__name__ == "UnslothSchema" + + +def test_extra_forbid_is_inherited(): + assert AutomodelJobInput.model_config.get("extra") == "forbid" + assert AutomodelJobInput.model_json_schema()["additionalProperties"] is False + with pytest.raises(ValidationError): + # Extra field rejected (validated via dict so the extra key isn't a + # static type error); training defaults, so ``bogus`` is the only fault. + AutomodelJobInput.model_validate({"training": {}, "bogus": 1}) + + +def test_schema_namespace_is_not_a_field(): + assert "__schema_namespace__" not in AutomodelJobInput.model_fields + assert AutomodelSchema.__schema_namespace__ == "Automodel" + + +def test_plain_namespaced_model_without_namespace_is_unprefixed(): + class Bare(NamespacedModel): + x: int = 0 + + assert Bare.__name__ == "Bare" + assert Bare.model_json_schema()["title"] == "Bare" + + +def test_subclass_with_own_model_config_still_inherits_extra_forbid(): + """A subclass that declares its OWN ``model_config`` must still inherit + ``extra='forbid'`` from the base. + + ``RlJobInput``, ``_TrainingBase`` and ``RlJobOutput`` each set + ``ConfigDict(protected_namespaces=())`` and rely on pydantic *merging* (not + replacing) the base config to keep rejecting unknown fields. If that ever + regressed, ``additionalProperties: false`` would silently vanish from those + request bodies and this suite would still pass without this guard. + """ + from pydantic import ConfigDict + + class WithOwnConfig(AutomodelSchema): + model_config = ConfigDict(protected_namespaces=()) + + value: int = 0 + + # Subclass's own key applied *and* the base's ``extra='forbid'`` preserved. + assert WithOwnConfig.model_config.get("protected_namespaces") == () + assert WithOwnConfig.model_config.get("extra") == "forbid" + assert WithOwnConfig.model_json_schema()["additionalProperties"] is False + with pytest.raises(ValidationError): + WithOwnConfig.model_validate({"value": 1, "bogus": 2}) diff --git a/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py b/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py index d306d05169..740787c4f5 100644 --- a/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py +++ b/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py @@ -8,7 +8,8 @@ from typing import Literal, Self from nemo_platform_plugin.integrations import IntegrationsSpec -from pydantic import BaseModel, ConfigDict, Field, model_validator +from nmp.customization_common.schema import NamespacedModel +from pydantic import Field, model_validator __all__ = [ "AutomodelJobInput", @@ -30,9 +31,16 @@ class ValidationError(ValueError): """Raised when automodel job input validation fails.""" -class LoRAParams(BaseModel): - model_config = ConfigDict(extra="forbid") +class AutomodelSchema(NamespacedModel): + """Backend base: every Automodel-owned model emits an ``Automodel``-prefixed + OpenAPI schema name (``TrainingSpec`` -> ``AutomodelTrainingSpec``), so it + can't collide with another backend's same-named model in the merged + ``/apis/customization`` spec. ``extra='forbid'`` is inherited from the base.""" + __schema_namespace__ = "Automodel" + + +class LoRAParams(AutomodelSchema): rank: int = Field(default=16, gt=0) alpha: int = Field(default=32, gt=0) dropout: float = Field(default=0.0, ge=0.0, le=1.0, description="LoRA dropout probability for regularization.") @@ -44,17 +52,13 @@ class LoRAParams(BaseModel): use_triton: bool = Field(default=True, description="Use the optimized Triton LoRA kernel.") -class DatasetSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class DatasetSpec(AutomodelSchema): training: str = Field(description="Training fileset as 'name' or 'workspace/name'.") validation: str | None = None prompt_template: str | None = None -class TrainingSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class TrainingSpec(AutomodelSchema): training_type: Literal["sft", "distillation"] = "sft" finetuning_type: Literal["lora", "all_weights", "lora_merged"] = "lora" lora: LoRAParams | None = None @@ -83,18 +87,14 @@ def _training_type_fields(self) -> Self: return self -class ScheduleSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class ScheduleSpec(AutomodelSchema): epochs: int = Field(default=1, gt=0) max_steps: int | None = Field(default=None, gt=0) val_check_interval: float | None = None seed: int | None = None -class BatchSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class BatchSpec(AutomodelSchema): global_batch_size: int = Field(default=8, gt=0) micro_batch_size: int = Field(default=1, gt=0) sequence_packing: bool = False @@ -103,9 +103,7 @@ class BatchSpec(BaseModel): ) -class OptimizerSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class OptimizerSpec(AutomodelSchema): learning_rate: float = Field(default=5e-6, gt=0.0) min_learning_rate: float | None = Field( default=None, ge=0.0, description="Minimum learning rate for the cosine decay schedule." @@ -121,9 +119,7 @@ class OptimizerSpec(BaseModel): ) -class ParallelismSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class ParallelismSpec(AutomodelSchema): num_nodes: int = Field(default=1, gt=0) num_gpus_per_node: int = Field(default=1, gt=0) tensor_parallel_size: int = Field(default=1, gt=0) @@ -133,27 +129,21 @@ class ParallelismSpec(BaseModel): sequence_parallel: bool = Field(default=False, description="Enable sequence parallelism.") -class OutputRequest(BaseModel): - model_config = ConfigDict(extra="forbid") - +class OutputRequest(AutomodelSchema): name: str description: str | None = None -class OutputResponse(BaseModel): - model_config = ConfigDict(extra="forbid") - +class OutputResponse(AutomodelSchema): name: str type: Literal["model", "adapter"] fileset: str description: str | None = None -class AutomodelJobInput(BaseModel): +class AutomodelJobInput(AutomodelSchema): """POST body / CLI JSON.""" - model_config = ConfigDict(extra="forbid") - name: str | None = None model: str dataset: DatasetSpec @@ -173,11 +163,9 @@ def reject_legacy_fields(cls, data: object) -> object: return data -class AutomodelJobOutput(BaseModel): +class AutomodelJobOutput(AutomodelSchema): """Stored canonical spec after ``to_spec()``.""" - model_config = ConfigDict(extra="forbid") - name: str | None = None model: str dataset: DatasetSpec diff --git a/plugins/nemo-customizer/openapi/openapi.yaml b/plugins/nemo-customizer/openapi/openapi.yaml index 4afce2002a..d852b91bc8 100644 --- a/plugins/nemo-customizer/openapi/openapi.yaml +++ b/plugins/nemo-customizer/openapi/openapi.yaml @@ -1142,6 +1142,49 @@ paths: $ref: '#/components/schemas/HTTPValidationError' components: schemas: + AutomodelBatchSpec: + properties: + global_batch_size: + type: integer + exclusiveMinimum: 0.0 + title: Global Batch Size + default: 8 + micro_batch_size: + type: integer + exclusiveMinimum: 0.0 + title: Micro Batch Size + default: 1 + sequence_packing: + type: boolean + title: Sequence Packing + default: false + sequence_packing_max_samples: + type: integer + exclusiveMinimum: 0.0 + title: Sequence Packing Max Samples + description: Samples analyzed to estimate the optimal pack size when packing + is enabled. + default: 1000 + additionalProperties: false + type: object + title: AutomodelBatchSpec + AutomodelDatasetSpec: + properties: + training: + type: string + title: Training + description: Training fileset as 'name' or 'workspace/name'. + validation: + title: Validation + type: string + prompt_template: + title: Prompt Template + type: string + additionalProperties: false + type: object + required: + - training + title: AutomodelDatasetSpec AutomodelJobInput: properties: name: @@ -1151,19 +1194,19 @@ components: type: string title: Model dataset: - $ref: '#/components/schemas/DatasetSpec' + $ref: '#/components/schemas/AutomodelDatasetSpec' training: - $ref: '#/components/schemas/TrainingSpec' + $ref: '#/components/schemas/AutomodelTrainingSpec' schedule: - $ref: '#/components/schemas/ScheduleSpec' + $ref: '#/components/schemas/AutomodelScheduleSpec' batch: - $ref: '#/components/schemas/BatchSpec' + $ref: '#/components/schemas/AutomodelBatchSpec' optimizer: - $ref: '#/components/schemas/OptimizerSpec' + $ref: '#/components/schemas/AutomodelOptimizerSpec' parallelism: - $ref: '#/components/schemas/ParallelismSpec' + $ref: '#/components/schemas/AutomodelParallelismSpec' output: - $ref: '#/components/schemas/OutputRequest' + $ref: '#/components/schemas/AutomodelOutputRequest' integrations: $ref: '#/components/schemas/IntegrationsSpecInput' additionalProperties: false @@ -1183,19 +1226,19 @@ components: type: string title: Model dataset: - $ref: '#/components/schemas/DatasetSpec' + $ref: '#/components/schemas/AutomodelDatasetSpec' training: - $ref: '#/components/schemas/TrainingSpec' + $ref: '#/components/schemas/AutomodelTrainingSpec' schedule: - $ref: '#/components/schemas/ScheduleSpec' + $ref: '#/components/schemas/AutomodelScheduleSpec' batch: - $ref: '#/components/schemas/BatchSpec' + $ref: '#/components/schemas/AutomodelBatchSpec' optimizer: - $ref: '#/components/schemas/OptimizerSpec' + $ref: '#/components/schemas/AutomodelOptimizerSpec' parallelism: - $ref: '#/components/schemas/ParallelismSpec' + $ref: '#/components/schemas/AutomodelParallelismSpec' output: - $ref: '#/components/schemas/OutputResponse' + $ref: '#/components/schemas/AutomodelOutputResponse' integrations: $ref: '#/components/schemas/IntegrationsSpecOutput' additionalProperties: false @@ -1349,201 +1392,283 @@ components: - updated_at - -updated_at title: AutomodelJobsJobsSortField - BatchSpec: + AutomodelLoRAParams: properties: - global_batch_size: + rank: type: integer exclusiveMinimum: 0.0 - title: Global Batch Size - default: 8 - micro_batch_size: + title: Rank + default: 16 + alpha: type: integer exclusiveMinimum: 0.0 - title: Micro Batch Size - default: 1 - sequence_packing: + title: Alpha + default: 32 + dropout: + type: number + maximum: 1.0 + minimum: 0.0 + title: Dropout + description: LoRA dropout probability for regularization. + default: 0.0 + merge: type: boolean - title: Sequence Packing + title: Merge default: false - sequence_packing_max_samples: - type: integer - exclusiveMinimum: 0.0 - title: Sequence Packing Max Samples - description: Samples analyzed to estimate the optimal pack size when packing - is enabled. - default: 1000 + target_modules: + title: Target Modules + items: + type: string + type: array + exclude_modules: + title: Exclude Modules + description: Module name patterns to exclude from LoRA (e.g. ['*.out_proj']). + items: + type: string + type: array + use_triton: + type: boolean + title: Use Triton + description: Use the optimized Triton LoRA kernel. + default: true additionalProperties: false type: object - title: BatchSpec - DPOTraining: + title: AutomodelLoRAParams + AutomodelOptimizerSpec: properties: - optimizer_type: - allOf: - - $ref: '#/components/schemas/OptimizerType' - description: "Optimizer + LR-scheduler combination (AdamW/Adam \xD7 cosine-annealing/flat-LR).\ - \ Defaults to AdamW with cosine annealing." learning_rate: type: number + exclusiveMinimum: 0.0 title: Learning Rate - description: Peak learning rate. - default: 0.0001 + default: 5.0e-06 min_learning_rate: title: Min Learning Rate - description: Minimum LR for cosine decay. + description: Minimum learning rate for the cosine decay schedule. type: number + minimum: 0.0 weight_decay: type: number + minimum: 0.0 title: Weight Decay - description: Weight decay coefficient. default: 0.01 adam_beta1: type: number + exclusiveMaximum: 1.0 + minimum: 0.0 title: Adam Beta1 - description: Adam beta1. + description: Adam optimizer beta1. default: 0.9 adam_beta2: type: number + exclusiveMaximum: 1.0 + minimum: 0.0 title: Adam Beta2 - description: Adam beta2. + description: Adam optimizer beta2. default: 0.999 - adam_eps: - type: number - exclusiveMinimum: 0.0 - title: Adam Eps - description: Adam epsilon (numerical stability term). - default: 1.0e-05 warmup_steps: type: integer minimum: 0.0 title: Warmup Steps - description: Linear warmup steps. default: 0 - epochs: + adam_eps: + type: number + exclusiveMinimum: 0.0 + title: Adam Eps + description: Adam/AdamW epsilon for numerical stability. + default: 1.0e-08 + optimizer: + type: string + enum: + - Adam + - AdamW + title: Optimizer + description: Optimizer algorithm. + default: Adam + lr_decay_style: + type: string + enum: + - cosine + - linear + - constant + title: Lr Decay Style + description: Learning-rate decay schedule. + default: cosine + additionalProperties: false + type: object + title: AutomodelOptimizerSpec + AutomodelOutputRequest: + properties: + name: + type: string + title: Name + description: + title: Description + type: string + additionalProperties: false + type: object + required: + - name + title: AutomodelOutputRequest + AutomodelOutputResponse: + properties: + name: + type: string + title: Name + type: + type: string + enum: + - model + - adapter + title: Type + fileset: + type: string + title: Fileset + description: + title: Description + type: string + additionalProperties: false + type: object + required: + - name + - type + - fileset + title: AutomodelOutputResponse + AutomodelParallelismSpec: + properties: + num_nodes: type: integer exclusiveMinimum: 0.0 - title: Epochs - description: Number of passes through the dataset. + title: Num Nodes default: 1 - max_steps: - title: Max Steps - description: Max training steps (overrides epochs if set). + num_gpus_per_node: type: integer exclusiveMinimum: 0.0 - val_check_interval: - title: Val Check Interval - description: Validation interval. Float <= 1.0 is fraction of epoch; > 1.0 - is step count. - type: number - val_at_end: - type: boolean - title: Val At End - description: Run a final validation pass after the last training step. Keep - enabled so the final checkpoint carries validation metrics and best-checkpoint - selection works; set False only to skip the extra eval. - default: true - keep_top_k: + title: Num Gpus Per Node + default: 1 + tensor_parallel_size: type: integer exclusiveMinimum: 0.0 - title: Keep Top K - description: Number of best checkpoints to retain (ranked by validation - loss). + title: Tensor Parallel Size default: 1 - batch_size: + pipeline_parallel_size: type: integer exclusiveMinimum: 0.0 - title: Batch Size - description: Global batch size across all GPUs. - default: 32 - micro_batch_size: + title: Pipeline Parallel Size + default: 1 + context_parallel_size: type: integer exclusiveMinimum: 0.0 - title: Micro Batch Size - description: Per-GPU micro batch size. + title: Context Parallel Size default: 1 - activation_checkpointing: + expert_parallel_size: + title: Expert Parallel Size + type: integer + exclusiveMinimum: 0.0 + sequence_parallel: type: boolean - title: Activation Checkpointing - description: Recompute activations during the backward pass to reduce memory - at the cost of compute. Enable to fit larger models or longer sequences. + title: Sequence Parallel + description: Enable sequence parallelism. default: false - max_seq_length: + additionalProperties: false + type: object + title: AutomodelParallelismSpec + AutomodelScheduleSpec: + properties: + epochs: type: integer exclusiveMinimum: 0.0 - title: Max Seq Length - description: Maximum token sequence length for training. - default: 2048 + title: Epochs + default: 1 + max_steps: + title: Max Steps + type: integer + exclusiveMinimum: 0.0 + val_check_interval: + title: Val Check Interval + type: number seed: title: Seed - description: Random seed for reproducibility. type: integer - parallelism: - $ref: '#/components/schemas/ParallelismParams' - execution_profile: - title: Execution Profile - description: Execution profile for the GPU training step (operator-configured). - Falls back to the service default when omitted. - type: string - minLength: 1 - type: + additionalProperties: false + type: object + title: AutomodelScheduleSpec + AutomodelTrainingSpec: + properties: + training_type: type: string - const: dpo - title: Type - default: dpo - ref_policy_kl_penalty: - type: number - minimum: 0.0 - title: Ref Policy Kl Penalty - description: KL penalty coefficient (beta in the DPO paper). - default: 0.05 - preference_average_log_probs: - type: boolean - title: Preference Average Log Probs - description: Average log probabilities for preference loss calculation. - default: false - sft_average_log_probs: - type: boolean - title: Sft Average Log Probs - description: Average log probabilities for SFT regularization loss. - default: false - preference_loss_weight: - type: number - minimum: 0.0 - title: Preference Loss Weight - description: Weight for the preference (DPO) loss term. - default: 1.0 - sft_loss_weight: + enum: + - sft + - distillation + title: Training Type + default: sft + finetuning_type: + type: string + enum: + - lora + - all_weights + - lora_merged + title: Finetuning Type + default: lora + lora: + $ref: '#/components/schemas/AutomodelLoRAParams' + max_seq_length: + type: integer + exclusiveMinimum: 0.0 + title: Max Seq Length + default: 2048 + precision: + title: Precision + description: Model precision for training. Auto-detected from the checkpoint + when unset. + type: string + enum: + - bf16 + - fp16 + - fp32 + - fp8 + attn_implementation: + type: string + enum: + - sdpa + - flash_attention_2 + - eager + title: Attn Implementation + description: 'Attention backend: ''sdpa'' (PyTorch native), ''flash_attention_2'', + or ''eager''.' + default: sdpa + execution_profile: + title: Execution Profile + type: string + minLength: 1 + teacher_model: + title: Teacher Model + type: string + distillation_ratio: type: number + maximum: 1.0 minimum: 0.0 - title: Sft Loss Weight - description: Weight for SFT regularization loss (0 = disabled). - default: 0.0 - max_grad_norm: + title: Distillation Ratio + default: 0.5 + distillation_temperature: type: number - minimum: 0.0 - title: Max Grad Norm - description: Maximum gradient norm for clipping. + exclusiveMinimum: 0.0 + title: Distillation Temperature default: 1.0 - additionalProperties: false - type: object - title: DPOTraining - description: "Direct Preference Optimization (full-weight only \u2014 PEFT unsupported)." - DatasetSpec: - properties: - training: - type: string - title: Training - description: Training fileset as 'name' or 'workspace/name'. - validation: - title: Validation - type: string - prompt_template: - title: Prompt Template + teacher_precision: type: string + enum: + - bf16 + - fp16 + - fp32 + title: Teacher Precision + default: bf16 + offload_teacher: + type: boolean + title: Offload Teacher + default: false additionalProperties: false type: object - required: - - training - title: DatasetSpec + title: AutomodelTrainingSpec DatetimeFilter: additionalProperties: false properties: @@ -1559,51 +1684,6 @@ components: type: string title: DatetimeFilter type: object - DeploymentParams: - properties: - gpu: - type: integer - title: Gpu - description: Number of GPUs required for the deployment. - default: 1 - additional_envs: - title: Additional Envs - description: Additional environment variables for the deployment. - additionalProperties: - type: string - type: object - disk_size: - title: Disk Size - description: Disk size for the deployment. - type: string - image_name: - title: Image Name - description: Container image name from NGC. If not specified, defaults to - multi-llm. - type: string - image_tag: - title: Image Tag - description: Container image tag from NGC. - type: string - lora_enabled: - type: boolean - title: Lora Enabled - description: When auto-deploying a full SFT training, setting this true - allows subsequent LoRA adapters to be deployed against it. - default: true - tool_call_config: - allOf: - - $ref: '#/components/schemas/ToolCallParams' - description: Tool calling configuration override for the NIM deployment. - additionalProperties: false - type: object - title: DeploymentParams - description: 'Inline deployment parameters for auto-deploying a trained model. - - - Used in :class:`UnslothJobInput.deployment_config` and passed through to - - the model_entity task at compile time. When unset, no deployment is launched.' FileStorageType: type: string enum: @@ -1618,24 +1698,6 @@ components: title: Detail type: object title: HTTPValidationError - HardwareSpec: - properties: - gpus: - title: Gpus - description: Comma-separated GPU indices ('0' or '0,1') for CUDA_VISIBLE_DEVICES. - Selection, not reservation. - type: string - precision: - type: string - enum: - - bf16 - - fp16 - title: Precision - description: Mixed-precision dtype for training. bf16 recommended for Ampere+. - default: bf16 - additionalProperties: false - type: object - title: HardwareSpec IntegrationsSpecInput: properties: wandb: @@ -1680,48 +1742,6 @@ components: disable it. Activation at training time still requires credentials/URIs (see compile-time warnings and runtime builders).' - LoRAParams: - properties: - rank: - type: integer - exclusiveMinimum: 0.0 - title: Rank - default: 16 - alpha: - type: integer - exclusiveMinimum: 0.0 - title: Alpha - default: 32 - dropout: - type: number - maximum: 1.0 - minimum: 0.0 - title: Dropout - description: LoRA dropout probability for regularization. - default: 0.0 - merge: - type: boolean - title: Merge - default: false - target_modules: - title: Target Modules - items: - type: string - type: array - exclude_modules: - title: Exclude Modules - description: Module name patterns to exclude from LoRA (e.g. ['*.out_proj']). - items: - type: string - type: array - use_triton: - type: boolean - title: Use Triton - description: Use the optimized Triton LoRA kernel. - default: true - additionalProperties: false - type: object - title: LoRAParams MlflowIntegration: properties: experiment_name: @@ -1755,190 +1775,26 @@ components: To enable MLflow, provide a non-null ``mlflow`` object on :class:`IntegrationsSpec`.' - ModelLoadSpec: + OptimizerType: + type: string + enum: + - adamw_with_cosine_annealing + - adam_with_cosine_annealing + - adamw_with_flat_lr + - adam_with_flat_lr + title: OptimizerType + description: Optimizer and scheduler combination types. + OutputNameType: + type: string + enum: + - adapter + - model + title: OutputNameType + description: "Output artifact type \u2014 adapter (LoRA only) or model (merged\ + \ / full)." + PaginationData: properties: - name: - type: string - title: Name - description: Model entity reference. Accepts 'name' (uses the job's workspace) - or 'workspace/name'. The plugin's run resolves this to a local path before - training. - max_seq_length: - type: integer - exclusiveMinimum: 0.0 - title: Max Seq Length - default: 2048 - load_in_4bit: - type: boolean - title: Load In 4Bit - description: bitsandbytes 4-bit. Mutex with load_in_8bit. Default for Unsloth's - headline path. - default: true - load_in_8bit: - type: boolean - title: Load In 8Bit - default: false - dtype: - type: string - enum: - - auto - - bfloat16 - - float16 - - float32 - title: Dtype - default: auto - trust_remote_code: - type: boolean - title: Trust Remote Code - default: false - device_map: - anyOf: - - type: string - - type: integer - - additionalProperties: - type: integer - type: object - title: Device Map - description: "Device placement forwarded to FastLanguageModel.from_pretrained.\ - \ Omit (null) to pin the whole model to the single visible GPU ({'': 0})\ - \ \u2014 the right default for this single-GPU backend, and it avoids\ - \ accelerate's auto-placement under-sizing GPU memory on unified-memory\ - \ parts (e.g. GB10 / DGX Spark), which otherwise spills layers to CPU\ - \ and aborts 4-bit loads. Set 'auto', 'balanced', 'sequential', a device\ - \ index, or a custom map for multi-device experiments." - rope_scaling: - title: Rope Scaling - description: 'RoPE scaling config for long-context extension, passed to - FastLanguageModel.from_pretrained (e.g. {''type'': ''linear'', ''factor'': - 2.0}). None uses the model''s native context length.' - additionalProperties: true - type: object - additionalProperties: false - type: object - required: - - name - title: ModelLoadSpec - description: 'Args to ``FastLanguageModel.from_pretrained``. - - - ``name`` is a NeMo Platform model entity reference (``"name"`` or - - ``"workspace/name"``). The plugin''s run orchestration resolves the - - entity, downloads its fileset to a local path, and hands that path - - to :func:`train_sft`.' - OptimizerSpec: - properties: - learning_rate: - type: number - exclusiveMinimum: 0.0 - title: Learning Rate - default: 5.0e-06 - min_learning_rate: - title: Min Learning Rate - description: Minimum learning rate for the cosine decay schedule. - type: number - minimum: 0.0 - weight_decay: - type: number - minimum: 0.0 - title: Weight Decay - default: 0.01 - adam_beta1: - type: number - exclusiveMaximum: 1.0 - minimum: 0.0 - title: Adam Beta1 - description: Adam optimizer beta1. - default: 0.9 - adam_beta2: - type: number - exclusiveMaximum: 1.0 - minimum: 0.0 - title: Adam Beta2 - description: Adam optimizer beta2. - default: 0.999 - warmup_steps: - type: integer - minimum: 0.0 - title: Warmup Steps - default: 0 - adam_eps: - type: number - exclusiveMinimum: 0.0 - title: Adam Eps - description: Adam/AdamW epsilon for numerical stability. - default: 1.0e-08 - optimizer: - type: string - enum: - - Adam - - AdamW - title: Optimizer - description: Optimizer algorithm. - default: Adam - lr_decay_style: - type: string - enum: - - cosine - - linear - - constant - title: Lr Decay Style - description: Learning-rate decay schedule. - default: cosine - additionalProperties: false - type: object - title: OptimizerSpec - OptimizerType: - type: string - enum: - - adamw_with_cosine_annealing - - adam_with_cosine_annealing - - adamw_with_flat_lr - - adam_with_flat_lr - title: OptimizerType - description: Optimizer and scheduler combination types. - OutputRequest: - properties: - name: - type: string - title: Name - description: - title: Description - type: string - additionalProperties: false - type: object - required: - - name - title: OutputRequest - OutputResponse: - properties: - name: - type: string - title: Name - type: - type: string - enum: - - model - - adapter - title: Type - fileset: - type: string - title: Fileset - description: - title: Description - type: string - additionalProperties: false - type: object - required: - - name - - type - - fileset - title: OutputResponse - PaginationData: - properties: - page: + page: type: integer title: Page description: The current page number. @@ -1966,93 +1822,6 @@ components: - total_pages - total_results title: PaginationData - ParallelismParams: - properties: - num_gpus_per_node: - type: integer - exclusiveMinimum: 0.0 - title: Num Gpus Per Node - description: Number of GPUs per node. - default: 1 - num_nodes: - type: integer - exclusiveMinimum: 0.0 - title: Num Nodes - description: "Number of nodes (>1 \u2192 multi-node Ray cluster)." - default: 1 - tensor_parallel_size: - type: integer - exclusiveMinimum: 0.0 - title: Tensor Parallel Size - description: Tensor parallel size. - default: 1 - pipeline_parallel_size: - type: integer - exclusiveMinimum: 0.0 - title: Pipeline Parallel Size - description: Pipeline parallel size. - default: 1 - context_parallel_size: - type: integer - exclusiveMinimum: 0.0 - title: Context Parallel Size - description: Context parallel size. - default: 1 - sequence_parallel: - type: boolean - title: Sequence Parallel - description: Enable sequence parallelism. - default: false - additionalProperties: false - type: object - title: ParallelismParams - description: 'Distributed training parallelism configuration. - - - Single-node multi-GPU uses ``num_nodes=1`` with ``num_gpus_per_node>1``; - - multi-node sets ``num_nodes>1`` and the compiler emits a distributed-GPU - - executor (see :mod:`nmp.rl.app.jobs.compiler`).' - ParallelismSpec: - properties: - num_nodes: - type: integer - exclusiveMinimum: 0.0 - title: Num Nodes - default: 1 - num_gpus_per_node: - type: integer - exclusiveMinimum: 0.0 - title: Num Gpus Per Node - default: 1 - tensor_parallel_size: - type: integer - exclusiveMinimum: 0.0 - title: Tensor Parallel Size - default: 1 - pipeline_parallel_size: - type: integer - exclusiveMinimum: 0.0 - title: Pipeline Parallel Size - default: 1 - context_parallel_size: - type: integer - exclusiveMinimum: 0.0 - title: Context Parallel Size - default: 1 - expert_parallel_size: - title: Expert Parallel Size - type: integer - exclusiveMinimum: 0.0 - sequence_parallel: - type: boolean - title: Sequence Parallel - description: Enable sequence parallelism. - default: false - additionalProperties: false - type: object - title: ParallelismSpec PlatformJobListResultResponse: properties: data: @@ -2295,40 +2064,192 @@ components: - created_at - updated_at title: PlatformJobTaskStatusResponse - RlJobInput: + RlDPOTraining: properties: - name: - title: Name - type: string - model: - type: string - title: Model - description: Model entity reference ('name' or 'workspace/name'). - dataset: - type: string - title: Dataset - description: Preference dataset fileset reference. Must contain training.jsonl - + validation.jsonl. - training: + optimizer_type: allOf: - - $ref: '#/components/schemas/DPOTraining' - description: DPO training method and hyperparameters. - integrations: - $ref: '#/components/schemas/IntegrationsSpecInput' - output: - $ref: '#/components/schemas/OutputRequest' - additionalProperties: false - type: object - required: - - model - - dataset - - training - title: RlJobInput - description: POST body / CLI JSON for ``nemo customization rl submit``. - RlJobOutput: - properties: - name: - title: Name + - $ref: '#/components/schemas/OptimizerType' + description: "Optimizer + LR-scheduler combination (AdamW/Adam \xD7 cosine-annealing/flat-LR).\ + \ Defaults to AdamW with cosine annealing." + learning_rate: + type: number + title: Learning Rate + description: Peak learning rate. + default: 0.0001 + min_learning_rate: + title: Min Learning Rate + description: Minimum LR for cosine decay. + type: number + weight_decay: + type: number + title: Weight Decay + description: Weight decay coefficient. + default: 0.01 + adam_beta1: + type: number + title: Adam Beta1 + description: Adam beta1. + default: 0.9 + adam_beta2: + type: number + title: Adam Beta2 + description: Adam beta2. + default: 0.999 + adam_eps: + type: number + exclusiveMinimum: 0.0 + title: Adam Eps + description: Adam epsilon (numerical stability term). + default: 1.0e-05 + warmup_steps: + type: integer + minimum: 0.0 + title: Warmup Steps + description: Linear warmup steps. + default: 0 + epochs: + type: integer + exclusiveMinimum: 0.0 + title: Epochs + description: Number of passes through the dataset. + default: 1 + max_steps: + title: Max Steps + description: Max training steps (overrides epochs if set). + type: integer + exclusiveMinimum: 0.0 + val_check_interval: + title: Val Check Interval + description: Validation interval. Float <= 1.0 is fraction of epoch; > 1.0 + is step count. + type: number + val_at_end: + type: boolean + title: Val At End + description: Run a final validation pass after the last training step. Keep + enabled so the final checkpoint carries validation metrics and best-checkpoint + selection works; set False only to skip the extra eval. + default: true + keep_top_k: + type: integer + exclusiveMinimum: 0.0 + title: Keep Top K + description: Number of best checkpoints to retain (ranked by validation + loss). + default: 1 + batch_size: + type: integer + exclusiveMinimum: 0.0 + title: Batch Size + description: Global batch size across all GPUs. + default: 32 + micro_batch_size: + type: integer + exclusiveMinimum: 0.0 + title: Micro Batch Size + description: Per-GPU micro batch size. + default: 1 + activation_checkpointing: + type: boolean + title: Activation Checkpointing + description: Recompute activations during the backward pass to reduce memory + at the cost of compute. Enable to fit larger models or longer sequences. + default: false + max_seq_length: + type: integer + exclusiveMinimum: 0.0 + title: Max Seq Length + description: Maximum token sequence length for training. + default: 2048 + seed: + title: Seed + description: Random seed for reproducibility. + type: integer + parallelism: + $ref: '#/components/schemas/RlParallelismParams' + execution_profile: + title: Execution Profile + description: Execution profile for the GPU training step (operator-configured). + Falls back to the service default when omitted. + type: string + minLength: 1 + type: + type: string + const: dpo + title: Type + default: dpo + ref_policy_kl_penalty: + type: number + minimum: 0.0 + title: Ref Policy Kl Penalty + description: KL penalty coefficient (beta in the DPO paper). + default: 0.05 + preference_average_log_probs: + type: boolean + title: Preference Average Log Probs + description: Average log probabilities for preference loss calculation. + default: false + sft_average_log_probs: + type: boolean + title: Sft Average Log Probs + description: Average log probabilities for SFT regularization loss. + default: false + preference_loss_weight: + type: number + minimum: 0.0 + title: Preference Loss Weight + description: Weight for the preference (DPO) loss term. + default: 1.0 + sft_loss_weight: + type: number + minimum: 0.0 + title: Sft Loss Weight + description: Weight for SFT regularization loss (0 = disabled). + default: 0.0 + max_grad_norm: + type: number + minimum: 0.0 + title: Max Grad Norm + description: Maximum gradient norm for clipping. + default: 1.0 + additionalProperties: false + type: object + title: RlDPOTraining + description: "Direct Preference Optimization (full-weight only \u2014 PEFT unsupported)." + RlJobInput: + properties: + name: + title: Name + type: string + model: + type: string + title: Model + description: Model entity reference ('name' or 'workspace/name'). + dataset: + type: string + title: Dataset + description: Preference dataset fileset reference. Must contain training.jsonl + + validation.jsonl. + training: + allOf: + - $ref: '#/components/schemas/RlDPOTraining' + description: DPO training method and hyperparameters. + integrations: + $ref: '#/components/schemas/IntegrationsSpecInput' + output: + $ref: '#/components/schemas/RlOutputRequest' + additionalProperties: false + type: object + required: + - model + - dataset + - training + title: RlJobInput + description: POST body / CLI JSON for ``nemo customization rl submit``. + RlJobOutput: + properties: + name: + title: Name description: Optional job name; auto-generated when omitted. type: string model: @@ -2341,7 +2262,7 @@ components: description: Preference dataset fileset reference ('name' or 'workspace/name'). training: allOf: - - $ref: '#/components/schemas/DPOTraining' + - $ref: '#/components/schemas/RlDPOTraining' description: Training method and hyperparameters (DPO). integrations: allOf: @@ -2349,8 +2270,9 @@ components: description: W&B / MLflow integrations. output: allOf: - - $ref: '#/components/schemas/OutputResponse' + - $ref: '#/components/schemas/RlOutputResponse' description: Output artifact created by this job. + additionalProperties: false type: object required: - model @@ -2504,26 +2426,91 @@ components: - updated_at - -updated_at title: RlJobsJobsSortField - ScheduleSpec: + RlOutputRequest: properties: - epochs: + name: + title: Name + type: string + additionalProperties: false + type: object + title: RlOutputRequest + description: Submitter-facing output preferences. ``name`` is auto-derived if + omitted. + RlOutputResponse: + properties: + name: + type: string + maxLength: 255 + title: Name + description: Name of the output artifact. Used to identify it during deployment + and inference. + examples: + - my-dpo-llama + type: + allOf: + - $ref: '#/components/schemas/OutputNameType' + description: Output artifact type. DPO is full-weight, so always `model`. + default: model + fileset: + type: string + maxLength: 255 + title: Fileset + description: FileSet name where output artifacts are stored. + additionalProperties: false + type: object + required: + - name + - fileset + title: RlOutputResponse + description: Resolved output artifact details. + RlParallelismParams: + properties: + num_gpus_per_node: type: integer exclusiveMinimum: 0.0 - title: Epochs + title: Num Gpus Per Node + description: Number of GPUs per node. default: 1 - max_steps: - title: Max Steps + num_nodes: type: integer exclusiveMinimum: 0.0 - val_check_interval: - title: Val Check Interval - type: number - seed: - title: Seed + title: Num Nodes + description: "Number of nodes (>1 \u2192 multi-node Ray cluster)." + default: 1 + tensor_parallel_size: + type: integer + exclusiveMinimum: 0.0 + title: Tensor Parallel Size + description: Tensor parallel size. + default: 1 + pipeline_parallel_size: + type: integer + exclusiveMinimum: 0.0 + title: Pipeline Parallel Size + description: Pipeline parallel size. + default: 1 + context_parallel_size: type: integer + exclusiveMinimum: 0.0 + title: Context Parallel Size + description: Context parallel size. + default: 1 + sequence_parallel: + type: boolean + title: Sequence Parallel + description: Enable sequence parallelism. + default: false additionalProperties: false type: object - title: ScheduleSpec + title: RlParallelismParams + description: 'Distributed training parallelism configuration. + + + Single-node multi-GPU uses ``num_nodes=1`` with ``num_gpus_per_node>1``; + + multi-node sets ``num_nodes>1`` and the compiler emits a distributed-GPU + + executor (see :mod:`nmp.rl.app.jobs.compiler`).' SecretRef: type: string pattern: ^[a-z0-9_-]+(/[a-z0-9_-]+)?$ @@ -2555,133 +2542,156 @@ components: type: array title: StringFilter type: object - ToolCallParams: + UnslothBatchSpec: properties: - tool_call_parser: - title: Tool Call Parser - description: Name of the tool call parser to use (e.g., 'openai', 'hermes', - 'pythonic', 'llama3_json', 'mistral'). - type: string - tool_call_plugin: - title: Tool Call Plugin - description: 'Reference to a fileset containing the custom tool call plugin - Python file. Expected format: ''{workspace}/{fileset_name}''.' - type: string - pattern: ^[\w\-.]+/[\w\-.]+$ - auto_tool_choice: - title: Auto Tool Choice - description: Whether to enable automatic tool choice. - type: boolean + per_device_train_batch_size: + type: integer + exclusiveMinimum: 0.0 + title: Per Device Train Batch Size + default: 1 + gradient_accumulation_steps: + type: integer + exclusiveMinimum: 0.0 + title: Gradient Accumulation Steps + default: 1 additionalProperties: false type: object - title: ToolCallParams - description: Tool calling configuration for NIM deployments. - TrainingSpec: + title: UnslothBatchSpec + UnslothDatasetSpec: properties: - training_type: + path: type: string - enum: - - sft - - distillation - title: Training Type - default: sft - finetuning_type: + title: Path + description: 'Training fileset reference: ''name'' (uses the job''s workspace) + or ''workspace/name''. Resolved to a local path by the plugin run.' + text_field: type: string - enum: - - lora - - all_weights - - lora_merged - title: Finetuning Type - default: lora - lora: - $ref: '#/components/schemas/LoRAParams' - max_seq_length: + title: Text Field + description: Row field consumed by SFTTrainer. + default: text + apply_chat_template: + type: boolean + title: Apply Chat Template + description: If True, expects rows with a 'messages' field and applies tokenizer.apply_chat_template + at training time. + default: false + validation_path: + title: Validation Path + description: Optional validation fileset reference (same format as 'path'). + Downloaded under the same scheme. + type: string + packing: + type: boolean + title: Packing + description: trl.SFTTrainer packing flag. + default: false + additionalProperties: false + type: object + required: + - path + title: UnslothDatasetSpec + description: 'Training data location + shape. + + + ``path`` and ``validation_path`` are platform fileset references + + (``"name"`` or ``"workspace/name"``). The plugin''s run downloads + + each fileset before training; ``train_sft`` only ever sees a local + + filesystem path.' + UnslothDeploymentParams: + properties: + gpu: type: integer - exclusiveMinimum: 0.0 - title: Max Seq Length - default: 2048 - precision: - title: Precision - description: Model precision for training. Auto-detected from the checkpoint - when unset. + title: Gpu + description: Number of GPUs required for the deployment. + default: 1 + additional_envs: + title: Additional Envs + description: Additional environment variables for the deployment. + additionalProperties: + type: string + type: object + disk_size: + title: Disk Size + description: Disk size for the deployment. type: string - enum: - - bf16 - - fp16 - - fp32 - - fp8 - attn_implementation: + image_name: + title: Image Name + description: Container image name from NGC. If not specified, defaults to + multi-llm. type: string - enum: - - sdpa - - flash_attention_2 - - eager - title: Attn Implementation - description: 'Attention backend: ''sdpa'' (PyTorch native), ''flash_attention_2'', - or ''eager''.' - default: sdpa - execution_profile: - title: Execution Profile + image_tag: + title: Image Tag + description: Container image tag from NGC. type: string - minLength: 1 - teacher_model: - title: Teacher Model + lora_enabled: + type: boolean + title: Lora Enabled + description: When auto-deploying a full SFT training, setting this true + allows subsequent LoRA adapters to be deployed against it. + default: true + tool_call_config: + allOf: + - $ref: '#/components/schemas/UnslothToolCallParams' + description: Tool calling configuration override for the NIM deployment. + additionalProperties: false + type: object + title: UnslothDeploymentParams + description: 'Inline deployment parameters for auto-deploying a trained model. + + + Used in :class:`UnslothJobInput.deployment_config` and passed through to + + the model_entity task at compile time. When unset, no deployment is launched.' + UnslothHardwareSpec: + properties: + gpus: + title: Gpus + description: Comma-separated GPU indices ('0' or '0,1') for CUDA_VISIBLE_DEVICES. + Selection, not reservation. type: string - distillation_ratio: - type: number - maximum: 1.0 - minimum: 0.0 - title: Distillation Ratio - default: 0.5 - distillation_temperature: - type: number - exclusiveMinimum: 0.0 - title: Distillation Temperature - default: 1.0 - teacher_precision: + precision: type: string enum: - bf16 - fp16 - - fp32 - title: Teacher Precision + title: Precision + description: Mixed-precision dtype for training. bf16 recommended for Ampere+. default: bf16 - offload_teacher: - type: boolean - title: Offload Teacher - default: false additionalProperties: false type: object - title: TrainingSpec + title: UnslothHardwareSpec UnslothJobInput: properties: name: title: Name type: string model: - $ref: '#/components/schemas/ModelLoadSpec' + $ref: '#/components/schemas/UnslothModelLoadSpec' dataset: - $ref: '#/components/schemas/DatasetSpec' + $ref: '#/components/schemas/UnslothDatasetSpec' training: - $ref: '#/components/schemas/TrainingSpec' + $ref: '#/components/schemas/UnslothTrainingSpec' schedule: - $ref: '#/components/schemas/ScheduleSpec' + $ref: '#/components/schemas/UnslothScheduleSpec' batch: - $ref: '#/components/schemas/BatchSpec' + $ref: '#/components/schemas/UnslothBatchSpec' optimizer: - $ref: '#/components/schemas/OptimizerSpec' + $ref: '#/components/schemas/UnslothOptimizerSpec' hardware: - $ref: '#/components/schemas/HardwareSpec' + $ref: '#/components/schemas/UnslothHardwareSpec' integrations: $ref: '#/components/schemas/IntegrationsSpecInput' output: - $ref: '#/components/schemas/OutputRequest' + $ref: '#/components/schemas/UnslothOutputRequest' deployment_config: anyOf: - type: string title: Reference - description: A reference to DeploymentParams. - - $ref: '#/components/schemas/DeploymentParams' + description: A reference to UnslothDeploymentParams. + - $ref: '#/components/schemas/UnslothDeploymentParams' title: Deployment Config description: Deployment configuration for auto-deploying the model after training. Pass a string to reference an existing ModelDeploymentConfig @@ -2700,29 +2710,29 @@ components: title: Name type: string model: - $ref: '#/components/schemas/ModelLoadSpec' + $ref: '#/components/schemas/UnslothModelLoadSpec' dataset: - $ref: '#/components/schemas/DatasetSpec' + $ref: '#/components/schemas/UnslothDatasetSpec' training: - $ref: '#/components/schemas/TrainingSpec' + $ref: '#/components/schemas/UnslothTrainingSpec' schedule: - $ref: '#/components/schemas/ScheduleSpec' + $ref: '#/components/schemas/UnslothScheduleSpec' batch: - $ref: '#/components/schemas/BatchSpec' + $ref: '#/components/schemas/UnslothBatchSpec' optimizer: - $ref: '#/components/schemas/OptimizerSpec' + $ref: '#/components/schemas/UnslothOptimizerSpec' hardware: - $ref: '#/components/schemas/HardwareSpec' + $ref: '#/components/schemas/UnslothHardwareSpec' integrations: $ref: '#/components/schemas/IntegrationsSpecOutput' output: - $ref: '#/components/schemas/OutputResponse' + $ref: '#/components/schemas/UnslothOutputResponse' deployment_config: anyOf: - type: string title: Reference - description: A reference to DeploymentParams. - - $ref: '#/components/schemas/DeploymentParams' + description: A reference to UnslothDeploymentParams. + - $ref: '#/components/schemas/UnslothDeploymentParams' title: Deployment Config description: Deployment configuration for auto-deploying the model after training. Pass a string to reference an existing ModelDeploymentConfig @@ -2884,6 +2894,417 @@ components: - updated_at - -updated_at title: UnslothJobsJobsSortField + UnslothLoRAParams: + properties: + rank: + type: integer + exclusiveMinimum: 0.0 + title: Rank + description: LoRA rank. + default: 16 + alpha: + type: integer + exclusiveMinimum: 0.0 + title: Alpha + description: LoRA scaling factor (alpha). + default: 16 + dropout: + type: number + exclusiveMaximum: 1.0 + minimum: 0.0 + title: Dropout + default: 0.0 + target_modules: + items: + type: string + type: array + title: Target Modules + bias: + type: string + enum: + - none + - all + - lora_only + title: Bias + default: none + use_rslora: + type: boolean + title: Use Rslora + default: false + random_state: + type: integer + title: Random State + default: 3407 + use_dora: + type: boolean + title: Use Dora + description: DoRA (weight-decomposed LoRA). Improves quality at low ranks; + adds training overhead. + default: false + loftq_config: + title: Loftq Config + description: LoftQ initialization config for quantized bases. None disables + LoftQ. + additionalProperties: true + type: object + modules_to_save: + title: Modules To Save + description: Extra non-LoRA modules to train and save in full (e.g. ['embed_tokens', + 'lm_head']). Needed for vocab changes / continued pretraining. + items: + type: string + type: array + layers_to_transform: + anyOf: + - type: integer + - items: + type: integer + type: array + title: Layers To Transform + description: Restrict LoRA to specific layer index(es). None applies to + all layers. + layer_replication: + title: Layer Replication + description: Layer-replication ranges for stacking, e.g. [[0, 16], [8, 24]]. + None disables. + items: + items: + type: integer + type: array + type: array + init_lora_weights: + anyOf: + - type: boolean + - type: string + enum: + - gaussian + - pissa + - olora + - loftq + title: Init Lora Weights + description: LoRA weight init scheme. True = PEFT default; 'pissa'/'olora'/'loftq' + for advanced inits. + default: true + additionalProperties: false + type: object + title: UnslothLoRAParams + description: Args to ``FastLanguageModel.get_peft_model``. + UnslothModelLoadSpec: + properties: + name: + type: string + title: Name + description: Model entity reference. Accepts 'name' (uses the job's workspace) + or 'workspace/name'. The plugin's run resolves this to a local path before + training. + max_seq_length: + type: integer + exclusiveMinimum: 0.0 + title: Max Seq Length + default: 2048 + load_in_4bit: + type: boolean + title: Load In 4Bit + description: bitsandbytes 4-bit. Mutex with load_in_8bit. Default for Unsloth's + headline path. + default: true + load_in_8bit: + type: boolean + title: Load In 8Bit + default: false + dtype: + type: string + enum: + - auto + - bfloat16 + - float16 + - float32 + title: Dtype + default: auto + trust_remote_code: + type: boolean + title: Trust Remote Code + default: false + device_map: + anyOf: + - type: string + - type: integer + - additionalProperties: + type: integer + type: object + title: Device Map + description: "Device placement forwarded to FastLanguageModel.from_pretrained.\ + \ Omit (null) to pin the whole model to the single visible GPU ({'': 0})\ + \ \u2014 the right default for this single-GPU backend, and it avoids\ + \ accelerate's auto-placement under-sizing GPU memory on unified-memory\ + \ parts (e.g. GB10 / DGX Spark), which otherwise spills layers to CPU\ + \ and aborts 4-bit loads. Set 'auto', 'balanced', 'sequential', a device\ + \ index, or a custom map for multi-device experiments." + rope_scaling: + title: Rope Scaling + description: 'RoPE scaling config for long-context extension, passed to + FastLanguageModel.from_pretrained (e.g. {''type'': ''linear'', ''factor'': + 2.0}). None uses the model''s native context length.' + additionalProperties: true + type: object + additionalProperties: false + type: object + required: + - name + title: UnslothModelLoadSpec + description: 'Args to ``FastLanguageModel.from_pretrained``. + + + ``name`` is a NeMo Platform model entity reference (``"name"`` or + + ``"workspace/name"``). The plugin''s run orchestration resolves the + + entity, downloads its fileset to a local path, and hands that path + + to :func:`train_sft`.' + UnslothOptimizerSpec: + properties: + learning_rate: + type: number + exclusiveMinimum: 0.0 + title: Learning Rate + default: 0.0002 + weight_decay: + type: number + minimum: 0.0 + title: Weight Decay + default: 0.0 + optim: + type: string + enum: + - adamw_torch + - adamw_torch_fused + - adamw_8bit + - paged_adamw_8bit + - sgd + title: Optim + default: adamw_8bit + adam_beta1: + type: number + exclusiveMaximum: 1.0 + minimum: 0.0 + title: Adam Beta1 + description: Adam/AdamW beta1. + default: 0.9 + adam_beta2: + type: number + exclusiveMaximum: 1.0 + minimum: 0.0 + title: Adam Beta2 + description: Adam/AdamW beta2. + default: 0.999 + adam_epsilon: + type: number + exclusiveMinimum: 0.0 + title: Adam Epsilon + description: Adam/AdamW epsilon for numerical stability. + default: 1.0e-08 + max_grad_norm: + type: number + minimum: 0.0 + title: Max Grad Norm + description: Gradient-clipping max norm (TRL default 1.0). + default: 1.0 + label_smoothing_factor: + type: number + exclusiveMaximum: 1.0 + minimum: 0.0 + title: Label Smoothing Factor + description: Label smoothing for the cross-entropy loss. 0.0 disables. + default: 0.0 + neftune_noise_alpha: + title: Neftune Noise Alpha + description: NEFTune embedding-noise alpha (quality boost). None disables. + type: number + minimum: 0.0 + additionalProperties: false + type: object + title: UnslothOptimizerSpec + UnslothOutputRequest: + properties: + name: + title: Name + type: string + description: + title: Description + type: string + save_method: + type: string + enum: + - lora + - merged_16bit + - merged_4bit + title: Save Method + default: lora + additionalProperties: false + type: object + title: UnslothOutputRequest + description: Submitter-facing output preferences. ``name`` is auto-derived if + omitted. + UnslothOutputResponse: + properties: + name: + type: string + title: Name + type: + type: string + enum: + - adapter + - model + title: Type + save_method: + type: string + enum: + - lora + - merged_16bit + - merged_4bit + title: Save Method + fileset: + type: string + title: Fileset + description: Platform fileset name the trained checkpoint will be uploaded + to. Defaults to the entity name. + description: + title: Description + type: string + additionalProperties: false + type: object + required: + - name + - type + - save_method + - fileset + title: UnslothOutputResponse + description: 'Stored on the canonical UnslothJobOutput. Output naming is resolved + during ``to_spec``. + + + ``type`` is the high-level shape (``adapter`` for a saved LoRA, ``model`` + + for a merged checkpoint). ``save_method`` keeps the original Unsloth + + save verb so the training driver can dispatch correctly without + + re-deriving it. ``fileset`` is the platform fileset name the trained + + artefacts will be uploaded to (the plugin''s ``transform`` defaults + + this to ``name``).' + UnslothScheduleSpec: + properties: + epochs: + type: integer + exclusiveMinimum: 0.0 + title: Epochs + default: 1 + max_steps: + title: Max Steps + type: integer + exclusiveMinimum: 0.0 + warmup_steps: + type: integer + minimum: 0.0 + title: Warmup Steps + default: 0 + warmup_ratio: + title: Warmup Ratio + type: number + maximum: 1.0 + minimum: 0.0 + lr_scheduler_type: + type: string + enum: + - linear + - cosine + - constant + - constant_with_warmup + - cosine_with_restarts + title: Lr Scheduler Type + default: linear + logging_steps: + type: integer + exclusiveMinimum: 0.0 + title: Logging Steps + default: 1 + save_steps: + title: Save Steps + type: integer + exclusiveMinimum: 0.0 + eval_steps: + title: Eval Steps + type: integer + exclusiveMinimum: 0.0 + seed: + type: integer + title: Seed + default: 3407 + lr_scheduler_kwargs: + title: Lr Scheduler Kwargs + description: 'Extra kwargs for the LR scheduler, e.g. {''num_cycles'': 3} + for cosine_with_restarts. None uses scheduler defaults.' + additionalProperties: true + type: object + additionalProperties: false + type: object + title: UnslothScheduleSpec + description: Training schedule, scheduler, logging cadence. + UnslothToolCallParams: + properties: + tool_call_parser: + title: Tool Call Parser + description: Name of the tool call parser to use (e.g., 'openai', 'hermes', + 'pythonic', 'llama3_json', 'mistral'). + type: string + tool_call_plugin: + title: Tool Call Plugin + description: 'Reference to a fileset containing the custom tool call plugin + Python file. Expected format: ''{workspace}/{fileset_name}''.' + type: string + pattern: ^[\w\-.]+/[\w\-.]+$ + auto_tool_choice: + title: Auto Tool Choice + description: Whether to enable automatic tool choice. + type: boolean + additionalProperties: false + type: object + title: UnslothToolCallParams + description: Tool calling configuration for NIM deployments. + UnslothTrainingSpec: + properties: + training_type: + type: string + const: sft + title: Training Type + default: sft + finetuning_type: + type: string + enum: + - lora + - all_weights + title: Finetuning Type + default: lora + lora: + allOf: + - $ref: '#/components/schemas/UnslothLoRAParams' + description: Required when finetuning_type='lora'. Auto-filled with defaults + if omitted. + use_gradient_checkpointing: + type: string + enum: + - unsloth + - 'true' + - 'false' + title: Use Gradient Checkpointing + default: unsloth + additionalProperties: false + type: object + title: UnslothTrainingSpec + description: Algorithm + adapter shape selectors. ValidationError: properties: loc: diff --git a/plugins/nemo-customizer/pyproject.toml b/plugins/nemo-customizer/pyproject.toml index 70c23e582a..bc0f0aa564 100644 --- a/plugins/nemo-customizer/pyproject.toml +++ b/plugins/nemo-customizer/pyproject.toml @@ -33,6 +33,10 @@ build-backend = "hatchling.build" packages = ["src/nemo_customizer"] [tool.nemo.openapi] +# nemo-customizer merges every customization backend (automodel, unsloth, …) +# into one /apis/customization app, so enforce that no two backends emit the +# same OpenAPI schema name with differing content (see NamespacedModel). +strict_schema_collisions = true [tool.uv.sources] diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/schema.py b/plugins/nemo-rl/src/nemo_rl_plugin/schema.py index 2cc4fa07db..9221514c24 100644 --- a/plugins/nemo-rl/src/nemo_rl_plugin/schema.py +++ b/plugins/nemo-rl/src/nemo_rl_plugin/schema.py @@ -18,9 +18,10 @@ OutputResponse, ParallelismParams, RlJobOutput, + RlSchema, TrainingMethod, ) -from pydantic import BaseModel, ConfigDict, Field +from pydantic import ConfigDict, Field __all__ = [ "DPOTraining", @@ -33,18 +34,18 @@ ] -class OutputRequest(BaseModel): +class OutputRequest(RlSchema): """Submitter-facing output preferences. ``name`` is auto-derived if omitted.""" - model_config = ConfigDict(extra="forbid") - name: str | None = None -class RlJobInput(BaseModel): +class RlJobInput(RlSchema): """POST body / CLI JSON for ``nemo customization rl submit``.""" - model_config = ConfigDict(extra="forbid", protected_namespaces=()) + # extra="forbid" inherited from RlSchema; protected_namespaces=() kept for the + # ``model`` field. + model_config = ConfigDict(protected_namespaces=()) name: str | None = None model: str = Field(description="Model entity reference ('name' or 'workspace/name').") diff --git a/plugins/nemo-unsloth/src/nemo_unsloth_plugin/schema.py b/plugins/nemo-unsloth/src/nemo_unsloth_plugin/schema.py index 298d2e594d..594fce571d 100644 --- a/plugins/nemo-unsloth/src/nemo_unsloth_plugin/schema.py +++ b/plugins/nemo-unsloth/src/nemo_unsloth_plugin/schema.py @@ -37,8 +37,9 @@ ToolCallParams, TrainingSpec, UnslothJobOutput, + UnslothSchema, ) -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import Field, model_validator __all__ = [ "BatchSpec", @@ -58,21 +59,17 @@ ] -class OutputRequest(BaseModel): +class OutputRequest(UnslothSchema): """Submitter-facing output preferences. ``name`` is auto-derived if omitted.""" - model_config = ConfigDict(extra="forbid") - name: str | None = None description: str | None = None save_method: Literal["lora", "merged_16bit", "merged_4bit"] = "lora" -class UnslothJobInput(BaseModel): +class UnslothJobInput(UnslothSchema): """POST body / CLI JSON for ``nemo customization unsloth run``.""" - model_config = ConfigDict(extra="forbid") - name: str | None = None model: ModelLoadSpec dataset: DatasetSpec diff --git a/script/generate_openapi_spec.py b/script/generate_openapi_spec.py index d832843511..5b190d5184 100644 --- a/script/generate_openapi_spec.py +++ b/script/generate_openapi_spec.py @@ -466,7 +466,7 @@ def extract_plugin_specs_with_process_pool(plugins: List[PluginConfig]) -> None: print_green(f"All {len(plugins)} plugin(s) completed successfully!") -def apply_schema_fixes(spec_files: List[str], apply_reorder: bool = True) -> None: +def apply_schema_fixes(spec_files: List[str], apply_reorder: bool = True, strict_collisions: bool = False) -> None: """Apply schema fixes to a list of OpenAPI spec files.""" print_green("=== Applying fixes to OpenAPI schemas ===") # Endpoints stripped from the public OpenAPI spec (not exposed in SDK). @@ -517,7 +517,7 @@ def apply_schema_fixes(spec_files: List[str], apply_reorder: bool = True) -> Non spec = fix_openai_streaming_endpoints(spec) # Apply the standard fix-schema logic - spec = tweak_spec(spec) + spec = tweak_spec(spec, strict_collisions=strict_collisions) spec = hoist_nested_defs(spec) spec = remove_unused_schemas(spec) spec = remove_invalid_components(spec) @@ -747,7 +747,20 @@ def process_plugin_specs() -> None: # Reuse the platform schema-fix pipeline. Branches in apply_schema_fixes # gated on `"platform" in spec_file` are no-ops because plugin paths live # under plugins//openapi/ — no false positives by inspection. - apply_schema_fixes(plugin_spec_files) + # + # A plugin opts into strict_collisions via + # [tool.nemo.openapi].strict_schema_collisions when its spec merges multiple + # sub-apps (e.g. nemo-customizer, which mounts every customization backend + # under /apis/customization): there, two backends defining a same-named + # model with differing content is always a bug, so fail the build loudly + # instead of silently shipping a wrong contract. Plugin specs are + # independent files, so processing the two groups separately is safe. + strict_spec_files = [p.output_path() for p in plugins if p.strict_schema_collisions] + lenient_spec_files = [p.output_path() for p in plugins if not p.strict_schema_collisions] + if lenient_spec_files: + apply_schema_fixes(lenient_spec_files) + if strict_spec_files: + apply_schema_fixes(strict_spec_files, strict_collisions=True) fix_ref_not_allowed_errors(plugin_spec_files) validate_final_specs(plugin_spec_files) diff --git a/script/openapi_helper/plugin_config.py b/script/openapi_helper/plugin_config.py index 79bc312412..e4057c093c 100644 --- a/script/openapi_helper/plugin_config.py +++ b/script/openapi_helper/plugin_config.py @@ -22,6 +22,12 @@ class PluginConfig: env_vars: Optional[Dict[str, str]] = None factory_override: Optional[str] = None # "module:callable" escape hatch data_designer_plugin_allowlist: Optional[List[str]] = None + # Opt in when the plugin's spec merges multiple sub-apps into one (e.g. + # nemo-customizer mounts every customization backend under + # /apis/customization). Then two backends defining a same-named model with + # differing content is always a real bug, and spec generation should fail + # loudly instead of silently collapsing them. Off => warn-and-collapse. + strict_schema_collisions: bool = False @classmethod def from_pyproject(cls, pyproject_path: Path) -> Optional["PluginConfig"]: @@ -49,6 +55,9 @@ def from_pyproject(cls, pyproject_path: Path) -> Optional["PluginConfig"]: raise ValueError( f"plugin '{plugin_dir}': [tool.nemo.openapi].data_designer_plugin_allowlist must be a list of strings" ) + strict_schema_collisions = opts.get("strict_schema_collisions", False) + if not isinstance(strict_schema_collisions, bool): + raise ValueError(f"plugin '{plugin_dir}': [tool.nemo.openapi].strict_schema_collisions must be a boolean") return cls( dir=plugin_dir, @@ -56,6 +65,7 @@ def from_pyproject(cls, pyproject_path: Path) -> Optional["PluginConfig"]: env_vars=env_vars, factory_override=opts.get("factory_override"), data_designer_plugin_allowlist=data_designer_plugin_allowlist, + strict_schema_collisions=strict_schema_collisions, ) def resolve_service_name(self) -> str: diff --git a/services/core/models/src/nmp/core/models/schemas.py b/services/core/models/src/nmp/core/models/schemas.py index b3ba68a8f9..d1edb28202 100644 --- a/services/core/models/src/nmp/core/models/schemas.py +++ b/services/core/models/src/nmp/core/models/schemas.py @@ -8,6 +8,7 @@ from jinja2 import Environment from jinja2 import nodes as jinja_nodes +from nemo_platform_plugin.inference_middleware import BackendFormat from nmp.common.auth import AuthContext from nmp.common.entities import Filter, constants from nmp.common.entities.utils import get_random_id @@ -85,13 +86,6 @@ class FinetuningType(str, Enum): GRPO = "grpo" -class BackendFormat(str, Enum): - """Inference backend API wire formats.""" - - OPENAI_CHAT = "OPENAI_CHAT" - ANTHROPIC_MESSAGES = "ANTHROPIC_MESSAGES" - - class MoEConfig(BaseModel): """Mixture of Experts configuration.""" diff --git a/services/rl/src/nmp/rl/schemas.py b/services/rl/src/nmp/rl/schemas.py index 348bde37d2..1499ee6e2c 100644 --- a/services/rl/src/nmp/rl/schemas.py +++ b/services/rl/src/nmp/rl/schemas.py @@ -18,12 +18,32 @@ from typing import Literal, Self from nemo_platform_plugin.integrations import IntegrationsSpec +from nmp.customization_common.schema import NamespacedModel from nmp.customization_common.schemas.values import OutputNameType from nmp.rl.app.jobs.training.schemas import OptimizerType -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import ConfigDict, Field, model_validator -class ParallelismParams(BaseModel): +class RlSchema(NamespacedModel): + """Backend base: every RL-owned model emits an ``Rl``-prefixed OpenAPI schema + name (``DPOTraining`` -> ``RlDPOTraining``), so it can't collide with another + backend's same-named model in the merged ``/apis/customization`` spec. + ``extra='forbid'`` is inherited from the base. Defined here (with the + canonical schemas) and re-imported by the plugin's submitter-facing module so + both prefix identically. + + Inheriting the base also applies ``extra='forbid'`` uniformly across *all* RL + models, including the canonical/output ones (``RlJobOutput`` and + ``OutputResponse``) that previously used pydantic's default ``extra='ignore'``. + This tightening is intentional: it matches the automodel and unsloth backends, + whose ``*JobOutput`` models already forbid extras and are rehydrated via the + same ``model_validate(spec.model_dump())`` path (see + ``nemo_rl_plugin.jobs.jobs.compile``) without carrying foreign fields.""" + + __schema_namespace__ = "Rl" + + +class ParallelismParams(RlSchema): """Distributed training parallelism configuration. Single-node multi-GPU uses ``num_nodes=1`` with ``num_gpus_per_node>1``; @@ -31,8 +51,6 @@ class ParallelismParams(BaseModel): executor (see :mod:`nmp.rl.app.jobs.compiler`). """ - model_config = ConfigDict(extra="forbid") - num_gpus_per_node: int = Field(default=1, gt=0, description="Number of GPUs per node.") num_nodes: int = Field(default=1, gt=0, description="Number of nodes (>1 → multi-node Ray cluster).") tensor_parallel_size: int = Field(default=1, gt=0, description="Tensor parallel size.") @@ -41,14 +59,16 @@ class ParallelismParams(BaseModel): sequence_parallel: bool = Field(default=False, description="Enable sequence parallelism.") -class _TrainingBase(BaseModel): +class _TrainingBase(RlSchema): """Common training configuration shared by all RL methods. Flat hyperparameters match the ML-practitioner mental model (HuggingFace ``TrainingArguments`` / TRL configs). Only parallelism is grouped. """ - model_config = ConfigDict(protected_namespaces=(), extra="forbid") + # extra="forbid" inherited from RlSchema; protected_namespaces=() kept so the + # ``model``-adjacent fields don't trip pydantic's protected-namespace guard. + model_config = ConfigDict(protected_namespaces=()) # --- Optimizer --- optimizer_type: OptimizerType | None = Field( @@ -134,7 +154,7 @@ class DPOTraining(_TrainingBase): TrainingMethod = DPOTraining -class _OutputBase(BaseModel): +class _OutputBase(RlSchema): name: str = Field( max_length=255, description="Name of the output artifact. Used to identify it during deployment and inference.", @@ -159,7 +179,7 @@ class OutputResponse(_OutputBase): ) -class RlJobOutput(BaseModel): +class RlJobOutput(RlSchema): """Canonical NeMo-RL job spec (output of the plugin transform). The ``dataset`` fileset must contain ``training.jsonl`` and ``validation.jsonl`` diff --git a/services/unsloth/src/nmp/unsloth/schemas.py b/services/unsloth/src/nmp/unsloth/schemas.py index 4d00ca5f97..2c60bd2ab4 100644 --- a/services/unsloth/src/nmp/unsloth/schemas.py +++ b/services/unsloth/src/nmp/unsloth/schemas.py @@ -32,10 +32,23 @@ from typing import Any, Literal, Self from nemo_platform_plugin.integrations import IntegrationsSpec -from pydantic import BaseModel, ConfigDict, Field, model_validator +from nmp.customization_common.schema import NamespacedModel +from pydantic import Field, model_validator -class ModelLoadSpec(BaseModel): +class UnslothSchema(NamespacedModel): + """Backend base: every Unsloth-owned model emits an ``Unsloth``-prefixed + OpenAPI schema name (``TrainingSpec`` -> ``UnslothTrainingSpec``), so it + can't collide with another backend's same-named model in the merged + ``/apis/customization`` spec. ``extra='forbid'`` is inherited from the base. + + Defined here (with the canonical schemas) and re-imported by the plugin's + submitter-facing module so both prefix identically.""" + + __schema_namespace__ = "Unsloth" + + +class ModelLoadSpec(UnslothSchema): """Args to ``FastLanguageModel.from_pretrained``. ``name`` is a NeMo Platform model entity reference (``"name"`` or @@ -44,8 +57,6 @@ class ModelLoadSpec(BaseModel): to :func:`train_sft`. """ - model_config = ConfigDict(extra="forbid") - name: str = Field( description=( "Model entity reference. Accepts 'name' (uses the job's workspace) " @@ -84,11 +95,9 @@ class ModelLoadSpec(BaseModel): ) -class LoRAParams(BaseModel): +class LoRAParams(UnslothSchema): """Args to ``FastLanguageModel.get_peft_model``.""" - model_config = ConfigDict(extra="forbid") - rank: int = Field(default=16, gt=0, description="LoRA rank.") alpha: int = Field(default=16, gt=0, description="LoRA scaling factor (alpha).") dropout: float = Field(default=0.0, ge=0.0, lt=1.0) @@ -136,11 +145,9 @@ class LoRAParams(BaseModel): ) -class TrainingSpec(BaseModel): +class TrainingSpec(UnslothSchema): """Algorithm + adapter shape selectors.""" - model_config = ConfigDict(extra="forbid") - training_type: Literal["sft"] = "sft" finetuning_type: Literal["lora", "all_weights"] = "lora" lora: LoRAParams | None = Field( @@ -166,7 +173,7 @@ def _enforce_lora_invariant(self) -> Self: return self -class DatasetSpec(BaseModel): +class DatasetSpec(UnslothSchema): """Training data location + shape. ``path`` and ``validation_path`` are platform fileset references @@ -175,8 +182,6 @@ class DatasetSpec(BaseModel): filesystem path. """ - model_config = ConfigDict(extra="forbid") - path: str = Field( description=( "Training fileset reference: 'name' (uses the job's workspace) " @@ -199,11 +204,9 @@ class DatasetSpec(BaseModel): packing: bool = Field(default=False, description="trl.SFTTrainer packing flag.") -class ScheduleSpec(BaseModel): +class ScheduleSpec(UnslothSchema): """Training schedule, scheduler, logging cadence.""" - model_config = ConfigDict(extra="forbid") - # Consistent with Automodel: train for ``epochs`` (default 1) unless ``max_steps`` # is set, in which case the trainer caps training at that many steps. epochs: int = Field(default=1, gt=0) @@ -230,16 +233,12 @@ class ScheduleSpec(BaseModel): ) -class BatchSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class BatchSpec(UnslothSchema): per_device_train_batch_size: int = Field(default=1, gt=0) gradient_accumulation_steps: int = Field(default=1, gt=0) -class OptimizerSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class OptimizerSpec(UnslothSchema): learning_rate: float = Field(default=2e-4, gt=0.0) weight_decay: float = Field(default=0.0, ge=0.0) # Unsloth's notebooks default to adamw_8bit (bitsandbytes-backed) — much @@ -266,9 +265,7 @@ class OptimizerSpec(BaseModel): ) -class HardwareSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - +class HardwareSpec(UnslothSchema): gpus: str | None = Field( default=None, description=( @@ -281,7 +278,7 @@ class HardwareSpec(BaseModel): ) -class OutputResponse(BaseModel): +class OutputResponse(UnslothSchema): """Stored on the canonical UnslothJobOutput. Output naming is resolved during ``to_spec``. ``type`` is the high-level shape (``adapter`` for a saved LoRA, ``model`` @@ -292,8 +289,6 @@ class OutputResponse(BaseModel): this to ``name``). """ - model_config = ConfigDict(extra="forbid") - name: str type: Literal["adapter", "model"] save_method: Literal["lora", "merged_16bit", "merged_4bit"] @@ -303,11 +298,9 @@ class OutputResponse(BaseModel): description: str | None = None -class ToolCallParams(BaseModel): +class ToolCallParams(UnslothSchema): """Tool calling configuration for NIM deployments.""" - model_config = ConfigDict(extra="forbid") - tool_call_parser: str | None = Field( default=None, description=( @@ -328,15 +321,13 @@ class ToolCallParams(BaseModel): ) -class DeploymentParams(BaseModel): +class DeploymentParams(UnslothSchema): """Inline deployment parameters for auto-deploying a trained model. Used in :class:`UnslothJobInput.deployment_config` and passed through to the model_entity task at compile time. When unset, no deployment is launched. """ - model_config = ConfigDict(extra="forbid") - gpu: int = Field(default=1, description="Number of GPUs required for the deployment.") additional_envs: dict[str, str] | None = Field( default=None, @@ -361,7 +352,7 @@ class DeploymentParams(BaseModel): ) -class UnslothJobOutput(BaseModel): +class UnslothJobOutput(UnslothSchema): """Canonical spec stored after the plugin's ``to_spec()`` resolves output naming. Defaults match :class:`~nemo_unsloth_plugin.schema.UnslothJobInput` so SDK @@ -370,8 +361,6 @@ class UnslothJobOutput(BaseModel): resolved input values through, so these defaults never override real input. """ - model_config = ConfigDict(extra="forbid") - name: str | None = None model: ModelLoadSpec dataset: DatasetSpec