From 6d27eedab2d07f37c55d28cba047f17fc34ba30d Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 13:30:01 -0400 Subject: [PATCH 01/17] add hyperparameter_search.py --- .../inverse_models/hyperparameter_search.py | 513 ++++++++++++++++++ 1 file changed, 513 insertions(+) create mode 100644 atomgpt/inverse_models/hyperparameter_search.py diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py new file mode 100644 index 0000000..2aa02af --- /dev/null +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -0,0 +1,513 @@ +#!/usr/bin/env python +# hyperparameter_search.py + +from __future__ import annotations + +import argparse, json, os, random, shutil, tempfile, time, csv, logging +from functools import partial +from pathlib import Path +from typing import Literal, Optional, Dict, List, Callable, Any + +import numpy as np, optuna, torch +from optuna.pruners import MedianPruner +from datasets import load_dataset +from optuna.trial import Trial +from pydantic_settings import BaseSettings +from transformers import IntervalStrategy, TrainingArguments, TrainerCallback +from peft import PeftModel + +from atomgpt.inverse_models.loader import FastLanguageModel +from atomgpt.inverse_models.custom_trainer import CustomSFTTrainer +from atomgpt.inverse_models.inverse_models import ( + evaluate, + make_alpaca_json, + formatting_prompts_func, + load_model, +) +from jarvis.db.jsonutils import dumpjson, loadjson +from jarvis.core.atoms import Atoms + +# ═════════════════════════════ Logging ══════════════════════════════ +""" +export ATOMGPT_DEBUG="true" to see debug lines in the console +""" +_DEBUG = os.getenv("ATOMGPT_DEBUG", "").lower() in {"1", "true", "yes", "y"} +logging.basicConfig( + level=logging.DEBUG if _DEBUG else logging.INFO, + format="%(asctime)s | %(levelname)-7s | %(message)s", + datefmt="%H:%M:%S", +) +log = logging.getLogger("hp_search") + +# ═════════════════════════════ Config ═══════════════════════════════ +""" +_SharedConfig contains fields in the normal AtomGPT config.json and the +new hyperparameter_search.py hp_search_cfg.json +""" +class _SharedConfig(BaseSettings): + instruction: str = "Below is a description of a material." + alpaca_prompt: str = "### Instruction:\n{}\n### Input:\n{}\n### Output:\n{}" + chem_info: str = "formula" + id_tag: str = "id" + prop: str = "Tc_supercon" + separator: str = "," + output_prompt: str = ( + " Generate atomic structure description with lattice lengths, angles, coordinates and atom types." + ) + +""" +TrainingConfig contains values and hyperparameters unique to the normal +AtomGPT config.json +""" +class TrainingConfig(_SharedConfig): + hp_cfg_path: str + id_prop_path: str + model_name: str + output_dir: str = "outputs" + model_save_path: str = "atomgpt_lora_model" + csv_out: str = "eval_results.csv" + file_format: Literal["poscar", "xyz", "pdb"] = "poscar" + prefix: str = "atomgpt_run" + + num_epochs: int = 2 + per_device_train_batch_size: int = 2 + gradient_accumulation_steps: int = 4 + learning_rate: float = 2e-4 + lora_rank: int = 16 + lora_alpha: int = 16 + max_seq_length: int = 2048 + optim: str = "adamw_8bit" + lr_scheduler_type: str = "linear" + warmup_ratio: float = 0.03 + logging_steps: int = 10 + seed_val: int = 42 + dataset_num_proc: int = 2 + dtype: str | None = None + load_in_4bit: bool = True + + val_ratio: float = 0.10 + test_ratio: float = 0.20 + num_train: int | None = None + num_test: int | None = None + +""" +OptunaSearchConfig is a schema for defining hyperparemeters +and values specific to a hyperparameter search study conducted +with this script +""" +class OptunaSearchConfig(BaseSettings): + parameters: Dict[str, Dict] + n_trials: int = 30 + objective_metric: str | None = None + objective_metrics: List[str] | None = None + study_direction: str | None = None + study_directions: List[str] | None = None + time_repeats: int = 1 + + +# ═════════════════════════ Metrics helpers ═══════════════════════════ +""" +last_value() returns the final value of the optimization parameter. +""" +def last_value(xs: List[float]) -> float: + return float("inf") if not xs else xs[-1] + + +""" +area_under_curve() returns the area under the curve of a set of +optimization parameter values. +""" +def area_under_curve(xs: List[float]) -> float: + return float("inf") if not xs else np.trapz(xs) + + +""" +trend_slope() returns the slope of the line of best fit for a +set of optimization parameter values. +""" +def trend_slope(xs: List[float]) -> float: + return float("inf") if len(xs) < 2 else abs(np.polyfit(range(len(xs)), xs, 1)[0]) + + +METRIC_EVALUATORS: Dict[str, Callable[[Dict[str, float]], float]] = { + "training_time": lambda m: m["training_time"], + "final_train_loss": lambda m: m["final_train_loss"], + "final_eval_loss": lambda m: m["final_eval_loss"], + "auc_train_loss": lambda m: m["auc_train_loss"], + "auc_eval_loss": lambda m: m["auc_eval_loss"], + "slope_train_loss": lambda m: m["slope_train_loss"], + "slope_eval_loss": lambda m: m["slope_eval_loss"], +} + + +def _auto_direction(metric: str) -> str: + return "maximize" if metric.lower() in {"accuracy", "f1"} else "minimize" + + +# ═════════════════════ Search-space sampler ══════════════════════════ +class SearchSpaceSampler: + _SUGGEST = { + "float": lambda t, k, s: t.suggest_float( + k, s["low"], s["high"], log=s.get("log", False) + ), + "int": lambda t, k, s: t.suggest_int(k, s["low"], s["high"]), + "categorical": lambda t, k, s: t.suggest_categorical(k, s["choices"]), + } + + def __init__(self, space: Dict[str, Dict]): + self.space = space + + def sample(self, trial: Trial) -> Dict[str, Any]: + sampled = {} + for k, spec in self.space.items(): + if not spec.get("include", True) or "condition" in spec: + continue + sampled[k] = self._SUGGEST[spec["type"]](trial, k, spec) + for k, spec in self.space.items(): + cond = spec.get("condition") + if cond and sampled.get(cond["param"]) == cond["value"]: + sampled[k] = self._SUGGEST[spec["type"]](trial, k, spec) + if _DEBUG: + log.debug("Trial %d — sampled params: %s", trial.number, sampled) + return sampled + + +# ═════════════════════ Optuna pruning callback ═══════════════════════ +""" +Ends the current trial if deemed unpromising. +""" +class OptunaPruningCallback(TrainerCallback): + def __init__(self, trial: Trial, key: str): + self.trial, self.key = trial, key + + def on_evaluate(self, *_, metrics=None, **__): + if metrics and self.key in metrics: + step = metrics.get("epoch", 0) + self.trial.report(metrics[self.key], step) + if self.trial.should_prune(): + raise optuna.TrialPruned() + + +# ═════════════════════ Split helpers ═════════════════════════════════ +def _set_seeds(seed: int): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +""" +Deterministically shuffle IDs and split into train/val/test lists. +""" +def train_val_test_split_ids( + data: List[dict], id_tag: str, seed: int, val_ratio: float, test_ratio: float +): + if val_ratio + test_ratio >= 1.0: + raise ValueError("val_ratio + test_ratio must be < 1.") + ids = [r[id_tag] for r in data] + rng = np.random.default_rng(seed) + rng.shuffle(ids) + + n = len(ids) + n_val = max(1, int(n * val_ratio)) + n_test = max(1, int(n * test_ratio)) + n_train = n - n_val - n_test + if n_train <= 0: + raise ValueError("Split sizes invalid – make dataset larger or ratios smaller.") + + val_ids = ids[:n_val] + test_ids = ids[n_val : n_val + n_test] + train_ids = ids[n_val + n_test :] + return train_ids, val_ids, test_ids + + +# ═════════════════════ Single train-pass ═════════════════════════════ +def _train_once( + cfg: TrainingConfig, + train_json: Path, + val_json: Path, + prune_cb: TrainerCallback | None, +) -> Dict[str, float]: + + model, tok, _ = load_model(path=cfg.model_name, config=cfg) + if not isinstance(model, PeftModel): + model = FastLanguageModel.get_peft_model( + model, + r=cfg.lora_rank, + target_modules=[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + lora_alpha=cfg.lora_alpha, + lora_dropout=0, + bias="none", + use_gradient_checkpointing=True, + ) + # ─── datasets ──────────────────────────────────────────────────── + train_ds = load_dataset("json", data_files=str(train_json), split="train") + val_ds = load_dataset("json", data_files=str(val_json), split="train") + + fmt = lambda e: formatting_prompts_func(e, cfg.alpaca_prompt) + train_ds = train_ds.map(fmt, batched=True) + val_ds = val_ds.map(fmt, batched=True) + + if _DEBUG: + log.debug( + "Prepared datasets — train: %d samples, val: %d samples", + len(train_ds), + len(val_ds), + ) + + trainer = CustomSFTTrainer( + model=model, + tokenizer=tok, + train_dataset=train_ds, + eval_dataset=val_ds, + dataset_text_field="text", + max_seq_length=cfg.max_seq_length, + dataset_num_proc=cfg.dataset_num_proc, + packing=False, + args=TrainingArguments( + output_dir=cfg.output_dir, + num_train_epochs=cfg.num_epochs, + per_device_train_batch_size=cfg.per_device_train_batch_size, + gradient_accumulation_steps=cfg.gradient_accumulation_steps, + learning_rate=cfg.learning_rate, + warmup_ratio=cfg.warmup_ratio, + lr_scheduler_type=cfg.lr_scheduler_type, + optim=cfg.optim, + fp16=not torch.cuda.is_bf16_supported(), + bf16=torch.cuda.is_bf16_supported(), + logging_steps=cfg.logging_steps, + evaluation_strategy=IntervalStrategy.EPOCH, + save_strategy=IntervalStrategy.NO, + report_to="none", + seed=cfg.seed_val, + ), + ) + if prune_cb: + trainer.add_callback(prune_cb) + + start = time.perf_counter() + trainer.train() + trainer.save_model(cfg.model_save_path) + + runtime = trainer.state.log_history[-1].get( + "train_runtime", time.perf_counter() - start + ) + + tl = [ + e["loss"] + for e in trainer.state.log_history + if "loss" in e and e.get("step") is None + ] + el = [e["eval_loss"] for e in trainer.state.log_history if "eval_loss" in e] + + metrics = { + "training_time": runtime, + "final_train_loss": last_value(tl), + "final_eval_loss": last_value(el), + "auc_train_loss": area_under_curve(tl), + "auc_eval_loss": area_under_curve(el), + "slope_train_loss": trend_slope(tl), + "slope_eval_loss": trend_slope(el), + } + if _DEBUG: + log.debug("Single pass metrics: %s", metrics) + + del model, tok, trainer + torch.cuda.empty_cache() + return metrics + + +# ═════════════════════ Optuna objective ══════════════════════════════ +def objective( + trial: Trial, + train_cfg: TrainingConfig, + hp_cfg: OptunaSearchConfig, + sampler: SearchSpaceSampler, + train_json: Path, + val_json: Path, + test_json: Path, + objective_metrics: List[str], +) -> List[float] | float: + + _set_seeds(train_cfg.seed_val + trial.number) + cfg = train_cfg.copy(deep=True) + + for k, v in sampler.sample(trial).items(): + setattr(cfg, k, v) + + work = Path(tempfile.mkdtemp(prefix="optuna_")) + cfg.output_dir = str(work / "out") + cfg.model_save_path = str(work / "model") + cfg.csv_out = str(work / "eval.csv") + os.makedirs(cfg.output_dir, exist_ok=True) + + try: + metrics_avgs = [] + for _ in range(hp_cfg.time_repeats): + prune_cb = OptunaPruningCallback(trial, objective_metrics[0]) + metrics_avgs.append(_train_once(cfg, train_json, val_json, prune_cb)) + + metrics = { + k: float(np.mean([d[k] for d in metrics_avgs])) for k in metrics_avgs[0] + } + + model, tok = FastLanguageModel.from_pretrained(cfg.model_save_path) + + for k, v in metrics.items(): + trial.set_user_attr(k, v) + trial.set_user_attr("metrics_vec", [metrics[m] for m in objective_metrics]) + + log_path = Path("logs") + log_path.mkdir(exist_ok=True) + with open(log_path / "optuna_trials.jsonl", "a") as f: + f.write( + json.dumps( + {"number": trial.number, "params": trial.params, "metrics": metrics} + ) + + "\n" + ) + + if _DEBUG: + log.debug("Trial %d finished — metrics: %s", trial.number, metrics) + + out = [METRIC_EVALUTATORS[m](metrics) for m in objective_metrics] + return out[0] if len(out) == 1 else tuple(out) + + finally: + shutil.rmtree(work, ignore_errors=True) + torch.cuda.empty_cache() + + +# ═══════════════════ id_prop.csv loader ════════════════════ +def _load_id_prop_data(id_prop_csv: str, cfg: TrainingConfig) -> List[dict]: + """ + Read a standard id_prop.csv file and accompanying structure files, + returning records compatible with `make_alpaca_json`. + """ + base = Path(id_prop_csv).parent + with open(id_prop_csv) as fh: + rows = list(csv.reader(fh)) + + records: list[dict] = [] + for row in rows: + rid, *vals = row + prop_val = ( + cfg.separator.join(map(str, map(float, vals))) + if len(vals) > 1 + else str(float(vals[0])) + ) + + fpath = base / rid + if cfg.file_format == "poscar": + atoms = Atoms.from_poscar(fpath) + elif cfg.file_format == "xyz": + atoms = Atoms.from_xyz(fpath) + elif cfg.file_format == "pdb": + atoms = Atoms.from_pdb(fpath) + else: + raise ValueError(f"Unsupported file_format '{cfg.file_format}'") + + records.append( + { + cfg.id_tag: rid, + cfg.prop: prop_val, + "atoms": atoms.to_dict(), + } + ) + return records + + +# ═════════════════════ CLI / study orchestration ═════════════════════ +def main() -> None: + """ + Entrypoint: run an Optuna HPO study for AtomGPT fine-tuning. + """ + p = argparse.ArgumentParser() + p.add_argument("--config_name", required=True, help="Path to a TrainingConfig JSON") + args = p.parse_args() + + train_cfg = TrainingConfig(**json.load(open(args.config_name))) + hp_cfg = OptunaSearchConfig(**json.load(open(train_cfg.hp_cfg_path))) + + objective_metrics = hp_cfg.objective_metrics or ( + [hp_cfg.objective_metric] if hp_cfg.objective_metric else ["final_eval_loss"] + ) + directions = hp_cfg.study_directions or ( + [hp_cfg.study_direction] if hp_cfg.study_direction else None + ) + if directions is None: + directions = [_auto_direction(k) for k in objective_metrics] + + if _DEBUG: + log.debug("Objectives: %s | Directions: %s", objective_metrics, directions) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA GPU required") + + data = _load_id_prop_data(train_cfg.id_prop_path, train_cfg) + + train_ids, val_ids, test_ids = train_val_test_split_ids( + data, + train_cfg.id_tag, + train_cfg.seed_val, + train_cfg.val_ratio, + train_cfg.test_ratio, + ) + + if _DEBUG: + log.debug( + "Dataset split sizes — train: %d | val: %d | test: %d", + len(train_ids), + len(val_ids), + len(test_ids), + ) + + tmp = Path(tempfile.mkdtemp(prefix="optuna_data_")) + train_j = tmp / "train.json" + val_j = tmp / "val.json" + test_j = tmp / "test.json" + dumpjson(make_alpaca_json(data, train_ids, config=train_cfg), train_j) + dumpjson(make_alpaca_json(data, val_ids, config=train_cfg), val_j) + dumpjson(make_alpaca_json(data, test_ids, config=train_cfg), test_j) + + sampler = SearchSpaceSampler(hp_cfg.parameters) + pruner = optuna.pruners.MedianPruner(n_warmup_steps=1) + study = optuna.create_study(directions=directions, pruner=pruner) + + wall = time.time() + study.optimize( + partial( + objective, + train_cfg=train_cfg, + hp_cfg=hp_cfg, + sampler=sampler, + train_json=train_j, + val_json=val_j, + test_json=test_j, + objective_metrics=objective_metrics, + ), + n_trials=hp_cfg.n_trials, + ) + runtime = time.time() - wall + print("\nStudy finished in %.1fs" % runtime) + if len(objective_metrics) == 1: + print("Best value :", study.best_value) + else: + print("Best values:", study.best_values) + print("Best params :", study.best_params) + + if _DEBUG: + log.debug("Full study completed in %.1fs", runtime) + + +# ═════════════════════════════════════════════════════════════════════ +if __name__ == "__main__": + main() From 850af0b30b8d94e5db363559337a99f63ca4215c Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 13:52:18 -0400 Subject: [PATCH 02/17] add an example hp_search_config.json --- .../examples/inverse_model/hp_search_config.json | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 atomgpt/examples/inverse_model/hp_search_config.json diff --git a/atomgpt/examples/inverse_model/hp_search_config.json b/atomgpt/examples/inverse_model/hp_search_config.json new file mode 100644 index 0000000..b94468a --- /dev/null +++ b/atomgpt/examples/inverse_model/hp_search_config.json @@ -0,0 +1,13 @@ +{ + "n_trials": 25, + "time_repeats": 1, + "objective_metric": "final_eval_loss", + + "parameters": { + "learning_rate": { "type": "float", "low": 5e-6, "high": 5e-4, "log": true,"include":true }, + "per_device_train_batch_size": { "type": "categorical", "choices": [1, 2, 4],"include":true }, + "gradient_accumulation_steps": { "type": "categorical", "choices": [2, 4, 8],"include":false }, + "lora_rank": { "type": "categorical", "choices": [8, 16, 32],"include":false }, + "lora_alpha": { "type": "categorical", "choices": [8, 16, 32], "include":false } + } +} From b26541dfabd89a60da8515947998f92553cccfee Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 13:54:46 -0400 Subject: [PATCH 03/17] add a field that points to the hp_search_config.json --- atomgpt/examples/inverse_model/config.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/atomgpt/examples/inverse_model/config.json b/atomgpt/examples/inverse_model/config.json index 4c3f03c..eab3438 100644 --- a/atomgpt/examples/inverse_model/config.json +++ b/atomgpt/examples/inverse_model/config.json @@ -25,5 +25,6 @@ "load_in_4bit": true, "instruction": "Below is a description of a superconductor material.", "alpaca_prompt": "### Instruction:\n{}\n### Input:\n{}\n### Output:\n{}", - "output_prompt": " Generate atomic structure description with lattice lengths, angles, coordinates and atom types." + "output_prompt": " Generate atomic structure description with lattice lengths, angles, coordinates and atom types.", + "hp_cfg_path": "hp_search_config.json" } From 4e36a22beb95a0c74e3d2f82f7f362bd7626c1f0 Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 13:57:36 -0400 Subject: [PATCH 04/17] add optuna>=3.5 --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index fedbb26..b7cb9c1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -59,6 +59,7 @@ nvidia-cusparselt-cu12==0.6.3 nvidia-nccl-cu12==2.26.2 nvidia-nvjitlink-cu12==12.6.85 nvidia-nvtx-cu12==12.6.77 +optuna>=3.5 packaging==25.0 pandas==2.2.3 pathspec==0.12.1 From 258a75fb9400921adf8dfa8b4ffcf2d89898c425 Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 13:59:43 -0400 Subject: [PATCH 05/17] fix typo and replace best_values with best_trials --- atomgpt/inverse_models/hyperparameter_search.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index 2aa02af..16292fc 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -378,7 +378,7 @@ def objective( if _DEBUG: log.debug("Trial %d finished — metrics: %s", trial.number, metrics) - out = [METRIC_EVALUTATORS[m](metrics) for m in objective_metrics] + out = [METRIC_EVALUATORS[m](metrics) for m in objective_metrics] return out[0] if len(out) == 1 else tuple(out) finally: @@ -500,9 +500,11 @@ def main() -> None: print("\nStudy finished in %.1fs" % runtime) if len(objective_metrics) == 1: print("Best value :", study.best_value) + print("Best params:", study.best_params) else: - print("Best values:", study.best_values) - print("Best params :", study.best_params) + print("Pareto front (top 5 shown):") + for i, t in enumerate(study.best_trials[:5]): + print(f" Trial {t.number}: values={t.values}, params={t.params}") if _DEBUG: log.debug("Full study completed in %.1fs", runtime) From 43722fe966fbe863a90cd92c6d1088b95ed769e7 Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 14:01:40 -0400 Subject: [PATCH 06/17] optuna>=3.5,<4 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index b7cb9c1..7e4d84a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -59,7 +59,7 @@ nvidia-cusparselt-cu12==0.6.3 nvidia-nccl-cu12==2.26.2 nvidia-nvjitlink-cu12==12.6.85 nvidia-nvtx-cu12==12.6.77 -optuna>=3.5 +optuna>=3.5,<4 packaging==25.0 pandas==2.2.3 pathspec==0.12.1 From 3013988156c3a77b2e85055b809d2323b3c398fb Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 14:21:07 -0400 Subject: [PATCH 07/17] import TrainingPropConfig in hp_search.py instead of making a new one --- .../inverse_models/hyperparameter_search.py | 61 ++----------------- atomgpt/inverse_models/inverse_models.py | 1 + 2 files changed, 7 insertions(+), 55 deletions(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index 16292fc..4ed692c 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -23,6 +23,7 @@ make_alpaca_json, formatting_prompts_func, load_model, + TrainingPropConfig ) from jarvis.db.jsonutils import dumpjson, loadjson from jarvis.core.atoms import Atoms @@ -40,56 +41,6 @@ log = logging.getLogger("hp_search") # ═════════════════════════════ Config ═══════════════════════════════ -""" -_SharedConfig contains fields in the normal AtomGPT config.json and the -new hyperparameter_search.py hp_search_cfg.json -""" -class _SharedConfig(BaseSettings): - instruction: str = "Below is a description of a material." - alpaca_prompt: str = "### Instruction:\n{}\n### Input:\n{}\n### Output:\n{}" - chem_info: str = "formula" - id_tag: str = "id" - prop: str = "Tc_supercon" - separator: str = "," - output_prompt: str = ( - " Generate atomic structure description with lattice lengths, angles, coordinates and atom types." - ) - -""" -TrainingConfig contains values and hyperparameters unique to the normal -AtomGPT config.json -""" -class TrainingConfig(_SharedConfig): - hp_cfg_path: str - id_prop_path: str - model_name: str - output_dir: str = "outputs" - model_save_path: str = "atomgpt_lora_model" - csv_out: str = "eval_results.csv" - file_format: Literal["poscar", "xyz", "pdb"] = "poscar" - prefix: str = "atomgpt_run" - - num_epochs: int = 2 - per_device_train_batch_size: int = 2 - gradient_accumulation_steps: int = 4 - learning_rate: float = 2e-4 - lora_rank: int = 16 - lora_alpha: int = 16 - max_seq_length: int = 2048 - optim: str = "adamw_8bit" - lr_scheduler_type: str = "linear" - warmup_ratio: float = 0.03 - logging_steps: int = 10 - seed_val: int = 42 - dataset_num_proc: int = 2 - dtype: str | None = None - load_in_4bit: bool = True - - val_ratio: float = 0.10 - test_ratio: float = 0.20 - num_train: int | None = None - num_test: int | None = None - """ OptunaSearchConfig is a schema for defining hyperparemeters and values specific to a hyperparameter search study conducted @@ -223,7 +174,7 @@ def train_val_test_split_ids( # ═════════════════════ Single train-pass ═════════════════════════════ def _train_once( - cfg: TrainingConfig, + cfg: TrainingPropConfig, train_json: Path, val_json: Path, prune_cb: TrainerCallback | None, @@ -328,7 +279,7 @@ def _train_once( # ═════════════════════ Optuna objective ══════════════════════════════ def objective( trial: Trial, - train_cfg: TrainingConfig, + train_cfg: TrainingPropConfig, hp_cfg: OptunaSearchConfig, sampler: SearchSpaceSampler, train_json: Path, @@ -387,7 +338,7 @@ def objective( # ═══════════════════ id_prop.csv loader ════════════════════ -def _load_id_prop_data(id_prop_csv: str, cfg: TrainingConfig) -> List[dict]: +def _load_id_prop_data(id_prop_csv: str, cfg: TrainingPropConfig) -> List[dict]: """ Read a standard id_prop.csv file and accompanying structure files, returning records compatible with `make_alpaca_json`. @@ -431,10 +382,10 @@ def main() -> None: Entrypoint: run an Optuna HPO study for AtomGPT fine-tuning. """ p = argparse.ArgumentParser() - p.add_argument("--config_name", required=True, help="Path to a TrainingConfig JSON") + p.add_argument("--config_name", required=True, help="Path to a TrainingPropConfig JSON") args = p.parse_args() - train_cfg = TrainingConfig(**json.load(open(args.config_name))) + train_cfg = TrainingPropConfig(**json.load(open(args.config_name))) hp_cfg = OptunaSearchConfig(**json.load(open(train_cfg.hp_cfg_path))) objective_metrics = hp_cfg.objective_metrics or ( diff --git a/atomgpt/inverse_models/inverse_models.py b/atomgpt/inverse_models/inverse_models.py index 1b7ff27..8c7d733 100644 --- a/atomgpt/inverse_models/inverse_models.py +++ b/atomgpt/inverse_models/inverse_models.py @@ -99,6 +99,7 @@ class TrainingPropConfig(BaseSettings): " Generate atomic structure description with lattice lengths, angles, coordinates and atom types." ) # num_val: Optional[int] = 2 + hp_cfg_path: Optional[str] = "hp_search_config.json" def get_input(config=None, chem="", val=10): From 287d7c0ab759a0563cbda47e7715b2cc1d24b687 Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 14:40:36 -0400 Subject: [PATCH 08/17] send SFTTrainer hardcoded args to config and add hp_opt training config parameters to config --- atomgpt/inverse_models/inverse_models.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/atomgpt/inverse_models/inverse_models.py b/atomgpt/inverse_models/inverse_models.py index 8c7d733..7778ec3 100644 --- a/atomgpt/inverse_models/inverse_models.py +++ b/atomgpt/inverse_models/inverse_models.py @@ -64,14 +64,14 @@ class TrainingPropConfig(BaseSettings): gradient_accumulation_steps: int = 4 num_train: Optional[int] = None num_test: Optional[int] = None - test_ratio: Optional[float] = 0.2 + test_ratio: Optional[float] = 0.1 + val_ratio: Optional[float] = 0.1 model_save_path: str = "atomgpt_lora_model" lora_rank: Optional[int] = 16 lora_alpha: Optional[int] = 16 loss_type: str = "default" optim: str = "adamw_8bit" id_tag: str = "id" - save_strategy: str = "st" lr_scheduler_type: str = "linear" separator: str = "," prop: str = "Tc_supercon" @@ -100,6 +100,13 @@ class TrainingPropConfig(BaseSettings): ) # num_val: Optional[int] = 2 hp_cfg_path: Optional[str] = "hp_search_config.json" + per_device_train_batch_size: int = 2 + gradient_accumulation_steps: int = 4 + warmup_steps: int = 3 + warmup_ratio: float = 0.0 + logging_steps: int = 10 + + def get_input(config=None, chem="", val=10): @@ -574,12 +581,13 @@ def tokenize_function(example): args=SFTConfig( dataset_text_field="text", max_seq_length=config.max_seq_length, - per_device_train_batch_size=2, - gradient_accumulation_steps=4, - warmup_steps=5, + per_device_train_batch_size=config.per_device_train_batch_size, + gradient_accumulation_steps=config.gradient_accumulation_steps, + warmup_steps=config.warmup_steps, overwrite_output_dir=True, + warmup_ratio=config.warmup_ratio, # max_steps=60, - logging_steps=1, + logging_steps=config.logging_steps, output_dir=config.output_dir, optim=config.optim, seed=config.seed_val, From a04c0ad071c3b0b8a26f3d5b37810ee8f808b0e6 Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 14:49:03 -0400 Subject: [PATCH 09/17] use strings for eval and save strategies --- atomgpt/inverse_models/hyperparameter_search.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index 4ed692c..e7d139f 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -235,8 +235,8 @@ def _train_once( fp16=not torch.cuda.is_bf16_supported(), bf16=torch.cuda.is_bf16_supported(), logging_steps=cfg.logging_steps, - evaluation_strategy=IntervalStrategy.EPOCH, - save_strategy=IntervalStrategy.NO, + eval_strategy="epoch", + save_strategy="no", report_to="none", seed=cfg.seed_val, ), From 5bb5daf074675fc9c130b4e9c36c65b1085bb5d7 Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 15:06:58 -0400 Subject: [PATCH 10/17] remove customsfttrainer and use regular hf sfttrainer --- atomgpt/inverse_models/hyperparameter_search.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index e7d139f..14388ba 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -17,7 +17,7 @@ from peft import PeftModel from atomgpt.inverse_models.loader import FastLanguageModel -from atomgpt.inverse_models.custom_trainer import CustomSFTTrainer +from trl import SFTTrainer from atomgpt.inverse_models.inverse_models import ( evaluate, make_alpaca_json, @@ -214,7 +214,7 @@ def _train_once( len(val_ds), ) - trainer = CustomSFTTrainer( + trainer = SFTTrainer( model=model, tokenizer=tok, train_dataset=train_ds, From 8595698e3ce1df03faf7d824c4e6f5da637a8441 Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 15:19:55 -0400 Subject: [PATCH 11/17] try-except-pass wrap for del m._flag_for_generation --- atomgpt/inverse_models/llama.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/atomgpt/inverse_models/llama.py b/atomgpt/inverse_models/llama.py index 5ea2280..f647a75 100644 --- a/atomgpt/inverse_models/llama.py +++ b/atomgpt/inverse_models/llama.py @@ -3458,8 +3458,15 @@ def _for_training(m): if hasattr(m, "_saved_temp_tokenizer"): m._saved_temp_tokenizer.padding_side = "right" # Set a flag for generation! - if hasattr(m, "_flag_for_generation"): - del m._flag_for_generation + if "_flag_for_generation" in getattr(m, "__dict__", {}): + m.__dict__.pop("_flag_for_generation", None) + else: + try: + # If it exists virtually, neutralize it rather than delete. + setattr(m, "_flag_for_generation", False) + except Exception: + print("TRY-EXCEPT-PASS TRIGGERED llama.py line 3468") + pass pass m = model From 9eca98b73b428d966cc42ef57e41c5c55b6c71d6 Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 15:30:57 -0400 Subject: [PATCH 12/17] replace tokenizer= with processing_class= --- atomgpt/inverse_models/hyperparameter_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index 14388ba..e7367d6 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -216,7 +216,7 @@ def _train_once( trainer = SFTTrainer( model=model, - tokenizer=tok, + processing_class=tok, train_dataset=train_ds, eval_dataset=val_ds, dataset_text_field="text", From d0be0a6e40708664b6ea6e38bd92ff4d23d1d63c Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 15:39:21 -0400 Subject: [PATCH 13/17] remove dataset_text_field --- atomgpt/inverse_models/hyperparameter_search.py | 1 - 1 file changed, 1 deletion(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index e7367d6..78a959f 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -219,7 +219,6 @@ def _train_once( processing_class=tok, train_dataset=train_ds, eval_dataset=val_ds, - dataset_text_field="text", max_seq_length=cfg.max_seq_length, dataset_num_proc=cfg.dataset_num_proc, packing=False, From b3e5ac0521f176d3ecd7ea681e78a1067cd08e0a Mon Sep 17 00:00:00 2001 From: crhysc Date: Thu, 28 Aug 2025 16:17:31 -0400 Subject: [PATCH 14/17] add sft_args to conform to atomgpt_compiled_cache --- .../inverse_models/hyperparameter_search.py | 51 +++++++++++-------- 1 file changed, 29 insertions(+), 22 deletions(-) diff --git a/atomgpt/inverse_models/hyperparameter_search.py b/atomgpt/inverse_models/hyperparameter_search.py index 78a959f..d49899e 100644 --- a/atomgpt/inverse_models/hyperparameter_search.py +++ b/atomgpt/inverse_models/hyperparameter_search.py @@ -17,7 +17,7 @@ from peft import PeftModel from atomgpt.inverse_models.loader import FastLanguageModel -from trl import SFTTrainer +from trl import SFTTrainer, SFTConfig from atomgpt.inverse_models.inverse_models import ( evaluate, make_alpaca_json, @@ -214,31 +214,38 @@ def _train_once( len(val_ds), ) + sft_args = SFTConfig( + # --- training --- + output_dir=cfg.output_dir, + num_train_epochs=cfg.num_epochs, + per_device_train_batch_size=cfg.per_device_train_batch_size, + gradient_accumulation_steps=cfg.gradient_accumulation_steps, + learning_rate=cfg.learning_rate, + warmup_ratio=cfg.warmup_ratio, + lr_scheduler_type=cfg.lr_scheduler_type, + optim=cfg.optim, + fp16=not torch.cuda.is_bf16_supported(), + bf16=torch.cuda.is_bf16_supported(), + logging_steps=cfg.logging_steps, + eval_strategy="epoch", + save_strategy="no", + report_to="none", + seed=cfg.seed_val, + + # --- data prep (must live in config for AtomGPT wrapper) --- + dataset_text_field="text", + dataset_num_proc=cfg.dataset_num_proc, + max_seq_length=cfg.max_seq_length, + packing=False, # or True if you really want packing; wrapper warns about it + ) + trainer = SFTTrainer( model=model, - processing_class=tok, + args=sft_args, + processing_class=tok, # <- replaces tokenizer=tok train_dataset=train_ds, eval_dataset=val_ds, - max_seq_length=cfg.max_seq_length, - dataset_num_proc=cfg.dataset_num_proc, - packing=False, - args=TrainingArguments( - output_dir=cfg.output_dir, - num_train_epochs=cfg.num_epochs, - per_device_train_batch_size=cfg.per_device_train_batch_size, - gradient_accumulation_steps=cfg.gradient_accumulation_steps, - learning_rate=cfg.learning_rate, - warmup_ratio=cfg.warmup_ratio, - lr_scheduler_type=cfg.lr_scheduler_type, - optim=cfg.optim, - fp16=not torch.cuda.is_bf16_supported(), - bf16=torch.cuda.is_bf16_supported(), - logging_steps=cfg.logging_steps, - eval_strategy="epoch", - save_strategy="no", - report_to="none", - seed=cfg.seed_val, - ), + # (optional) compute_metrics=..., callbacks=..., peft_config=..., formatting_func=... ) if prune_cb: trainer.add_callback(prune_cb) From 3b3ed5688c804bdf68ea8032109952fa3842ca8a Mon Sep 17 00:00:00 2001 From: Charles Campbell Date: Thu, 28 Aug 2025 16:25:14 -0400 Subject: [PATCH 15/17] update to a more robust example --- .../inverse_model/hp_search_config.json | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/atomgpt/examples/inverse_model/hp_search_config.json b/atomgpt/examples/inverse_model/hp_search_config.json index b94468a..f0e4897 100644 --- a/atomgpt/examples/inverse_model/hp_search_config.json +++ b/atomgpt/examples/inverse_model/hp_search_config.json @@ -2,12 +2,24 @@ "n_trials": 25, "time_repeats": 1, "objective_metric": "final_eval_loss", - "parameters": { - "learning_rate": { "type": "float", "low": 5e-6, "high": 5e-4, "log": true,"include":true }, - "per_device_train_batch_size": { "type": "categorical", "choices": [1, 2, 4],"include":true }, - "gradient_accumulation_steps": { "type": "categorical", "choices": [2, 4, 8],"include":false }, - "lora_rank": { "type": "categorical", "choices": [8, 16, 32],"include":false }, - "lora_alpha": { "type": "categorical", "choices": [8, 16, 32], "include":false } + "learning_rate": { "type": "float", "low": 5e-6, "high": 5e-4, "log": true, "include": true }, + "per_device_train_batch_size": { "type": "categorical", "choices": [1, 2, 4], "include": true }, + "num_epochs": { "type": "int", "low": 1, "high": 5, "include": false }, + "gradient_accumulation_steps": { "type": "categorical", "choices": [1, 2, 4, 8], "include": false }, + "lora_rank": { "type": "categorical", "choices": [8, 16, 32, 64], "include": false }, + "lora_alpha": { "type": "categorical", "choices": [8, 16, 32, 64], "include": false }, + "max_seq_length": { "type": "categorical", "choices": [1024, 1536, 2048], "include": false }, + "optim": { "type": "categorical", "choices": ["adamw_torch", "adamw_torch_fused", "adamw_bnb_8bit", "adamw_hf"], "include": false }, + "lr_scheduler_type": { "type": "categorical", "choices": ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant_with_warmup"], "include": false }, + "warmup_ratio": { "type": "float", "low": 0.0, "high": 0.1, "include": false }, + "logging_steps": { "type": "int", "low": 5, "high": 200, "include": false }, + "seed_val": { "type": "int", "low": 1, "high": 10000, "include": false }, + "dataset_num_proc": { "type": "categorical", "choices": [1, 2, 4, 8], "include": false }, + "val_ratio": { "type": "float", "low": 0.05, "high": 0.2, "include": false }, + "test_ratio": { "type": "float", "low": 0.1, "high": 0.3, "include": false }, + "num_train": { "type": "int", "low": 0, "high": 1000000, "include": false }, + "num_test": { "type": "int", "low": 0, "high": 100000, "include": false } } } + From 9e07f93f7c6ad3da0741225cdf2ad7c5a683e1b6 Mon Sep 17 00:00:00 2001 From: Kamal Choudhary Date: Sun, 14 Sep 2025 06:11:22 -0400 Subject: [PATCH 16/17] Update README.md --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index c7075ae..b4e7595 100644 --- a/README.md +++ b/README.md @@ -221,6 +221,8 @@ Funding support NIST-MGI (https://www.nist.gov/mgi) and CHIPS (https://www.nist.gov/chips) +Note: This project was originally developed under the github.com/usnistgov organization and is now maintained here by the lead developer. + Code of conduct -------------------- From f90521571b5bfeb6f29c396a0e2c974084b32542 Mon Sep 17 00:00:00 2001 From: Kamal Choudhary Date: Sun, 14 Sep 2025 16:27:33 -0400 Subject: [PATCH 17/17] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b4e7595..ef7a4c8 100644 --- a/README.md +++ b/README.md @@ -221,7 +221,7 @@ Funding support NIST-MGI (https://www.nist.gov/mgi) and CHIPS (https://www.nist.gov/chips) -Note: This project was originally developed under the github.com/usnistgov organization and is now maintained here by the lead developer. +Note: This project was originally developed under the github.com/usnistgov organization. New updates and developments will be carried out here. Code of conduct --------------------