from __future__ import annotations import argparse import json import os from collections import Counter, defaultdict from pathlib import Path from typing import Any import torch from datasets import Dataset from peft import LoraConfig, prepare_model_for_kbit_training from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from trl import SFTConfig, SFTTrainer DEFAULT_MODEL = "Qwen/Qwen3-4B-Instruct-2507" DEFAULT_DATASET = "outputs/datasets/grid2op_sft_v1.jsonl" DEFAULT_OUTPUT_DIR = "outputs/models/grid2op-qwen3-4b-sft-v1" QWEN_CHATML_TRAINING_TEMPLATE = """{%- for message in messages %} {{- '<|im_start|>' + message['role'] + '\n' }} {%- if message['role'] == 'assistant' %} {% generation %}{{- message['content'] }}{% endgeneration %} {%- else %} {{- message['content'] }} {%- endif %} {{- '<|im_end|>\n' }} {%- endfor %} {%- if add_generation_prompt %} {{- '<|im_start|>assistant\n' }} {%- endif %} """ def resolve_precision(precision: str) -> tuple[torch.dtype, bool, bool, str]: if precision == "auto": if torch.cuda.is_available() and torch.cuda.is_bf16_supported(): precision = "bf16" else: precision = "fp16" if precision == "bf16": return torch.bfloat16, True, False, precision if precision == "fp16": return torch.float16, False, True, precision if precision == "fp32": return torch.float32, False, False, precision raise ValueError(f"Unsupported precision: {precision}") def load_jsonl_rows(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for line_no, line in enumerate(path.read_text().splitlines(), 1): if not line.strip(): continue row = json.loads(line) messages = row.get("messages") if not isinstance(messages, list) or len(messages) < 2: raise ValueError(f"Row {line_no} has no valid messages list") if messages[-1].get("role") != "assistant": raise ValueError(f"Row {line_no} final message must be assistant") content = messages[-1].get("content") if not isinstance(content, str): raise ValueError(f"Row {line_no} assistant content must be a string") json.loads(content) rows.append(row) if not rows: raise ValueError(f"No rows found in {path}") return rows def summarize_rows(rows: list[dict[str, Any]]) -> dict[str, Any]: tasks: Counter[str] = Counter() actions: Counter[str] = Counter() tiers: Counter[str] = Counter() policies: Counter[str] = Counter() task_actions: dict[str, Counter[str]] = defaultdict(Counter) prompt_chars: list[int] = [] completion_chars: list[int] = [] for row in rows: metadata = row.get("metadata", {}) task_id = str(metadata.get("task_id", "unknown")) tasks[task_id] += 1 tiers[str(metadata.get("benchmark_tier", "unknown"))] += 1 policies[str(metadata.get("label_policy", "unknown"))] += 1 selected_action = metadata.get("selected_action", {}) action_type = selected_action_type(selected_action) actions[action_type] += 1 task_actions[task_id][action_type] += 1 messages = row["messages"] prompt_chars.append(sum(len(str(message.get("content", ""))) for message in messages[:-1])) completion_chars.append(len(str(messages[-1].get("content", "")))) return { "rows": len(rows), "tasks": dict(sorted(tasks.items())), "actions": dict(sorted(actions.items())), "benchmark_tiers": dict(sorted(tiers.items())), "label_policies": dict(sorted(policies.items())), "task_actions": { task: dict(sorted(counter.items())) for task, counter in sorted(task_actions.items()) }, "avg_prompt_chars": sum(prompt_chars) / len(prompt_chars), "max_prompt_chars": max(prompt_chars), "avg_completion_chars": sum(completion_chars) / len(completion_chars), "max_completion_chars": max(completion_chars), } def selected_action_type(action: dict[str, Any]) -> str: if action.get("do_nothing"): return "do_nothing" if action.get("redispatch"): return "redispatch" if action.get("line_set"): statuses = [int(value) for value in action["line_set"].values()] if statuses and statuses[0] == 1: return "reconnect_line" if statuses and statuses[0] == -1: return "disconnect_line" return "line_set" return "empty" def add_token_lengths(dataset: Dataset, tokenizer, max_samples: int) -> dict[str, Any]: lengths: list[int] = [] sample_count = min(len(dataset), max_samples) for index in range(sample_count): text = tokenizer.apply_chat_template( dataset[index]["messages"], tokenize=False, add_generation_prompt=False, ) lengths.append(len(tokenizer(text, add_special_tokens=False)["input_ids"])) if not lengths: return {"sampled_token_lengths": 0} lengths_sorted = sorted(lengths) p95_index = min(len(lengths_sorted) - 1, int(0.95 * (len(lengths_sorted) - 1))) return { "sampled_token_lengths": len(lengths), "avg_tokens": sum(lengths) / len(lengths), "max_tokens": max(lengths), "p95_tokens": lengths_sorted[p95_index], } def build_datasets(rows: list[dict[str, Any]], eval_ratio: float, seed: int) -> tuple[Dataset, Dataset | None]: dataset = Dataset.from_list([{"messages": row["messages"], "metadata": row.get("metadata", {})} for row in rows]) dataset = dataset.shuffle(seed=seed) if eval_ratio <= 0 or len(dataset) < 4: return dataset, None split = dataset.train_test_split(test_size=eval_ratio, seed=seed) return split["train"], split["test"] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="QLoRA SFT for Grid2Op chat-action data.") parser.add_argument("--dataset", type=Path, default=Path(DEFAULT_DATASET)) parser.add_argument("--model", default=DEFAULT_MODEL) parser.add_argument("--output-dir", type=Path, default=Path(DEFAULT_OUTPUT_DIR)) parser.add_argument("--run-name", default="grid2op-qwen3-4b-sft-v1") parser.add_argument("--wandb-project", default=os.environ.get("WANDB_PROJECT", "grid2op-openenv-sft")) parser.add_argument("--wandb-entity", default=os.environ.get("WANDB_ENTITY")) parser.add_argument("--eval-ratio", type=float, default=0.05) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--max-length", type=int, default=4096) parser.add_argument("--num-train-epochs", type=float, default=1.0) parser.add_argument("--max-steps", type=int, default=-1) parser.add_argument("--learning-rate", type=float, default=2e-4) parser.add_argument("--per-device-train-batch-size", type=int, default=1) parser.add_argument("--per-device-eval-batch-size", type=int, default=1) parser.add_argument("--gradient-accumulation-steps", type=int, default=8) parser.add_argument("--logging-steps", type=int, default=5) parser.add_argument("--save-steps", type=int, default=50) parser.add_argument("--eval-steps", type=int, default=50) parser.add_argument("--save-total-limit", type=int, default=3) parser.add_argument("--lora-r", type=int, default=32) parser.add_argument("--lora-alpha", type=int, default=64) parser.add_argument("--lora-dropout", type=float, default=0.05) parser.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--packing", action=argparse.BooleanOptionalAction, default=False) parser.add_argument("--assistant-only-loss", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--patch-qwen-training-template", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--precision", choices=["auto", "bf16", "fp16", "fp32"], default="auto") parser.add_argument("--use-4bit", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--use-liger-kernel", action=argparse.BooleanOptionalAction, default=False) parser.add_argument("--attn-implementation", default=None) parser.add_argument("--device-map", default="auto") parser.add_argument("--pad-to-multiple-of", type=int, default=8) parser.add_argument("--torch-empty-cache-steps", type=int, default=25) parser.add_argument("--token-length-samples", type=int, default=256) return parser.parse_args() def main() -> None: args = parse_args() if args.use_liger_kernel and args.device_map != "none": raise ValueError( "--use-liger-kernel is not safe with --device-map auto/sharded loading in this QLoRA setup. " "Use --no-use-liger-kernel for two-GPU device_map=auto training, or use --device-map none " "only when the full model fits on one GPU/process." ) os.environ["WANDB_PROJECT"] = args.wandb_project if args.wandb_entity: os.environ["WANDB_ENTITY"] = args.wandb_entity rows = load_jsonl_rows(args.dataset) dataset_summary = summarize_rows(rows) tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" if args.assistant_only_loss and args.patch_qwen_training_template and "Qwen" in args.model: tokenizer.chat_template = QWEN_CHATML_TRAINING_TEMPLATE train_dataset, eval_dataset = build_datasets( rows=rows, eval_ratio=args.eval_ratio, seed=args.seed, ) token_summary = add_token_lengths( Dataset.from_list([{"messages": row["messages"]} for row in rows]), tokenizer=tokenizer, max_samples=args.token_length_samples, ) model_dtype, bf16, fp16, resolved_precision = resolve_precision(args.precision) quantization_config = None if args.use_4bit: quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=model_dtype, bnb_4bit_use_double_quant=True, ) device_map = None if args.device_map == "none" else args.device_map model_kwargs: dict[str, Any] = { "trust_remote_code": True, "dtype": model_dtype, "device_map": device_map, } if quantization_config is not None: model_kwargs["quantization_config"] = quantization_config if args.attn_implementation: model_kwargs["attn_implementation"] = args.attn_implementation model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs) model.config.use_cache = False if args.use_4bit: model = prepare_model_for_kbit_training( model, use_gradient_checkpointing=args.gradient_checkpointing, ) lora_config = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, bias="none", task_type="CAUSAL_LM", target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], ) training_args = SFTConfig( output_dir=str(args.output_dir), run_name=args.run_name, report_to=["wandb"], max_length=args.max_length, num_train_epochs=args.num_train_epochs, max_steps=args.max_steps, learning_rate=args.learning_rate, per_device_train_batch_size=args.per_device_train_batch_size, per_device_eval_batch_size=args.per_device_eval_batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, logging_steps=args.logging_steps, logging_first_step=True, save_strategy="steps", save_steps=args.save_steps, save_total_limit=args.save_total_limit, eval_strategy="steps" if eval_dataset is not None else "no", eval_steps=args.eval_steps if eval_dataset is not None else None, gradient_checkpointing=args.gradient_checkpointing, gradient_checkpointing_kwargs={"use_reentrant": False}, optim="paged_adamw_8bit" if args.use_4bit else "adamw_torch_fused", warmup_ratio=0.03, weight_decay=0.01, lr_scheduler_type="cosine", fp16=fp16, bf16=bf16, packing=args.packing, pad_to_multiple_of=args.pad_to_multiple_of, assistant_only_loss=args.assistant_only_loss, use_liger_kernel=args.use_liger_kernel, torch_empty_cache_steps=args.torch_empty_cache_steps, remove_unused_columns=True, seed=args.seed, data_seed=args.seed, ) import wandb wandb.init( project=args.wandb_project, entity=args.wandb_entity, name=args.run_name, config={ "base_model": args.model, "dataset": str(args.dataset), "output_dir": str(args.output_dir), "dataset_summary": dataset_summary, "token_summary": token_summary, "lora": { "r": args.lora_r, "alpha": args.lora_alpha, "dropout": args.lora_dropout, }, "qlora_4bit": args.use_4bit, "max_length": args.max_length, "precision": resolved_precision, "torch_dtype": str(model_dtype), "device_map": args.device_map, "attn_implementation": args.attn_implementation, "use_liger_kernel": args.use_liger_kernel, "patch_qwen_training_template": args.patch_qwen_training_template, "pad_to_multiple_of": args.pad_to_multiple_of, "torch_empty_cache_steps": args.torch_empty_cache_steps, }, ) wandb.summary.update({f"dataset/{key}": value for key, value in dataset_summary.items() if not isinstance(value, dict)}) wandb.summary.update({f"tokens/{key}": value for key, value in token_summary.items()}) trainer = SFTTrainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, peft_config=lora_config, processing_class=tokenizer, ) trainer.train() trainer.save_model(str(args.output_dir)) tokenizer.save_pretrained(str(args.output_dir)) wandb.finish() if __name__ == "__main__": main()