cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: use Ornith-1.0-35B-FP8 with multi-GPU support

Browse Source
This commit is contained in:
Christian Medina
2026-06-30 15:57:54 -04:00
parent 215886d6a7
commit 2074d3d993
2 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
training/configs/llama2-7b-lora.yaml
View File

@@ -1,7 +1,7 @@
# LoRA Training Configuration for Llama-2-7b # LoRA Training Configuration for Llama-2-7b
# Dataset: cyron_summary_lora_dataset (20k examples) # Dataset: cyron_summary_lora_dataset (20k examples)
base_model: meta-llama/Llama-2-7b-hf base_model: deepreinforce-ai/Ornith-1.0-35B-FP8
model_type: LlamaForCausalLM model_type: LlamaForCausalLM
tokenizer_type: LlamaTokenizer tokenizer_type: LlamaTokenizer

5
training/scripts/train.py
View File

@@ -39,10 +39,13 @@ def train(config_path):
use_nested_quant=config.get("use_nested_quant", False), use_nested_quant=config.get("use_nested_quant", False),
) )
# Use all available GPUs
device_map = "auto" if torch.cuda.device_count() == 1 else "balanced"
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto", device_map=device_map,
) )
model = prepare_model_for_kbit_training(model) model = prepare_model_for_kbit_training(model)