Merge remote-tracking branch 'origin/nightly' into feature/canvas-lab

# Conflicts: # studio/frontend/bun.lock # studio/frontend/package.json
2026-02-24 09:02:54 +00:00 · 2026-02-24 09:02:54 +00:00 · c0f6012d77
commit c0f6012d77
parent a4b7d360de 8aca1cad29
11 changed files with 337 additions and 106 deletions
--- a/datasets/recipes/recipe_a28d33f9eab14660bba2437a6dfda290/builder_config.json
+++ b/datasets/recipes/recipe_a28d33f9eab14660bba2437a6dfda290/builder_config.json
@ -0,0 +1,72 @@
+{
+  "data_designer": {
+    "columns": [
+      {
+        "name": "generated_instruction",
+        "drop": false,
+        "allow_resize": false,
+        "column_type": "llm-text",
+        "prompt": "Based on this target answer:\n{{ output }}\n\nWrite one high-quality plain text short and brief user instruction that this answer would satisfy.\nReturn only the instruction.",
+        "model_alias": "ministral",
+        "system_prompt": null,
+        "multi_modal_context": null,
+        "tool_alias": null,
+        "with_trace": "none",
+        "extract_reasoning_content": false
+      },
+      {
+        "name": "instruction",
+        "drop": false,
+        "allow_resize": false,
+        "column_type": "seed-dataset"
+      },
+      {
+        "name": "input",
+        "drop": false,
+        "allow_resize": false,
+        "column_type": "seed-dataset"
+      },
+      {
+        "name": "output",
+        "drop": false,
+        "allow_resize": false,
+        "column_type": "seed-dataset"
+      }
+    ],
+    "model_configs": [
+      {
+        "alias": "ministral",
+        "model": "mistralai/ministral-8b-2512",
+        "inference_parameters": {
+          "generation_type": "chat-completion",
+          "max_parallel_requests": 4,
+          "timeout": null,
+          "extra_body": null,
+          "temperature": 0.7,
+          "top_p": null,
+          "max_tokens": 1024
+        },
+        "provider": "openai_provider",
+        "skip_health_check": false
+      }
+    ],
+    "tool_configs": [],
+    "seed_config": {
+      "source": {
+        "seed_type": "hf",
+        "path": "datasets/unsloth/alpaca-cleaned/**/*.json",
+        "token": null,
+        "endpoint": "https://huggingface.co"
+      },
+      "sampling_strategy": "ordered",
+      "selection_strategy": {
+        "start": 1,
+        "end": 100
+      }
+    },
+    "constraints": null,
+    "profilers": null,
+    "processors": null
+  },
+  "library_version": "0.5.1"
+}
--- a/datasets/recipes/recipe_a28d33f9eab14660bba2437a6dfda290/metadata.json
+++ b/datasets/recipes/recipe_a28d33f9eab14660bba2437a6dfda290/metadata.json
@ -0,0 +1,36 @@
+{
+  "actual_num_records": 50,
+  "buffer_size": 50,
+  "column_statistics": [
+    {
+      "column_name": "generated_instruction",
+      "column_type": "llm-text",
+      "input_tokens_mean": 170.52,
+      "input_tokens_median": 132.5,
+      "input_tokens_stddev": 119.91,
+      "num_null": 0,
+      "num_records": 50,
+      "num_unique": 50,
+      "output_tokens_mean": 39.74,
+      "output_tokens_median": 31.0,
+      "output_tokens_stddev": 39.58,
+      "pyarrow_dtype": "string",
+      "simple_dtype": "string"
+    }
+  ],
+  "dataset_name": "recipe_a28d33f9eab14660bba2437a6dfda290",
+  "file_paths": {
+    "parquet-files": [
+      "parquet-files/batch_00000.parquet"
+    ]
+  },
+  "num_completed_batches": 1,
+  "schema": {
+    "generated_instruction": "string",
+    "input": "string",
+    "instruction": "string",
+    "output": "string"
+  },
+  "target_num_records": 50,
+  "total_num_batches": 1
+}
--- a/datasets/recipes/recipe_a28d33f9eab14660bba2437a6dfda290/parquet-files/batch_00000.parquet
+++ b/datasets/recipes/recipe_a28d33f9eab14660bba2437a6dfda290/parquet-files/batch_00000.parquet
--- a/studio/frontend/bun.lock
+++ b/studio/frontend/bun.lock
@ -40,9 +40,16 @@
        "clsx": "^2.1.1",
        "cmdk": "^1.1.1",
        "date-fns": "^4.1.0",
