wip refactor
This commit is contained in:
parent
11ea9eeaa7
commit
5acb6809ed
5 changed files with 292 additions and 232 deletions
229
train_dynamic.py
229
train_dynamic.py
|
|
@ -1,156 +1,18 @@
|
|||
import transformers
|
||||
from transformers import AutoModelForCausalLM, get_scheduler
|
||||
from peft.utils import _get_submodules
|
||||
from transformers import get_scheduler
|
||||
|
||||
import torch
|
||||
from torch.utils import tensorboard
|
||||
import os
|
||||
import shutil
|
||||
import math
|
||||
import collections
|
||||
from tqdm.auto import tqdm
|
||||
from random import randint
|
||||
import collections
|
||||
|
||||
|
||||
from arguments import DataArguments, ModelArguments, TrainingArguments
|
||||
from datamodules import create_data_module_s2s, create_data_module
|
||||
from convertinglinear import ConvertingLinear
|
||||
from tokenizer import get_tokenizer
|
||||
|
||||
|
||||
def find_all_linear_module_names(model):
|
||||
module_names = set()
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, torch.nn.Linear) or isinstance(module, ConvertingLinear):
|
||||
module_names.add(name)
|
||||
|
||||
if 'lm_head' in module_names: # needed for 16-bit
|
||||
module_names.remove('lm_head')
|
||||
return list(module_names)
|
||||
|
||||
|
||||
def find_all_outher_module_names(model):
|
||||
module_names = set()
|
||||
for name, module in model.named_modules():
|
||||
if not (isinstance(module, torch.nn.Linear) or isinstance(module, ConvertingLinear)):
|
||||
module_names.add(name)
|
||||
return list(module_names)
|
||||
|
||||
|
||||
def get_model(model_args: ModelArguments, cache_dir, gradient_checkpointing):
|
||||
dtype = torch.float16 if training_args.fp16 or (training_args.storage_fp16 and model_args.max_instant_params > 0) else torch.float32
|
||||
print(f'loading base model {model_args.model_name_or_path} in {dtype}...')
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
cache_dir=cache_dir,
|
||||
torch_dtype=dtype if model_args.max_instant_params > 0 else torch.float32,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
device_map=None,
|
||||
attn_implementation="flash_attention_2"
|
||||
)
|
||||
|
||||
modules = dict(model.named_modules())
|
||||
keys = find_all_linear_module_names(model)
|
||||
for key in keys:
|
||||
modules[key].weight = modules[key].weight.to(dtype)
|
||||
if modules[key].bias is not None:
|
||||
modules[key].bias = modules[key].bias.to(dtype)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def recursive_setattr(obj, attr, value):
|
||||
attr = attr.split('.', 1)
|
||||
if len(attr) == 1:
|
||||
setattr(obj, attr[0], value)
|
||||
else:
|
||||
recursive_setattr(getattr(obj, attr[0]), attr[1], value)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def set_linear_module_frozen_simple(module, frozen: bool, dtype: torch.dtype, device: torch.device):
|
||||
new_module = torch.nn.Linear(module.in_features,
|
||||
module.out_features,
|
||||
module.bias is not None,
|
||||
module.weight.device,
|
||||
dtype)
|
||||
new_module.weight = torch.nn.Parameter(module.weight.detach().clone())
|
||||
new_module.bias = torch.nn.Parameter(module.bias.detach().clone()) if module.bias is not None else None
|
||||
new_module.weight.requires_grad = not frozen
|
||||
if new_module.bias is not None:
|
||||
new_module.bias.requires_grad = not frozen
|
||||
return new_module
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def set_linear_module_frozen(module, frozen: bool, dtype: torch.dtype, device: torch.device):
|
||||
if type(module) is torch.nn.Linear:
|
||||
if frozen:
|
||||
module.weight.requires_grad = False
|
||||
if module.bias is not None:
|
||||
module.bias.requires_grad = False
|
||||
return module.to(dtype).to(device)
|
||||
else:
|
||||
new_module = ConvertingLinear.fromLinear(module).to(dtype)
|
||||
new_module.weight.requires_grad = True
|
||||
if new_module.bias is not None:
|
||||
new_module.bias.requires_grad = True
|
||||
return new_module.to(device)
|
||||
elif type(module) is ConvertingLinear:
|
||||
if not frozen:
|
||||
module.weight.requires_grad = True
|
||||
if module.bias is not None:
|
||||
module.bias.requires_grad = True
|
||||
assert False
|
||||
return module.to(dtype).to(device)
|
||||
else:
|
||||
new_module = torch.nn.utils.skip_init(torch.nn.Linear, in_features=module.in_features,
|
||||
out_features=module.out_features,
|
||||
bias=module.bias is not None,
|
||||
device=module.weight.device,
|
||||
dtype=dtype)
|
||||
new_module.weight = torch.nn.Parameter(module.weight.to(dtype))
|
||||
new_module.bias = torch.nn.Parameter(module.bias.to(dtype)) if module.bias is not None else None
|
||||
new_module.weight.requires_grad = False
|
||||
if new_module.bias is not None:
|
||||
new_module.bias.requires_grad = False
|
||||
return new_module.to(device)
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def freeze_random_modules(model, target_params: int, frozen_dtype: torch.dtype, frozen_device: torch.device, active_device: torch.device):
|
||||
modules = dict(model.named_modules())
|
||||
linear_names = find_all_linear_module_names(model)
|
||||
|
||||
for key in linear_names:
|
||||
if modules[key].weight.dtype != frozen_dtype or modules[key].weight.requires_grad or modules[key].weight.requires_grad:
|
||||
parent, target, target_name = _get_submodules(model, key)
|
||||
setattr(parent, target_name, set_linear_module_frozen(modules[key], True, frozen_dtype, frozen_device))
|
||||
modules = dict(model.named_modules())
|
||||
|
||||
active_paramter_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
if active_paramter_count > target_params:
|
||||
raise RuntimeError("Enough paramters must be available to train at least one linear layer")
|
||||
|
||||
while active_paramter_count < target_params and len(linear_names) > 0:
|
||||
i = randint(0, len(linear_names) - 1)
|
||||
parent, target, target_name = _get_submodules(model, linear_names[i])
|
||||
new_module = set_linear_module_frozen(modules[linear_names[i]], False, torch.float32, active_device)
|
||||
setattr(parent, target_name, new_module)
|
||||
active_paramter_count += modules[linear_names[i]].weight.numel()
|
||||
if modules[linear_names[i]].bias is not None:
|
||||
active_paramter_count += modules[linear_names[i]].bias.numel()
|
||||
linear_names.pop(i)
|
||||
modules = dict()
|
||||
|
||||
assert active_paramter_count == sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
return active_paramter_count
|
||||
from dyntrainmodel import DyntrainModel
|
||||
|
||||
|
||||
def save_model(model, global_step: int, output_dir: str, max_checkpoints: int = 0):
|
||||
|
|
@ -172,26 +34,15 @@ def save_model(model, global_step: int, output_dir: str, max_checkpoints: int =
|
|||
shutil.rmtree(delete_checkpoit_dir)
|
||||
|
||||
|
||||
def get_optimizer(model, dynamic_module_names: list, static_module_names: list, lr: float, static_lr: float,
|
||||
def get_optimizer(dyamic_parameters: list[torch.nn.parameter], static_parameters: list[torch.nn.parameter], lr: float, static_lr: float,
|
||||
weight_decay: float, eps: float, adam8bit: bool):
|
||||
|
||||
all_keys = dynamic_module_names + static_module_names
|
||||
duplicated = [k for k, v in collections.Counter(all_keys).items() if v > 1]
|
||||
if len(duplicated) > 0:
|
||||
print("duplicated items:")
|
||||
for item in duplicated:
|
||||
print(item)
|
||||
raise ValueError("dynamic_module_names and or static_module_names contain duplicated paramters")
|
||||
|
||||
parameters = list()
|
||||
modules = dict(model.named_modules())
|
||||
for key in dynamic_module_names:
|
||||
parameters.extend({'params': p} for p in modules[key].parameters() if p.requires_grad)
|
||||
parameters.extend({'params': p} for p in dyamic_parameters if p.requires_grad)
|
||||
param_ids = set([id(p['params']) for p in parameters])
|
||||
for key in static_module_names:
|
||||
parameters.extend({'params': p, 'lr': static_lr} for p in modules[key].parameters() if p.requires_grad and id(p) not in param_ids)
|
||||
for p in modules[key].parameters():
|
||||
param_ids.add(id(p))
|
||||
for param in static_parameters:
|
||||
if param.requires_grad and id(param) not in param_ids:
|
||||
parameters.append({'params': param, 'lr': static_lr})
|
||||
param_ids.add(id(param))
|
||||
|
||||
if not adam8bit:
|
||||
optimizer = torch.optim.AdamW(parameters, weight_decay=weight_decay, lr=lr, eps=training_args.adam_epsilon)
|
||||
|
|
@ -204,18 +55,17 @@ def get_optimizer(model, dynamic_module_names: list, static_module_names: list,
|
|||
return optimizer
|
||||
|
||||
|
||||
def compute_dynamic_parameter_ratio(model):
|
||||
modules = dict(model.named_modules())
|
||||
active_linear_parameters = 0
|
||||
total_linear_parameters = 0
|
||||
for key in find_all_linear_module_names(model):
|
||||
active_linear_parameters += sum(p.numel() for p in modules[key].parameters() if p.requires_grad)
|
||||
total_linear_parameters += sum(p.numel() for p in modules[key].parameters())
|
||||
return math.ceil(total_linear_parameters / active_linear_parameters)
|
||||
def train(model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments):
|
||||
primary_device = torch.device(training_args.primary_device)
|
||||
secondary_device = torch.device(training_args.secondary_device)
|
||||
log_writer = tensorboard.SummaryWriter()
|
||||
|
||||
model = DyntrainModel(model_args.model_name_or_path, training_args.cache_dir, model_args.max_instant_params * 1e6, True, True)
|
||||
model = model.toDevices(primary_device, [secondary_device])
|
||||
|
||||
def prepare(model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments, primary_device: torch.device, secondary_device: torch.device) -> tuple:
|
||||
model = get_model(model_args, training_args.cache_dir, training_args.gradient_checkpointing).to(primary_device)
|
||||
paramter_count = sum(p.numel() for p in model.model.parameters())
|
||||
active_paramter_count = sum(p.numel() for p in model.model.parameters() if p.requires_grad)
|
||||
print(f"Training model with {paramter_count/1e6}m parameters and {active_paramter_count/1e6}m instantanous active paramters")
|
||||
|
||||
tokenizer = get_tokenizer(model, training_args.cache_dir, model_args)
|
||||
|
||||
|
|
@ -239,48 +89,24 @@ def prepare(model_args: ModelArguments, data_args: DataArguments, training_args:
|
|||
batch_size=training_args.per_device_train_batch_size
|
||||
) if dataset['eval_dataset'] is not None else None
|
||||
|
||||
if model_args.max_instant_params != 0:
|
||||
print(f"Target params {model_args.max_instant_params}m")
|
||||
freeze_random_modules(model, model_args.max_instant_params * 1e6,
|
||||
torch.float16 if training_args.storage_fp16 else torch.float32,
|
||||
frozen_device=primary_device, active_device=secondary_device)
|
||||
|
||||
paramter_count = sum(p.numel() for p in model.parameters())
|
||||
active_paramter_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
print(f"Training model with {paramter_count/1e6}m parameters and {active_paramter_count/1e6}m instantanous active paramters")
|
||||
|
||||
dynamic_param_ratio = compute_dynamic_parameter_ratio(model)
|
||||
print(f"dyanamic parameter ratio: 1/{dynamic_param_ratio}")
|
||||
|
||||
dynamic_param_ratio = (model.staticParamterCount() + model.dynamicParameterCount()) / model.dynamicParameterCount()
|
||||
steps_per_epoch = math.ceil(len(train_dataloader) / training_args.gradient_accumulation_steps)
|
||||
total_steps = steps_per_epoch * training_args.epochs
|
||||
|
||||
optimizer = get_optimizer(model, find_all_linear_module_names(model),
|
||||
find_all_outher_module_names(model) if training_args.train_non_linear_layers else list(),
|
||||
optimizer = get_optimizer(model.dynamicParameters(),
|
||||
model.staticParameters(),
|
||||
training_args.learning_rate,
|
||||
training_args.learning_rate / dynamic_param_ratio,
|
||||
training_args.weight_decay,
|
||||
training_args.adam_epsilon,
|
||||
training_args.adam8bit)
|
||||
|
||||
lr_scheduler = get_scheduler(
|
||||
name=training_args.lr_scheduler_type,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=training_args.warmup_steps,
|
||||
num_training_steps=total_steps
|
||||
)
|
||||
return model, optimizer, lr_scheduler, train_dataloader
|
||||
|
||||
|
||||
def train(model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments):
|
||||
primary_device = torch.device(training_args.primary_device)
|
||||
secondary_device = torch.device(training_args.secondary_device)
|
||||
log_writer = tensorboard.SummaryWriter()
|
||||
|
||||
model, optimizer, lr_scheduler, train_dataloader = prepare(model_args, data_args, training_args, primary_device, secondary_device)
|
||||
|
||||
steps_per_epoch = math.ceil(len(train_dataloader) / training_args.gradient_accumulation_steps)
|
||||
total_steps = steps_per_epoch * training_args.epochs
|
||||
dynamic_param_ratio = compute_dynamic_parameter_ratio(model)
|
||||
|
||||
if training_args.do_train:
|
||||
progress_bar = tqdm(range(total_steps))
|
||||
|
|
@ -307,13 +133,12 @@ def train(model_args: ModelArguments, data_args: DataArguments, training_args: T
|
|||
print(loss)
|
||||
|
||||
if global_step % 10 == 0 and model_args.max_instant_params != 0:
|
||||
param_count = freeze_random_modules(model, model_args.max_instant_params * 1e6,
|
||||
torch.float16 if training_args.storage_fp16 else torch.float32,
|
||||
frozen_device=primary_device,
|
||||
active_device=secondary_device)
|
||||
log_writer.add_scalar("Parameters/train", param_count, global_step)
|
||||
optimizer = get_optimizer(model, find_all_linear_module_names(model),
|
||||
find_all_outher_module_names(model) if training_args.train_non_linear_layers else list(),
|
||||
lr_scheduler.optimizer = None
|
||||
del optimizer
|
||||
model.reshuffleActive()
|
||||
log_writer.add_scalar("Parameters/train", model.activeParameterCount(), global_step)
|
||||
optimizer = get_optimizer(model.dynamicParameters(),
|
||||
model.staticParameters(),
|
||||
training_args.learning_rate,
|
||||
training_args.learning_rate / dynamic_param_ratio,
|
||||
training_args.weight_decay,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue