contimp-app / data /oumi_training_schema.json
lefft's picture
contimp-app v0.1.0
910dadd verified
Raw
History Blame Contribute Delete
10.2 kB
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "Oumi TrainingConfig (curated subset)",
"description": "Subset of oumi.core.configs.TrainingConfig supported by the config copilot. Generated by scripts/dump_oumi_schema.py from oumi 0.1.4.",
"type": "object",
"additionalProperties": false,
"required": [
"model",
"data",
"training"
],
"properties": {
"model": {
"type": "object",
"additionalProperties": false,
"properties": {
"model_name": {
"type": "string",
"description": "HF Hub id or local path of the base model (required)"
},
"model_max_length": {
"type": [
"integer",
"null"
],
"description": "Max sequence length; null lets the model default apply"
},
"torch_dtype_str": {
"type": "string",
"default": "float32",
"description": "Model dtype, e.g. \"auto\", \"bfloat16\", \"float16\", \"float32\""
},
"trust_remote_code": {
"type": "boolean",
"default": false,
"description": "Allow models with custom code from the Hub"
},
"chat_template": {
"type": [
"string",
"null"
],
"description": "Name of the chat template to apply; null uses the tokenizer default"
}
},
"required": [
"model_name"
]
},
"data": {
"type": "object",
"additionalProperties": false,
"required": [
"train"
],
"properties": {
"train": {
"type": "object",
"additionalProperties": false,
"properties": {
"pack": {
"type": "boolean",
"default": false,
"description": "Pack multiple short examples into each sequence"
},
"datasets": {
"type": "array",
"minItems": 1,
"items": {
"type": "object",
"additionalProperties": false,
"properties": {
"dataset_name": {
"type": "string",
"description": "Registered dataset name, HF Hub id, or format name for local files (required)"
},
"dataset_path": {
"type": [
"string",
"null"
],
"description": "Path to a local dataset file (e.g. a .jsonl), if not loading from the Hub"
},
"split": {
"type": "string",
"default": "train",
"description": "Dataset split to load, e.g. \"train\""
},
"sample_count": {
"type": [
"integer",
"null"
],
"description": "Cap the number of examples drawn from this dataset"
}
},
"required": [
"dataset_name"
]
},
"description": "Datasets to mix for this split"
}
},
"required": [
"datasets"
]
},
"validation": {
"type": "object",
"additionalProperties": false,
"properties": {
"pack": {
"type": "boolean",
"default": false,
"description": "Pack multiple short examples into each sequence"
},
"datasets": {
"type": "array",
"minItems": 1,
"items": {
"type": "object",
"additionalProperties": false,
"properties": {
"dataset_name": {
"type": "string",
"description": "Registered dataset name, HF Hub id, or format name for local files (required)"
},
"dataset_path": {
"type": [
"string",
"null"
],
"description": "Path to a local dataset file (e.g. a .jsonl), if not loading from the Hub"
},
"split": {
"type": "string",
"default": "train",
"description": "Dataset split to load, e.g. \"train\""
},
"sample_count": {
"type": [
"integer",
"null"
],
"description": "Cap the number of examples drawn from this dataset"
}
},
"required": [
"dataset_name"
]
},
"description": "Datasets to mix for this split"
}
},
"required": []
}
}
},
"training": {
"type": "object",
"additionalProperties": false,
"properties": {
"trainer_type": {
"enum": [
"HF",
"OUMI",
"TRL_DPO",
"TRL_SFT",
"hf",
"oumi",
"trl_dpo",
"trl_sft"
],
"default": "HF",
"description": "Which trainer implementation to use"
},
"use_peft": {
"type": "boolean",
"default": false,
"description": "Train with parameter-efficient fine-tuning (LoRA/QLoRA); pair with the peft section"
},
"output_dir": {
"type": "string",
"default": "output",
"description": "Directory where checkpoints and the final model are written"
},
"num_train_epochs": {
"type": "integer",
"default": 3,
"description": "Number of passes over the training data (ignored if max_steps > 0)"
},
"max_steps": {
"type": "integer",
"default": -1,
"description": "Hard cap on optimizer steps; -1 disables"
},
"learning_rate": {
"type": "number",
"default": 5e-05,
"description": "Peak learning rate"
},
"per_device_train_batch_size": {
"type": "integer",
"default": 8,
"description": "Micro-batch size per device"
},
"gradient_accumulation_steps": {
"type": "integer",
"default": 1,
"description": "Steps to accumulate before each optimizer update"
},
"lr_scheduler_type": {
"type": "string",
"default": "linear",
"description": "LR schedule, e.g. \"linear\", \"cosine\", \"constant\""
},
"warmup_steps": {
"type": [
"integer",
"null"
],
"description": "LR warmup steps (alternative to warmup_ratio)"
},
"warmup_ratio": {
"type": [
"number",
"null"
],
"description": "LR warmup as a fraction of total steps"
},
"optimizer": {
"type": "string",
"default": "adamw_torch",
"description": "Optimizer name, e.g. \"adamw_torch\", \"adamw_torch_fused\", \"sgd\""
},
"weight_decay": {
"type": "number",
"default": 0.0,
"description": "Weight decay coefficient"
},
"mixed_precision_dtype": {
"enum": [
"BF16",
"FP16",
"NONE",
"bf16",
"fp16",
"none"
],
"default": "NONE",
"description": "Mixed-precision mode"
},
"enable_gradient_checkpointing": {
"type": "boolean",
"default": false,
"description": "Trade compute for memory during backprop"
},
"eval_strategy": {
"type": "string",
"default": "no",
"description": "\"no\", \"steps\", or \"epoch\""
},
"eval_steps": {
"type": "integer",
"default": 500,
"description": "Evaluate every N steps (when eval_strategy=steps)"
},
"save_steps": {
"type": "integer",
"default": 500,
"description": "Checkpoint every N steps"
},
"save_final_model": {
"type": "boolean",
"default": true,
"description": "Save the model at the end of training"
},
"logging_steps": {
"type": "integer",
"default": 50,
"description": "Log metrics every N steps"
},
"seed": {
"type": "integer",
"default": 42,
"description": "Random seed"
},
"run_name": {
"type": [
"string",
"null"
],
"description": "Human-readable name for the run"
}
}
},
"peft": {
"type": "object",
"additionalProperties": false,
"properties": {
"lora_r": {
"type": "integer",
"default": 8,
"description": "LoRA rank"
},
"lora_alpha": {
"type": "integer",
"default": 8,
"description": "LoRA scaling alpha"
},
"lora_dropout": {
"type": "number",
"default": 0.0,
"description": "Dropout on LoRA layers"
},
"lora_target_modules": {
"type": [
"array",
"null"
],
"items": {
"type": "string"
},
"description": "Module names to adapt, e.g. [\"q_proj\", \"v_proj\"]; null lets oumi pick"
},
"q_lora": {
"type": "boolean",
"default": false,
"description": "Quantize the base model (QLoRA)"
},
"q_lora_bits": {
"type": "integer",
"default": 4,
"description": "Quantization bits for QLoRA (typically 4)"
}
}
}
}
}