+<<<<<<< HEAD
        "dexie": "^4.3.0",
        "framer-motion": "^11.18.2",
        "katex": "^0.16.28",
+=======
+        "dexie": "^4.2.1",
+        "framer-motion": "^11.15.0",
+        "js-yaml": "^4.1.1",
+        "katex": "^0.16.22",
+>>>>>>> origin/nightly
        "lucide-react": "^0.563.0",
        "mammoth": "^1.11.0",
        "motion": "^12.34.0",
@ -68,6 +75,7 @@
      "devDependencies": {
        "@biomejs/biome": "^1.9.4",
        "@eslint/js": "^9.39.1",
+        "@types/js-yaml": "^4.0.9",
        "@types/node": "^24.10.1",
        "@types/react": "^19.2.5",
        "@types/react-dom": "^19.2.3",
@ -767,6 +775,8 @@

    "@types/hast": ["@types/hast@3.0.4", "", { "dependencies": { "@types/unist": "*" } }, "sha512-WPs+bbQw5aCj+x6laNGWLH3wviHtoCv/P3+otBhbOhJgG8qtpdAMlTCxLtsTWA7LH1Oh/bFCHsBn0TPS5m30EQ=="],

+    "@types/js-yaml": ["@types/js-yaml@4.0.9", "", {}, "sha512-k4MGaQl5TGo/iipqb2UDG2UwjXziSWkh0uysQelTlJpX1qGlpUZYm8PnO4DxG1qBomtJUdYJ6qR6xdIah10JLg=="],
+
    "@types/json-schema": ["@types/json-schema@7.0.15", "", {}, "sha512-5+fP8P8MFNC+AyZCDxrB2pkZFPGzqQWUzpSeuuVLvm8VMcorNYavBqoFcxK8bQz4Qsbn4oUEEem4wDLfcysGHA=="],

    "@types/katex": ["@types/katex@0.16.8", "", {}, "sha512-trgaNyfU+Xh2Tc+ABIb44a5AYUpicB3uwirOioeOkNPPbmgRNtcWyDeeFRzjPZENO9Vq8gvVqfhaaXWLlevVwg=="],
@ -841,7 +851,7 @@

    "ansis": ["ansis@4.2.0", "", {}, "sha512-HqZ5rWlFjGiV0tDm3UxxgNRqsOTniqoKZu0pIAfh7TZQMGuZK+hH0drySty0si0QXj1ieop4+SkSfPZBPPkHig=="],

-    "argparse": ["argparse@1.0.10", "", { "dependencies": { "sprintf-js": "~1.0.2" } }, "sha512-o5Roy6tNG4SL/FOkCAN6RzjiakZS25RLYFrcMttJqbdd8BWrnA+fGz57iN5Pb06pvBGvl5gQ0B48dJlslXvoTg=="],
+    "argparse": ["argparse@2.0.1", "", {}, "sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q=="],

    "aria-hidden": ["aria-hidden@1.2.6", "", { "dependencies": { "tslib": "^2.0.0" } }, "sha512-ik3ZgC9dY/lYVVM++OISsaYDeg1tb0VtP5uL3ouh1koGOaUMDPpbFIei4JkFimWUFPn90sbMNMXQAIVOlnYKJA=="],

@ -2301,16 +2311,21 @@

    "fast-glob/glob-parent": ["glob-parent@5.1.2", "", { "dependencies": { "is-glob": "^4.0.1" } }, "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow=="],

+<<<<<<< HEAD
    "js-yaml/argparse": ["argparse@2.0.1", "", {}, "sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q=="],

    "langsmith/chalk": ["chalk@5.6.2", "", {}, "sha512-7NzBL0rN6fMUW+f7A6Io4h40qQlG+xGmtMxfbnH/K7TAtt8JQWVQK+6g0UXKMeVJoyV5EkkNsErQ8pVD3bLHbA=="],

    "langsmith/semver": ["semver@7.7.4", "", { "bin": { "semver": "bin/semver.js" } }, "sha512-vFKC2IEtQnVhpT78h1Yp8wzwrf8CM+MzKMHGJZfBtzhZNycRFnXsHk6E5TxIkkMsgNS7mdX3AGB7x2QM2di4lA=="],

+=======
+>>>>>>> origin/nightly
    "log-symbols/chalk": ["chalk@5.6.2", "", {}, "sha512-7NzBL0rN6fMUW+f7A6Io4h40qQlG+xGmtMxfbnH/K7TAtt8JQWVQK+6g0UXKMeVJoyV5EkkNsErQ8pVD3bLHbA=="],

    "log-symbols/is-unicode-supported": ["is-unicode-supported@1.3.0", "", {}, "sha512-43r2mRvz+8JRIKnWJ+3j8JtjRKZ6GmjzfaE/qiBJnikNnYv/6bagRJ1kUhNk8R5EX/GkobD+r+sfxCPJsiKBLQ=="],

+    "mammoth/argparse": ["argparse@1.0.10", "", { "dependencies": { "sprintf-js": "~1.0.2" } }, "sha512-o5Roy6tNG4SL/FOkCAN6RzjiakZS25RLYFrcMttJqbdd8BWrnA+fGz57iN5Pb06pvBGvl5gQ0B48dJlslXvoTg=="],
+
    "mdast-util-find-and-replace/escape-string-regexp": ["escape-string-regexp@5.0.0", "", {}, "sha512-/veY75JbMK4j1yjvuUxuVsiS/hr/4iHs9FTT6cgTexxdE0Ly/glccBAkloH/DofkjRbZU3bnoj38mOmhkZ0lHw=="],

    "mermaid/marked": ["marked@16.4.2", "", { "bin": { "marked": "bin/marked.js" } }, "sha512-TI3V8YYWvkVf3KJe1dRkpnjs68JUPyEa5vjKrp1XEEJUAOaQc+Qj+L1qWbPd0SJuAdQkFU0h73sXXqwDYxsiDA=="],
--- a/studio/frontend/package.json
+++ b/studio/frontend/package.json
@ -50,6 +50,7 @@
    "date-fns": "^4.1.0",
    "dexie": "^4.3.0",
    "framer-motion": "^11.18.2",
+    "js-yaml": "^4.1.1",
    "katex": "^0.16.28",
    "lucide-react": "^0.563.0",
    "mammoth": "^1.11.0",
@ -76,6 +77,7 @@
  "devDependencies": {
    "@biomejs/biome": "^1.9.4",
    "@eslint/js": "^9.39.1",
+    "@types/js-yaml": "^4.0.9",
    "@types/node": "^24.10.1",
    "@types/react": "^19.2.5",
    "@types/react-dom": "^19.2.3",
@ -88,4 +90,4 @@
    "typescript-eslint": "^8.55.0",
    "vite": "^7.3.1"
  }
-}
+}
--- a/studio/frontend/src/features/studio/sections/training-section.tsx
+++ b/studio/frontend/src/features/studio/sections/training-section.tsx
@ -2,23 +2,27 @@ import { SectionCard } from "@/components/section-card";
 import { Button } from "@/components/ui/button";
 import { ChartContainer } from "@/components/ui/chart";
 import type { ChartConfig } from "@/components/ui/chart";
