Gemma training script

tools/gemma_train/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "gemma_train"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+dependencies = [
+    "accelerate>=0.26.0",
+    "bitsandbytes>=0.45.0",
+    "llama-cpp-scripts @ file:///root/sia/modules/llama.cpp",
+    "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
+    "trl>=0.17.0",
+    "peft>=0.15.0",
+    "python-dotenv>=1.0.0",
+]
+
+[project.scripts]
+gemma_train = "gemma_train.__main__:main"

tools/gemma_train/src/gemma_train/__main__.py

@@ -0,0 +1,153 @@
+from llm_engine_utils.dataset import Dataset
+from pathlib import Path
+from peft import LoraConfig, AutoPeftModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
+from trl import SFTTrainer
+import os
+import sys
+import torch
+
+from .config import Config
+
+def main():
+    config = Config()
+    train(config)
+    merge(config)
+    os.system(f"cp -r {config.output_dir}/tokenizer/* {config.output_dir}/merged")
+    convert_to_gguf(config)
+
+def train(config: Config):
+    tokenizer = AutoTokenizer.from_pretrained(config.model, token=config.api_key)
+    tokenizer.save_pretrained(config.output_dir/"tokenizer")
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        config.model,
+        quantization_config=bnb_config,
+        device_map="auto",
+        token=config.api_key,
+        attn_implementation='eager',
+    )
+
+    dataset = Dataset(config.config_path)
+    dataset.validate()
+    dataset = dataset.to_transformers_dataset(tokenizer)
+
+    lora_config = LoraConfig(
+        r=4,
+        lora_alpha=4,
+        target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    training_args = TrainingArguments(
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=4,
+        warmup_steps=1,
+        max_steps=1,
+        learning_rate=1e-3,
+        fp16=True,
+        logging_steps=1,
+        save_strategy="steps",
+        save_steps=1,
+        output_dir=config.output_dir/"lora",
+        optim="paged_adamw_8bit",
+        seed=42,
+        group_by_length=True,
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=dataset,
+        args=training_args,
+        peft_config=lora_config,
+        formatting_func=format_sia_example,
+    )
+
+    trainer.train()
+
+    trainer.model.save_pretrained(config.output_dir/"lora_adapter")
+
+def merge(config: Config):
+    adapted_model = AutoPeftModelForCausalLM.from_pretrained(
+        config.output_dir/"lora_adapter",
+        torch_dtype=torch.float16,  # Use float16 for better compatibility
+        device_map="auto",
+        offload_folder="offload",
+        token=config.api_key,
+    )
+    merged_model = adapted_model.merge_and_unload()
+    merged_model.save_pretrained(
+        config.output_dir/"merged",
+        safe_serialization=True
+    )
+
+def convert_to_gguf(config: Config):
+    """Convert the merged model to GGUF format using llama.cpp's convert_hf_to_gguf script."""
+    print("Converting merged model to GGUF format...")
+    
+    # Add path to llama.cpp directory
+    sys.path.append("./llama.cpp")
+    
+    try:
+        # Import the necessary components from the conversion script
+        from convert_hf_to_gguf import ModelBase, gguf, ModelType
+        
+        # Set up paths
+        dir_model = config.output_dir / "merged"
+        fname_out = config.output_dir / "model.gguf"
+        output_type = gguf.LlamaFileType.MOSTLY_Q8_0  # Using Q8_0 quantization
+        
+        # Run the conversion with torch inference mode
+        with torch.inference_mode():
+            # Load hyperparameters
+            hparams = ModelBase.load_hparams(dir_model)
+            model_architecture = hparams["architectures"][0]
+            model_type = ModelType.TEXT
+            
+            print(f"Model architecture: {model_architecture}")
+            
+            try:
+                # Get the appropriate model class
+                model_class = ModelBase.from_model_architecture(model_architecture, model_type=model_type)
+                
+                # Create model instance
+                model_instance = model_class(
+                    dir_model, 
+                    output_type, 
+                    fname_out,
+                    is_big_endian=False, 
+                    use_temp_file=False, 
+                    eager=False
+                )
+                
+                # Export the model
+                print(f"Exporting model to GGUF format...")
+                model_instance.write()
+                print(f"Model successfully exported to {model_instance.fname_out}")
+                
+            except NotImplementedError:
+                print(f"Error: Model architecture {model_architecture} is not supported for GGUF conversion")
+                print("Skipping GGUF conversion")
+    
+    except ImportError as e:
+        print(f"Error importing conversion script: {e}")
+        print("Make sure llama.cpp is properly installed and accessible")
+        print("Skipping GGUF conversion")
+    except Exception as e:
+        print(f"Error during GGUF conversion: {e}")
+        print("Skipping GGUF conversion")
+
+def format_sia_example(example):
+    return example['messages'].removeprefix("<bos>")
+
+if __name__ == "__main__":
+    exit(main())

tools/gemma_train/src/gemma_train/config.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+import argparse
+import os
+
+class Config:
+    def __init__(self):
+        parser = argparse.ArgumentParser(description='Train Gemma model and convert to gguf')
+        parser.add_argument(
+            '--config',
+            type=Path,
+            default=Path('/root/sia/training/config.yaml'),
+            help='Path to config file (default: /root/sia/training/config.yaml)'
+        )
+        parser.add_argument(
+            '--model',
+            type=str,
+            default='google/gemma-3-1b-it',
+            help='Base model for fine-tuning (default: google/gemma-3-1b-it)'
+        )
+        parser.add_argument(
+            '--api-key',
+            type=str,
+            default=os.environ.get('SIA_HF_API_KEY'),
+            help='Huggingface API key (optional, env: SIA_HF_API_KEY)'
+        )
+        parser.add_argument(
+            '--output-dir',
+            type=Path,
+            default=Path('/root/models/current'),
+            help='Output directory for fine-tuned model and converted gguf (default: /root/models/current)'
+        )
+        self.args = parser.parse_args()
+
+    @property
+    def config_path(self) -> Path:
+        return self.args.config
+
+    @property
+    def model(self) -> str:
+        return self.args.model
+    
+    @property
+    def api_key(self) -> str:
+        return self.args.api_key
+    
+    @property
+    def output_dir(self) -> Path:
+        return self.args.output_dir

tools/notebook/requirements.txt

@@ -0,0 +1,3 @@
+jupyter
+ipykernel
+ipywidgets