mirror of
https://github.com/nomic-ai/gpt4all.git
synced 2025-08-09 03:47:25 +00:00
feat: adamw, fix training, log gradients
This commit is contained in:
parent
5c5f41ba36
commit
3efc19ebc5
84
train.py
84
train.py
@ -1,5 +1,5 @@
|
|||||||
import os
|
import os
|
||||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
from transformers import AutoModelForCausalLM, AutoTokenizer, AdamW, get_scheduler
|
||||||
from transformers.trainer_pt_utils import get_parameter_names
|
from transformers.trainer_pt_utils import get_parameter_names
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
@ -11,6 +11,7 @@ from peft import get_peft_model, LoraConfig, TaskType
|
|||||||
from data import load_data
|
from data import load_data
|
||||||
from torchmetrics import MeanMetric
|
from torchmetrics import MeanMetric
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
import wandb
|
||||||
|
|
||||||
|
|
||||||
def format_metrics(metrics, split, prefix=""):
|
def format_metrics(metrics, split, prefix=""):
|
||||||
@ -20,17 +21,12 @@ def format_metrics(metrics, split, prefix=""):
|
|||||||
return log
|
return log
|
||||||
|
|
||||||
|
|
||||||
def evaluate(config, model, val_dataloader):
|
def evaluate(model, val_dataloader):
|
||||||
model.eval()
|
model.eval()
|
||||||
val_loss = MeanMetric().to(model.device)
|
val_loss = MeanMetric().to(model.device)
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
for i, batch in enumerate(
|
for batch in tqdm(val_dataloader):
|
||||||
tqdm(val_dataloader),
|
|
||||||
):
|
|
||||||
if i == config["eval_steps"]:
|
|
||||||
break
|
|
||||||
|
|
||||||
loss = model(**batch).loss
|
loss = model(**batch).loss
|
||||||
|
|
||||||
loss_values = accelerator.gather_for_metrics({"loss": loss.detach()})
|
loss_values = accelerator.gather_for_metrics({"loss": loss.detach()})
|
||||||
@ -49,8 +45,7 @@ def train(accelerator, config):
|
|||||||
tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'])
|
tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'])
|
||||||
# llama has no pad token, set it to new token
|
# llama has no pad token, set it to new token
|
||||||
if tokenizer.pad_token is None:
|
if tokenizer.pad_token is None:
|
||||||
# these tokens are already in the vocab, just not mapped correctly
|
tokenizer.pad_token = tokenizer.eos_token
|
||||||
added_tokens = tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>", "pad_token": "<pad>"})
|
|
||||||
|
|
||||||
|
|
||||||
with accelerator.main_process_first():
|
with accelerator.main_process_first():
|
||||||
@ -61,10 +56,6 @@ def train(accelerator, config):
|
|||||||
model = AutoModelForCausalLM.from_pretrained(config["model_name"],
|
model = AutoModelForCausalLM.from_pretrained(config["model_name"],
|
||||||
use_cache=False if checkpoint else True,
|
use_cache=False if checkpoint else True,
|
||||||
trust_remote_code=True)
|
trust_remote_code=True)
|
||||||
|
|
||||||
if added_tokens > 0:
|
|
||||||
model.resize_token_embeddings(len(tokenizer))
|
|
||||||
|
|
||||||
if checkpoint:
|
if checkpoint:
|
||||||
model.gradient_checkpointing_enable()
|
model.gradient_checkpointing_enable()
|
||||||
|
|
||||||
@ -77,19 +68,55 @@ def train(accelerator, config):
|
|||||||
model.print_trainable_parameters()
|
model.print_trainable_parameters()
|
||||||
|
|
||||||
optimizer_cls = (
|
optimizer_cls = (
|
||||||
torch.optim.AdamW
|
AdamW
|
||||||
if accelerator.state.deepspeed_plugin is None
|
if accelerator.state.deepspeed_plugin is None
|
||||||
or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
|
or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
|
||||||
else DummyOptim
|
else DummyOptim
|
||||||
)
|
)
|
||||||
|
|
||||||
|
no_decay = ["bias", "LayerNorm.weight"]
|
||||||
|
optimizer_grouped_parameters = [
|
||||||
|
{
|
||||||
|
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
|
||||||
|
"weight_decay": config["weight_decay"],
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
|
||||||
|
"weight_decay": 0.0,
|
||||||
|
},
|
||||||
|
]
|
||||||
|
|
||||||
# karpathy doesn't decay embeddding, maybe we should exclude
|
# karpathy doesn't decay embeddding, maybe we should exclude
|
||||||
# https://github.com/karpathy/minGPT/commit/bbbdac74fa9b2e55574d70056163ffbae42310c1#diff-2075fa9c224b395be5bda85544dd36572b59c76c54562819eadadbf268602834R157s
|
# https://github.com/karpathy/minGPT/commit/bbbdac74fa9b2e55574d70056163ffbae42310c1#diff-2075fa9c224b395be5bda85544dd36572b59c76c54562819eadadbf268602834R157s
|
||||||
optimizer = optimizer_cls(model.parameters(), lr=config["lr"])
|
optimizer = optimizer_cls(optimizer_grouped_parameters, lr=config["lr"])
|
||||||
|
|
||||||
# scheduler defined in Deepspeed config
|
if accelerator.state.deepspeed_plugin is not None:
|
||||||
scheduler = DummyScheduler(
|
gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
|
||||||
optimizer, warmup_num_steps=config["warmup_steps"],
|
"gradient_accumulation_steps"
|
||||||
|
]
|
||||||
|
|
||||||
|
# decay to min_lr instead of 0
|
||||||
|
lr_ratio = config["min_lr"] / config["lr"]
|
||||||
|
accelerator.print(f"Len of train_dataloader: {len(train_dataloader)}")
|
||||||
|
total_num_steps = (len(train_dataloader) / gradient_accumulation_steps) * config["num_epochs"]
|
||||||
|
# instead of decaying to zero, decay to ratio of min_lr / lr
|
||||||
|
total_num_steps += int(total_num_steps * lr_ratio) + config["warmup_steps"]
|
||||||
|
accelerator.print(f"Total training steps: {total_num_steps}")
|
||||||
|
|
||||||
|
# Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
|
||||||
|
if (
|
||||||
|
accelerator.state.deepspeed_plugin is None
|
||||||
|
or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
|
||||||
|
):
|
||||||
|
scheduler = get_scheduler(
|
||||||
|
name="cosine",
|
||||||
|
optimizer=optimizer,
|
||||||
|
num_warmup_steps=config["warmup_steps"] * accelerator.num_processes,
|
||||||
|
num_training_steps=total_num_steps * accelerator.num_processes,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
scheduler = DummyScheduler(
|
||||||
|
optimizer, total_num_steps=config["warmup_steps"], warmup_num_steps=config["warmup_steps"]
|
||||||
)
|
)
|
||||||
|
|
||||||
model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
|
model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
|
||||||
@ -108,14 +135,13 @@ def train(accelerator, config):
|
|||||||
accelerator.skip_first_batches(train_dataloader, resume_step)
|
accelerator.skip_first_batches(train_dataloader, resume_step)
|
||||||
accelerator.print(f"Resuming from step {resume_step}")
|
accelerator.print(f"Resuming from step {resume_step}")
|
||||||
|
|
||||||
train_loss = MeanMetric().to(model.device)
|
|
||||||
|
|
||||||
if accelerator.state.deepspeed_plugin is not None:
|
# log gradients
|
||||||
gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
|
if accelerator.is_local_main_process and config["wandb"]:
|
||||||
"gradient_accumulation_steps"
|
wandb.watch(model, log_freq=config["log_grads_every"])
|
||||||
]
|
|
||||||
|
|
||||||
for epoch in range(config["num_epochs"]):
|
for epoch in range(config["num_epochs"]):
|
||||||
|
train_loss = MeanMetric().to(model.device)
|
||||||
for step, batch in enumerate(tqdm(train_dataloader)):
|
for step, batch in enumerate(tqdm(train_dataloader)):
|
||||||
model.train()
|
model.train()
|
||||||
outputs = model(**batch)
|
outputs = model(**batch)
|
||||||
@ -139,9 +165,10 @@ def train(accelerator, config):
|
|||||||
train_loss.update(loss_values["loss"])
|
train_loss.update(loss_values["loss"])
|
||||||
|
|
||||||
if step > 0 and step % config["save_every"] == 0:
|
if step > 0 and step % config["save_every"] == 0:
|
||||||
accelerator.save_state(f"{config['output_dir']}/step_{step}")
|
curr_step = step + epoch * len(train_dataloader)
|
||||||
|
accelerator.save_state(f"{config['output_dir']}/step_{curr_step}")
|
||||||
|
|
||||||
if step > 0 and step % config["eval_every"] == 0:
|
if step > 0 and (step % config["eval_every"] == 0 or step == len(train_dataloader) - 1):
|
||||||
val_loss = evaluate(config, model, val_dataloader)
|
val_loss = evaluate(config, model, val_dataloader)
|
||||||
|
|
||||||
log_train = {
|
log_train = {
|
||||||
@ -166,7 +193,7 @@ def train(accelerator, config):
|
|||||||
accelerator.wait_for_everyone()
|
accelerator.wait_for_everyone()
|
||||||
unwrapped_model = accelerator.unwrap_model(model)
|
unwrapped_model = accelerator.unwrap_model(model)
|
||||||
if accelerator.is_main_process:
|
if accelerator.is_main_process:
|
||||||
unwrapped_model.push_to_hub(config["save_name"] + "_first_epoch", private=True)
|
unwrapped_model.push_to_hub(config["save_name"] + f"-epoch_{epoch}", private=True)
|
||||||
|
|
||||||
|
|
||||||
accelerator.wait_for_everyone()
|
accelerator.wait_for_everyone()
|
||||||
@ -178,9 +205,6 @@ def train(accelerator, config):
|
|||||||
state_dict=accelerator.get_state_dict(model),
|
state_dict=accelerator.get_state_dict(model),
|
||||||
)
|
)
|
||||||
|
|
||||||
if accelerator.is_main_process:
|
|
||||||
unwrapped_model.push_to_hub(config["save_name"], private=True)
|
|
||||||
|
|
||||||
accelerator.end_training()
|
accelerator.end_training()
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user