CLI和web支持训练恢复功能

Author: Yesifan

darkit/core/trainer.py

@@ -332,23 +332,19 @@ def _auto_save_pretrained(self):
         """
         根据 Train 的相关参数，控制训练时的自动保存逻辑。
         """
-        max_step, current_step, save_step_interval = (
-            self.max_step,
-            self.current_step,
-            self.save_step_interval,
-        )
-
-        current_step_idx = current_step + 1
+        current_step_idx = self.current_step + 1
+        # 如果设置 save_step_interval 为 0，则不保存 checkpoint
         if self.save_step_interval > 0:
+            # 当当前步数（current_step_idx）为 max_step 或者是 save_step_interval 的倍数时保存模型
             if (
-                current_step_idx == max_step
-                or current_step_idx % save_step_interval == 0
+                current_step_idx == self.max_step
+                or current_step_idx % self.save_step_interval == 0
             ):
                 check_poinent = f"iter-{current_step_idx}-ckpt"
                 self.save_pretrained(check_poinent=check_poinent)
                 if self.save_directory:
                     print(
-                        f"Model saved epoch {current_step_idx}/{max_step} at {check_poinent}"
+                        f"Model saved epoch {current_step_idx}/{self.max_step} at {check_poinent}"
                     )
 
     def _auto_validate(self, val_dataloader):
@@ -407,7 +403,6 @@ def _save_model(self, checkpoint="complete"):
         if self.save_directory:
             save_path = self.save_directory / f"{checkpoint}.pth"
             current_step_idx = self.current_step + 1
