{
  "$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)"
        }
      }
    }
  }
}