-import { Checkbox } from "@/components/ui/checkbox";
 import {
-  Collapsible,
-  CollapsibleContent,
-  CollapsibleTrigger,
-} from "@/components/ui/collapsible";
-import { Input } from "@/components/ui/input";
-import { useTrainingActions, useTrainingConfigStore } from "@/features/training";
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import {
+  parseYamlConfig,
+  serializeConfigToYaml,
+  useTrainingActions,
+  useTrainingConfigStore,
+} from "@/features/training";
 import {
  Archive04Icon,
-  ArrowDown01Icon,
  ChartAverageIcon,
  CleanIcon,
+  CloudUploadIcon,
  Rocket01Icon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { useState } from "react";
+import { useRef } from "react";
+import { toast } from "sonner";
 import { CartesianGrid, Line, LineChart, XAxis, YAxis } from "recharts";

 const chartConfig = {
@ -37,10 +41,55 @@ const placeholderData = [
 export function TrainingSection() {
  const store = useTrainingConfigStore();
  const { isStarting, startError, startTrainingRun } = useTrainingActions();
-  const [logOpen, setLogOpen] = useState(false);
  const isIncompatible =
    !store.isVisionModel && store.isDatasetMultimodal === true;
+  const fileInputRef = useRef<HTMLInputElement>(null);

+  const handleFileUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
+    const file = e.target.files?.[0];
+    if (!file) return;
+    e.target.value = "";
+
+    const reader = new FileReader();
+    reader.onload = () => {
+      try {
+        const config = parseYamlConfig(reader.result as string);
+        store.applyConfigPatch(config);
+        toast.success("Config loaded", { description: file.name });
+      } catch (err) {
+        toast.error("Failed to load config", {
+          description:
+            err instanceof Error ? err.message : "Invalid YAML file",
+        });
+      }
+    };
+    reader.onerror = () => {
+      toast.error("Failed to read file");
+    };
+    reader.readAsText(file);
+  };
+
+  const handleSaveConfig = () => {
+    const yamlStr = serializeConfigToYaml(store, store.isVisionModel);
+    const blob = new Blob([yamlStr], { type: "text/yaml" });
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement("a");
+    a.href = url;
+
+    const model = (store.selectedModel ?? "model").split("/").pop();
+    const method = store.trainingMethod ?? "qlora";
+    const dataset = (store.dataset ?? "dataset").split("/").pop();
+    const timestamp = new Date().toISOString().replace(/[:T]/g, "-").slice(0, 19);
+    a.download = `${model}_${method}_${dataset}_${timestamp}.yaml`;
+
+    a.click();
+    URL.revokeObjectURL(url);
+  };
+
+  const handleResetConfig = () => {
+    store.resetToModelDefaults();
+    toast.success("Parameters reset to model defaults");
+  };

  return (
    <div data-tour="studio-training" className="col-span-1 xl:col-span-4">
@ -115,100 +164,61 @@ export function TrainingSection() {
          </p>
        )}

-        {/* Save / Clear */}
-        <div className="grid grid-cols-2 gap-2">
-          <Button
-            data-tour="studio-save"
-            variant="outline"
-            size="sm"
-            className="cursor-pointer"
-          >
-            <HugeiconsIcon icon={Archive04Icon} className="size-3.5" /> Save
-            Config
-          </Button>
-          <Button variant="outline" size="sm" className="cursor-pointer">
-            <HugeiconsIcon icon={CleanIcon} className="size-3.5" /> Clear
-          </Button>
+        {/* Upload / Save / Reset */}
+        <p className="text-xs text-muted-foreground">Training Config</p>
+        <div className="grid grid-cols-3 gap-2">
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                variant="outline"
+                size="sm"
+                className="cursor-pointer"
+                onClick={() => fileInputRef.current?.click()}
+              >
+                <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
+                Upload
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Load a saved YAML config</TooltipContent>
+          </Tooltip>
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                data-tour="studio-save"
+                variant="outline"
+                size="sm"
+                className="cursor-pointer"
+                onClick={handleSaveConfig}
+              >
+                <HugeiconsIcon icon={Archive04Icon} className="size-3.5" />
+                Save
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Download current config as YAML</TooltipContent>
+          </Tooltip>
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                variant="outline"
+                size="sm"
+                className="cursor-pointer"
+                onClick={handleResetConfig}
+                disabled={!store.selectedModel}
+              >
+                <HugeiconsIcon icon={CleanIcon} className="size-3.5" />
+                Reset
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Reset to model defaults</TooltipContent>
+          </Tooltip>
        </div>
-
-        {/* Logging */}
-        <Collapsible open={logOpen} onOpenChange={setLogOpen}>
-          <CollapsibleTrigger className="flex w-full cursor-pointer items-center gap-1.5 text-xs text-muted-foreground">
-            <HugeiconsIcon
-              icon={ArrowDown01Icon}
-              className={`size-3.5 transition-transform ${logOpen ? "rotate-180" : ""}`}
-            />
-            Logging
-          </CollapsibleTrigger>
-          <CollapsibleContent className="mt-3 flex flex-col gap-3">
-            {/* W&B */}
-            <div className="flex items-center gap-2">
-              <Checkbox
-                id="wandb"
-                checked={store.enableWandb}
-                onCheckedChange={(v) => store.setEnableWandb(!!v)}
-              />
-              <label
-                htmlFor="wandb"
-                className="text-xs cursor-pointer text-muted-foreground"
-              >
-                Weights & Biases
-              </label>
-            </div>
-            {store.enableWandb && (
-              <div className="flex flex-col gap-2 pl-6">
-                <Input
-                  placeholder="W&B API Token"
-                  type="password"
-                  value={store.wandbToken}
-                  onChange={(e) => store.setWandbToken(e.target.value)}
-                />
-                <Input
-                  placeholder="Project name"
-                  value={store.wandbProject}
-                  onChange={(e) => store.setWandbProject(e.target.value)}
-                />
-              </div>
-            )}
-
-            {/* TensorBoard */}
-            <div className="flex items-center gap-2">
-              <Checkbox
-                id="tensorboard"
-                checked={store.enableTensorboard}
-                onCheckedChange={(v) => store.setEnableTensorboard(!!v)}
-              />
-              <label
-                htmlFor="tensorboard"
-                className="text-xs cursor-pointer text-muted-foreground"
-              >
-                TensorBoard
-              </label>
-            </div>
-            {store.enableTensorboard && (
-              <div className="flex flex-col gap-2 pl-6">
-                <Input
-                  placeholder="Log directory"
-                  value={store.tensorboardDir}
-                  onChange={(e) => store.setTensorboardDir(e.target.value)}
-                />
-                <div className="flex items-center justify-between">
-                  <span className="text-xs font-medium text-muted-foreground">
-                    Log frequency
-                  </span>
-                  <Input
-                    type="number"
-                    value={store.logFrequency}
-                    onChange={(e) =>
-                      store.setLogFrequency(Number(e.target.value))
-                    }
-                    className="w-24"
-                  />
-                </div>
-              </div>
-            )}
-          </CollapsibleContent>
-        </Collapsible>
+        <input
+          ref={fileInputRef}
+          type="file"
+          accept=".yaml,.yml"
+          className="hidden"
+          onChange={handleFileUpload}
+        />
        </div>
      </SectionCard>
    </div>
--- a/studio/frontend/src/features/studio/tour/steps/save.tsx
+++ b/studio/frontend/src/features/studio/tour/steps/save.tsx
@ -6,8 +6,8 @@ export const studioSaveStep: TourStep = {
  title: "Save config",
  body: (
    <>
-      Save configs that worked. Re-running the same baseline makes it obvious
-      if a change helped (or if you just got lucky).
+      Save your training config as a YAML file. Re-running the same baseline
+      makes it obvious if a change helped (or if you just got lucky).
    </>
  ),
 };
--- a/studio/frontend/src/features/training/index.ts
+++ b/studio/frontend/src/features/training/index.ts
@ -10,3 +10,4 @@ export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-st
 export { listLocalModels } from "./api/models-api";
 export type { LocalModelInfo } from "./api/models-api";
 export type { TrainingPhase } from "./types/runtime";
+export { parseYamlConfig, serializeConfigToYaml } from "./lib/yaml-config";
--- a/studio/frontend/src/features/training/lib/yaml-config.ts
+++ b/studio/frontend/src/features/training/lib/yaml-config.ts
@ -0,0 +1,81 @@
+import * as yaml from "js-yaml";
+import type { BackendModelConfig } from "../api/models-api";
+import type { TrainingConfigState } from "../types/config";
+
+const EXPECTED_TOP_KEYS = new Set(["training", "lora", "logging", "inference"]);
+
+/**
+ * Parse a YAML string into a BackendModelConfig suitable for
+ * `mapBackendModelConfigToTrainingPatch`. Throws on invalid input.
+ */
+export function parseYamlConfig(text: string): BackendModelConfig {
+  const parsed = yaml.load(text);
+  if (parsed == null || typeof parsed !== "object" || Array.isArray(parsed)) {
+    throw new Error(
+      "Invalid config: expected a YAML mapping with training/lora/logging sections",
+    );
+  }
+
+  const raw = parsed as Record<string, unknown>;
+  const unknownKeys = Object.keys(raw).filter(
+    (k) => !EXPECTED_TOP_KEYS.has(k),
+  );
+  if (unknownKeys.length > 0) {
+    console.warn("Ignored unknown YAML keys:", unknownKeys.join(", "));
+  }
+
+  return {
+    training: (raw.training ?? undefined) as BackendModelConfig["training"],
+    lora: (raw.lora ?? undefined) as BackendModelConfig["lora"],
+    logging: (raw.logging ?? undefined) as BackendModelConfig["logging"],
+  };
+}
+
+/**
+ * Serialize the current training config state to a YAML string matching the
+ * backend model-defaults schema.
+ */
+export function serializeConfigToYaml(
+  state: TrainingConfigState,
+  includeVisionFields: boolean,
+): string {
+  const lora: Record<string, unknown> = {
+    lora_r: state.loraRank,
+    lora_alpha: state.loraAlpha,
+    lora_dropout: state.loraDropout,
+    target_modules: state.targetModules,
+    use_rslora: state.loraVariant === "rslora",
+    use_loftq: state.loraVariant === "loftq",
+  };
+
+  if (includeVisionFields) {
+    lora.finetune_vision_layers = state.finetuneVisionLayers;
+    lora.finetune_language_layers = state.finetuneLanguageLayers;
+    lora.finetune_attention_modules = state.finetuneAttentionModules;
+    lora.finetune_mlp_modules = state.finetuneMLPModules;
+  }
+
+  const config = {
+    training: {
+      max_seq_length: state.contextLength,
+      num_epochs: state.epochs,
+      learning_rate: state.learningRate,
+      batch_size: state.batchSize,
+      gradient_accumulation_steps: state.gradientAccumulation,
+      warmup_steps: state.warmupSteps,
+      max_steps: state.maxSteps,
+      save_steps: state.saveSteps,
+      eval_steps: state.evalSteps,
+      weight_decay: state.weightDecay,
+      random_seed: state.randomSeed,
+      packing: state.packing,
+      train_on_completions: state.trainOnCompletions,
+      gradient_checkpointing: state.gradientCheckpointing,
+      optim: state.optimizerType,
+      lr_scheduler_type: state.lrSchedulerType,
+    },
+    lora,
+  };
+
+  return yaml.dump(config, { lineWidth: -1, noRefs: true });
+}
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@ -5,6 +5,7 @@ import { persist } from "zustand/middleware";
 import { checkDatasetFormat } from "../api/datasets-api";
 import { checkVisionModel, getModelConfig } from "../api/models-api";
 import { mapBackendModelConfigToTrainingPatch } from "../lib/model-defaults";
+import type { BackendModelConfig } from "../api/models-api";
 import type { TrainingConfigState, TrainingConfigStore } from "../types/config";

 const MIN_STEP: StepNumber = 1;
@ -344,6 +345,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
        setTargetModules: (targetModules) => set({ targetModules }),
        canProceed: () => canProceedForStep(get()),
        reset: () => set(initialState),
+        resetToModelDefaults: () => {
+          const { selectedModel } = get();
+          if (!selectedModel) return;
+          set({ modelDefaultsAppliedFor: null });
+          loadAndApplyModelDefaults(selectedModel);
+        },
+        applyConfigPatch: (config: BackendModelConfig) => {
+          const patch = mapBackendModelConfigToTrainingPatch(config);
+          set(patch);
+        },
      };
    },
    {
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/studio/frontend/src/features/training/types/config.ts
@ -6,6 +6,7 @@ import type {
  StepNumber,
  TrainingMethod,
 } from "@/types/training";
+import type { BackendModelConfig } from "../api/models-api";

 export type LoraVariant = "lora" | "rslora" | "loftq";

@ -117,6 +118,8 @@ export interface TrainingConfigActions {
  setTargetModules: (value: string[]) => void;
  canProceed: () => boolean;
  reset: () => void;
+  resetToModelDefaults: () => void;
+  applyConfigPatch: (config: BackendModelConfig) => void;
 }

 export type TrainingConfigStore = TrainingConfigState & TrainingConfigActions;