-            print("current_step_idx", current_step_idx)
             save_dict = {
                 "model_class": self.model.__class__.__name__,
                 "state_dict": self.model,

darkit/core/utils/csv_logger.py

@@ -1,16 +1,17 @@
 import csv
 from pathlib import Path
-from typing import Any, Collection, Union
+from typing import Any, Collection
 from dataclasses import asdict
 
 
 class CSVLogger:
-    def __init__(self, filename: Union[str, Path], fieldnames: Collection[Any]):
+    def __init__(self, filename: Path, fieldnames: Collection[Any]):
         self.filename = filename
         self.fieldnames = fieldnames
-        self.file = open(filename, "w")
+        self.file = open(filename, "a")
         self.writer = csv.DictWriter(self.file, fieldnames=fieldnames)
-        self.writer.writeheader()
+        if self.file.tell() == 0:
+            self.writer.writeheader()
 
     def log(self, data):
         if hasattr(data, "__dataclass_fields__"):

darkit/core/web/messages/en.json

@@ -26,6 +26,7 @@
     "startPredict": "Start Predicting",
     "cxtLenTooltip": "Context length. The maximum sequence length that the model can consider when processing the input sequence.",
     "promptTooltip": "The prompt to use for generation.",
+    "resumeTooltip": "Resume training from the model checkpoint, Ex: mymodel or mymodel:epoch-1.",
     "modelVisualHead": "Select the model for which you want to view details.",
     "showDetails": "Show Details",
     "mDeleteModel": "Delete Model",

darkit/core/web/messages/zh.json

@@ -26,6 +26,7 @@
     "startPredict": "开始预测",
     "cxtLenTooltip": "上下文长度。模型在处理输入序列时所能考虑的最大序列长度",
     "promptTooltip": "输入的提示词，模型将根据提示词生成后续文本",
+    "resumeTooltip": "从指定模型检查点恢复训练，Ex: mymodel or mymodel:epoch-1",
     "modelVisualHead": "选择您想要查看其详细信息的模型。",
     "showDetails": "查看详情",
     "mDeleteModel": "删除模型",

darkit/core/web/src/lib/apis/lm.ts

@@ -1,4 +1,4 @@
-import { get } from '$lib/api';
+import { get, post } from '$lib/api';
 
 export type ModelLabel = {
 	default: string | number;
@@ -22,6 +22,7 @@ export interface TrainedModel {
 	config: ExternalConfig;
 }
 
+/** 获取以及训练完成的模型 */
 export const getTrainedNetworks = async () => {
 	const res = await get('/lm/models');
 	return res.json() as Promise<TrainedModel[]>;
@@ -36,3 +37,69 @@ export const getModelDetail = async (model: string) => {
 	const res = await get(`/lm/models/${model}`);
 	return res.json() as Promise<TrainedModel>;
 };
+
+export const getModelsOptions = async () => {
+	const res = await get('/lm/models/options');
+	return res.json() as Promise<{ [key: string]: { model: any; trainer: any } }>;
+};
+
+export const getTrainResources = async () => {
+	const res = await get('/lm/train/resources');
+	return res.json() as Promise<string[][]>;
+};
+
+export const getTrainCommand = async (
+	type: string,
+	fork: string | null,
+	resume: string | null,
+	dataset: string,
+	tokenizer: string,
+	modelOption: any,
+	trainerOption: any
+) => {
+	const res = await post(`/lm/model/train/command`, {
+		type: type,
+		fork: fork,
+		resume: resume,
+		dataset: dataset,
+		tokenizer: tokenizer,
+		m_conf: modelOption,
+		t_conf: trainerOption
+	});
+	return res.text() as Promise<string>;
+};
+
+export const startTrainModel = async (
+	type: string,
+	fork: string | null,
+	resume: string | null,
+	dataset: string,
+	tokenizer: string,
+	modelOption: any,
+	trainerOption: any
+) => {
+	const res = await post('/lm/v2/model/train/', {
+		type: type,
+		fork: fork,
+		resume: resume,
+		dataset: dataset,
+		tokenizer: tokenizer,
+		m_conf: modelOption,
+		t_conf: trainerOption
+	});
+	return res.json() as Promise<any>;
+};
+
+export const createForkNetwork = async (
+	name: string,
+	modelType: string,
+	modelOption: any,
+	trainerOption: any
+) => {
+	const res = await post(`/lm/edit/init/${name}`, {
+		model: modelType,
+		m_conf: modelOption,
+		t_conf: trainerOption
+	});
+	return res.json() as Promise<any>;
+};

darkit/core/web/src/routes/lm/train/+page.svelte

@@ -1,111 +1,75 @@
 <script lang="ts">
-	import { onMount } from 'svelte';
 	import { toast } from 'svelte-sonner';
 	import { gotoWithi18n } from '$lib/i18n';
-	import Select from '$lib/components/select.svelte';
 	import * as Alert from '$lib/components/ui/alert';
+	import Select from '$lib/components/select.svelte';
 	import ConfigForm from '$lib/components/config-form.svelte';
-	import Button, { buttonVariants } from '$lib/components/ui/button/button.svelte';
+	import { Button, buttonVariants } from '$lib/components/ui/button';
 	import Code from '$lib/components/code.svelte';
 	import * as m from '$lib/paraglide/messages';
-	import { get, post } from '../api';
+	import {
+		createForkNetwork as createForkNetworkApi,
+		getTrainCommand,
+		startTrainModel
+	} from '$lib/apis/lm';
 	import * as Dialog from '$lib/components/ui/dialog';
 	import Input from '$lib/components/input.svelte';
 
-	let modelType = $state('');
-	let modelsOptions = $state<{ [key: string]: { model: any; trainer: any } } | null>(null);
+	let { data } = $props();
 
-	let dataset = $state('');
-	let tokenizer = $state('');
-	let datasetList = $state<string[]>([]);
-	let tokenizerList = $state<string[]>([]);
+	let modelType = $state(Object.keys(data.modelsOptions)[0]);
+
+	let dataset = $state(data.resources[0][0]);
+	let tokenizer = $state(data.resources[1][0]);
+
+	let fork = $state<string | null>(null);
+	let resume = $state('');
 
-	let modelKeys = $derived(
-		modelsOptions ? Object.keys(modelsOptions).map((item) => ({ value: item, label: item })) : []
-	);
 	let modelOption = $state<any>({});
 	let trainerOption = $state<any>({});
-	let modelConfig = $derived(modelsOptions ? modelsOptions[modelType]?.model : null);
-	let trainerConfig = $derived(modelsOptions ? modelsOptions[modelType]?.trainer : null);
-
-	let forkName = $state<string | null>(null);
+	let modelKeys = $derived(
+		data.modelsOptions
+			? Object.keys(data.modelsOptions).map((item) => ({ value: item, label: item }))
+			: []
+	);
+	let resumeList = $derived(
+		data.trainedNetworks ? data.trainedNetworks.map((v) => ({ value: v.name, label: v.name })) : []
+	);
+	let datasetList = $derived(data.resources[0].map((item) => ({ value: item, label: item })));
+	let tokenizerList = $derived(data.resources[1].map((item) => ({ value: item, label: item })));
+	let modelConfig = $derived(data.modelsOptions ? data.modelsOptions[modelType]?.model : null);
+	let trainerConfig = $derived(data.modelsOptions ? data.modelsOptions[modelType]?.trainer : null);
 
-	let command = $derived.by(() => {
-		const baseCommand = `darkit lm train --tokenizer ${tokenizer} --dataset ${dataset} ${modelType}`;
-		return `${baseCommand} ${generateCommand(modelOption)} ${generateCommand(trainerOption)}`;
+	let command = $derived.by(async () => {
+		return getTrainCommand(modelType, fork, resume, dataset, tokenizer, modelOption, trainerOption);
 	});
 
-	async function startTrain(command: string) {
+	async function startTrain() {
 		try {
-			const res = await post(`/model/${modelType}/train`, {
-				command: command
-			});
-			const data = await res.json();
-			if (res.status !== 200) {
-				toast.error(data.detail);
-				return;
-			}
+			const type = modelType;
+			await startTrainModel(type, fork, resume, dataset, tokenizer, modelOption, trainerOption);
 			gotoWithi18n(`/lm/visual/${trainerOption.name}`);
 		} catch (e) {
 			console.error(e);
+			toast.error('Failed to start training');
 		}
 	}
 
 	async function createForkNetwork() {
-		try {
-			const res = await post(`/edit/init/${forkName}`, {
-				model: modelType,
-				m_conf: modelOption,
-				t_conf: trainerOption
-			});
-			const data = await res.json();
-			if (res.status !== 200) {
-				toast.error(data.detail);
-				return;
+		if (fork) {
+			try {
+				await createForkNetworkApi(fork, modelType, modelOption, trainerOption);
+				gotoWithi18n(`/lm/fork/${fork}`);
+			} catch (e) {
+				console.error(e);
+				toast.error('Failed to create fork network');
 			}
-
-			gotoWithi18n(`/lm/fork/${forkName}`);
-		} catch (e) {
-			console.error(e);
 		}
 	}
-
-	// Generate command
-	function generateCommand(config: any) {
-		return Object.entries(config)
-			.filter(([key, val]) => {
-				return val !== null && val !== '' && val !== undefined;
-			})
-			.map(([key, val]) => `--${key} ${val}`)
-			.join(' ');
-	}
-
-	onMount(async () => {
-		const [modelsOptionsJson, resources] = await Promise.all([
-			get('/models/options').then((res) => res.json()),
-			get('/train/resources').then((res) => res.json())
-		]);
-
-		modelsOptions = modelsOptionsJson;
-		if (modelsOptions) {
-			if (modelType === '' || !Object.keys(modelsOptions).includes(modelType)) {
-				modelType = Object.keys(modelsOptions)[0];
-			}
-		}
-		if (resources) {
-			[datasetList, tokenizerList] = resources;
-			if (dataset === '' || !datasetList.includes(dataset)) {
-				dataset = datasetList ? datasetList[0] : '';
-			}
-			if (tokenizer === '' || !tokenizerList.includes(tokenizer)) {
-				tokenizer = tokenizerList ? tokenizerList[0] : '';
-			}
-		}
-	});
 </script>
 
 <div class="h-full flex-1 overflow-y-auto overflow-x-hidden p-8">
-	{#if modelsOptions}
+	{#if data.modelsOptions}
 		<Alert.Root class="text-primary/60 [&:has(svg)]:pl-4">
 			<Alert.Description>
 				<span class="text-xl"> 🍾 </span>
@@ -120,30 +84,36 @@
 					class="flex-1"
 					label="Tokenizers"
 					bind:value={tokenizer}
-					options={tokenizerList.map((v) => ({ value: v, label: v }))}
+					options={tokenizerList}
 					tooltip={m.tokenizersTooltip()}
 				/>
+				<Select
+					class="flex-1"
+					label="Resume"
+					bind:value={resume}
+					options={resumeList}
+					tooltip={m.resumeTooltip()}
+				/>
 			</div>
 			<ConfigForm key={modelType} config={modelConfig} bind:option={modelOption} />
 		</section>
 		<section class="py-4">
 			<h3 class="mb-2 font-bold">Trainer Config</h3>
-
-			<Select
-				label={m.dataset()}
-				class="mb-4 w-96"
-				bind:value={dataset}
-				options={datasetList.map((v) => ({ value: v, label: v }))}
-			/>
-
+			<Select label={m.dataset()} class="mb-4 w-96" bind:value={dataset} options={datasetList} />
 			<ConfigForm key={modelType} config={trainerConfig} bind:option={trainerOption} />
 		</section>
 
 		<h3 class="mb-2 mt-8 scroll-m-20 text-2xl font-semibold tracking-tight">
 			{m.generatedCommand()}
 		</h3>
 
-		<Code class="mt-8 max-h-96" lang="bash" content={command} wrap />
+		{#await command}
+			<Code class="mt-8 max-h-96" lang="bash" content="Updating..." wrap />
+		{:then command}
+			<Code class="mt-8 max-h-96" lang="bash" content={command} wrap />
+		{:catch error}
+			<Code class="mt-8 max-h-96" lang="bash" content={error.detail} wrap />
+		{/await}
 
 		<div class="w-full text-right">
 			<Dialog.Root>
@@ -163,15 +133,15 @@
 							placeholder="Enter a name for the fork"
 							class="w-full"
 							variant="row"
-							bind:value={forkName}
+							bind:value={fork}
 						/>
 					</div>
 					<Dialog.Footer>
 						<Button type="submit" onclick={createForkNetwork}>Create fork</Button>
 					</Dialog.Footer>
 				</Dialog.Content>
 			</Dialog.Root>
-			<Button size="lg" class="mt-8" onclick={() => startTrain(command)}>
+			<Button size="lg" class="mt-8" onclick={startTrain}>
 				{m.startTrain()}
 			</Button>
 		</div>

darkit/core/web/src/routes/lm/train/+page.ts

@@ -0,0 +1,13 @@
+import { getModelsOptions, getTrainResources, getTrainedNetworks } from '$lib/apis/lm';
+
+export const load = async () => {
+	const resources = await getTrainResources();
+	const modelsOptions = await getModelsOptions();
+	const trainedNetworks = await getTrainedNetworks();
+
+	return {
+		resources,
+		modelsOptions,
+		trainedNetworks
+	};
+};

darkit/core/web/src/routes/lm/visual/[models]/+page.svelte

@@ -2,6 +2,7 @@
 	import type { ModelDetail } from '$lib/apis/lm';
 </script>
 
+<!-- #TODO: RESUME 适配 -->
 <script lang="ts">
 	import { toast } from 'svelte-sonner';
 	import { onDestroy, onMount } from 'svelte';