455 lines
17 KiB
Python
455 lines
17 KiB
Python
import os
|
|
import yaml
|
|
import torch
|
|
from torch import nn
|
|
import wandb
|
|
import json
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from torch.utils.data import Dataset, DataLoader
|
|
|
|
from torch.distributed import init_process_group, destroy_process_group, barrier
|
|
from torch.distributed.fsdp import (
|
|
FullyShardedDataParallel as FSDP,
|
|
FullStateDictConfig,
|
|
MixedPrecision,
|
|
ShardingStrategy,
|
|
StateDictType,
|
|
)
|
|
|
|
from .utils import Base, EXPECTED, EXPECTED_TRAIN
|
|
from .utils import create_folder_if_necessary, safe_save, load_or_fail
|
|
|
|
|
|
# pylint: disable=unused-argument
|
|
class WarpCore(ABC):
|
|
@dataclass(frozen=True)
|
|
class Config(Base):
|
|
experiment_id: str = EXPECTED_TRAIN
|
|
checkpoint_path: str = EXPECTED_TRAIN
|
|
output_path: str = EXPECTED_TRAIN
|
|
checkpoint_extension: str = "safetensors"
|
|
dist_file_subfolder: str = ""
|
|
allow_tf32: bool = True
|
|
|
|
wandb_project: str = None
|
|
wandb_entity: str = None
|
|
|
|
@dataclass() # not frozen, means that fields are mutable
|
|
class Info: # not inheriting from Base, because we don't want to enforce the default fields
|
|
wandb_run_id: str = None
|
|
total_steps: int = 0
|
|
iter: int = 0
|
|
|
|
@dataclass(frozen=True)
|
|
class Data(Base):
|
|
dataset: Dataset = EXPECTED
|
|
dataloader: DataLoader = EXPECTED
|
|
iterator: any = EXPECTED
|
|
|
|
@dataclass(frozen=True)
|
|
class Models(Base):
|
|
pass
|
|
|
|
@dataclass(frozen=True)
|
|
class Optimizers(Base):
|
|
pass
|
|
|
|
@dataclass(frozen=True)
|
|
class Schedulers(Base):
|
|
pass
|
|
|
|
@dataclass(frozen=True)
|
|
class Extras(Base):
|
|
pass
|
|
|
|
# ---------------------------------------
|
|
info: Info
|
|
config: Config
|
|
|
|
# FSDP stuff
|
|
fsdp_defaults = {
|
|
"sharding_strategy": ShardingStrategy.SHARD_GRAD_OP,
|
|
"cpu_offload": None,
|
|
"mixed_precision": MixedPrecision(
|
|
param_dtype=torch.bfloat16,
|
|
reduce_dtype=torch.bfloat16,
|
|
buffer_dtype=torch.bfloat16,
|
|
),
|
|
"limit_all_gathers": True,
|
|
}
|
|
fsdp_fullstate_save_policy = FullStateDictConfig(
|
|
offload_to_cpu=True, rank0_only=True
|
|
)
|
|
# ------------
|
|
|
|
# OVERRIDEABLE METHODS
|
|
|
|
# [optionally] setup extra stuff, will be called BEFORE the models & optimizers are setup
|
|
def setup_extras_pre(self) -> Extras:
|
|
return self.Extras()
|
|
|
|
# setup dataset & dataloader, return a dict contained dataser, dataloader and/or iterator
|
|
@abstractmethod
|
|
def setup_data(self, extras: Extras) -> Data:
|
|
raise NotImplementedError("This method needs to be overriden")
|
|
|
|
# return a dict with all models that are going to be used in the training
|
|
@abstractmethod
|
|
def setup_models(self, extras: Extras) -> Models:
|
|
raise NotImplementedError("This method needs to be overriden")
|
|
|
|
# return a dict with all optimizers that are going to be used in the training
|
|
@abstractmethod
|
|
def setup_optimizers(self, extras: Extras, models: Models) -> Optimizers:
|
|
raise NotImplementedError("This method needs to be overriden")
|
|
|
|
# [optionally] return a dict with all schedulers that are going to be used in the training
|
|
def setup_schedulers(
|
|
self, extras: Extras, models: Models, optimizers: Optimizers
|
|
) -> Schedulers:
|
|
return self.Schedulers()
|
|
|
|
# [optionally] setup extra stuff, will be called AFTER the models & optimizers are setup
|
|
def setup_extras_post(
|
|
self,
|
|
extras: Extras,
|
|
models: Models,
|
|
optimizers: Optimizers,
|
|
schedulers: Schedulers,
|
|
) -> Extras:
|
|
return self.Extras.from_dict(extras.to_dict())
|
|
|
|
# perform the training here
|
|
@abstractmethod
|
|
def train(
|
|
self,
|
|
data: Data,
|
|
extras: Extras,
|
|
models: Models,
|
|
optimizers: Optimizers,
|
|
schedulers: Schedulers,
|
|
):
|
|
raise NotImplementedError("This method needs to be overriden")
|
|
|
|
# ------------
|
|
|
|
def setup_info(self, full_path=None) -> Info:
|
|
if full_path is None:
|
|
full_path = (
|
|
f"{self.config.checkpoint_path}/{self.config.experiment_id}/info.json"
|
|
)
|
|
info_dict = load_or_fail(full_path, wandb_run_id=None) or {}
|
|
info_dto = self.Info(**info_dict)
|
|
if info_dto.total_steps > 0 and self.is_main_node:
|
|
print(">>> RESUMING TRAINING FROM ITER ", info_dto.total_steps)
|
|
return info_dto
|
|
|
|
def setup_config(
|
|
self, config_file_path=None, config_dict=None, training=True
|
|
) -> Config:
|
|
if config_file_path is not None:
|
|
if config_file_path.endswith(".yml") or config_file_path.endswith(".yaml"):
|
|
with open(config_file_path, "r", encoding="utf-8") as file:
|
|
loaded_config = yaml.safe_load(file)
|
|
elif config_file_path.endswith(".json"):
|
|
with open(config_file_path, "r", encoding="utf-8") as file:
|
|
loaded_config = json.load(file)
|
|
else:
|
|
raise ValueError(
|
|
"Config file must be either a .yml|.yaml or .json file"
|
|
)
|
|
return self.Config.from_dict({**loaded_config, "training": training})
|
|
if config_dict is not None:
|
|
return self.Config.from_dict({**config_dict, "training": training})
|
|
return self.Config(training=training)
|
|
|
|
def setup_ddp(self, experiment_id, single_gpu=False):
|
|
if not single_gpu:
|
|
local_rank = int(os.environ.get("SLURM_LOCALID"))
|
|
process_id = int(os.environ.get("SLURM_PROCID"))
|
|
world_size = int(os.environ.get("SLURM_NNODES")) * torch.cuda.device_count()
|
|
|
|
self.process_id = process_id
|
|
self.is_main_node = process_id == 0
|
|
self.device = torch.device(local_rank)
|
|
self.world_size = world_size
|
|
|
|
dist_file_path = f"{os.getcwd()}/{self.config.dist_file_subfolder}dist_file_{experiment_id}"
|
|
# if os.path.exists(dist_file_path) and self.is_main_node:
|
|
# os.remove(dist_file_path)
|
|
|
|
torch.cuda.set_device(local_rank)
|
|
init_process_group(
|
|
backend="nccl",
|
|
rank=process_id,
|
|
world_size=world_size,
|
|
init_method=f"file://{dist_file_path}",
|
|
)
|
|
print(f"[GPU {process_id}] READY")
|
|
else:
|
|
pass
|
|
# print("Running in single thread, DDP not enabled.")
|
|
|
|
def setup_wandb(self):
|
|
if self.is_main_node and self.config.wandb_project is not None:
|
|
self.info.wandb_run_id = self.info.wandb_run_id or wandb.util.generate_id()
|
|
wandb.init(
|
|
project=self.config.wandb_project,
|
|
entity=self.config.wandb_entity,
|
|
name=self.config.experiment_id,
|
|
id=self.info.wandb_run_id,
|
|
resume="allow",
|
|
config=self.config.to_dict(),
|
|
)
|
|
|
|
if self.info.total_steps > 0:
|
|
wandb.alert(
|
|
title=f"Training {self.info.wandb_run_id} resumed",
|
|
text=f"Training {self.info.wandb_run_id} resumed from step {self.info.total_steps}",
|
|
)
|
|
else:
|
|
wandb.alert(
|
|
title=f"Training {self.info.wandb_run_id} started",
|
|
text=f"Training {self.info.wandb_run_id} started",
|
|
)
|
|
|
|
# LOAD UTILITIES ----------
|
|
def load_model(self, model, model_id=None, full_path=None, strict=True):
|
|
# print("in line 181 load model", type(model), model_id, full_path, strict)
|
|
if model_id is not None and full_path is None:
|
|
full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{model_id}.{self.config.checkpoint_extension}"
|
|
elif full_path is None and model_id is None:
|
|
raise ValueError(
|
|
"This method expects either 'model_id' or 'full_path' to be defined"
|
|
)
|
|
|
|
checkpoint = load_or_fail(
|
|
full_path,
|
|
wandb_run_id=self.info.wandb_run_id if self.is_main_node else None,
|
|
)
|
|
if checkpoint is not None:
|
|
model.load_state_dict(checkpoint, strict=strict)
|
|
del checkpoint
|
|
|
|
return model
|
|
|
|
def load_optimizer(self, optim, optim_id=None, full_path=None, fsdp_model=None):
|
|
if optim_id is not None and full_path is None:
|
|
full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{optim_id}.pt"
|
|
elif full_path is None and optim_id is None:
|
|
raise ValueError(
|
|
"This method expects either 'optim_id' or 'full_path' to be defined"
|
|
)
|
|
|
|
checkpoint = load_or_fail(
|
|
full_path,
|
|
wandb_run_id=self.info.wandb_run_id if self.is_main_node else None,
|
|
)
|
|
if checkpoint is not None:
|
|
try:
|
|
if fsdp_model is not None:
|
|
sharded_optimizer_state_dict = (
|
|
FSDP.scatter_full_optim_state_dict( # <---- FSDP
|
|
(
|
|
checkpoint
|
|
if (
|
|
self.is_main_node
|
|
or self.fsdp_defaults["sharding_strategy"]
|
|
== ShardingStrategy.NO_SHARD
|
|
)
|
|
else None
|
|
),
|
|
fsdp_model,
|
|
)
|
|
)
|
|
optim.load_state_dict(sharded_optimizer_state_dict)
|
|
del checkpoint, sharded_optimizer_state_dict
|
|
else:
|
|
optim.load_state_dict(checkpoint)
|
|
# pylint: disable=broad-except
|
|
except Exception as e:
|
|
print("!!! Failed loading optimizer, skipping... Exception:", e)
|
|
|
|
return optim
|
|
|
|
# SAVE UTILITIES ----------
|
|
def save_info(self, info, suffix=""):
|
|
full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/info{suffix}.json"
|
|
create_folder_if_necessary(full_path)
|
|
if self.is_main_node:
|
|
safe_save(vars(self.info), full_path)
|
|
|
|
def save_model(self, model, model_id=None, full_path=None, is_fsdp=False):
|
|
if model_id is not None and full_path is None:
|
|
full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{model_id}.{self.config.checkpoint_extension}"
|
|
elif full_path is None and model_id is None:
|
|
raise ValueError(
|
|
"This method expects either 'model_id' or 'full_path' to be defined"
|
|
)
|
|
create_folder_if_necessary(full_path)
|
|
if is_fsdp:
|
|
with FSDP.summon_full_params(model):
|
|
pass
|
|
with FSDP.state_dict_type(
|
|
model, StateDictType.FULL_STATE_DICT, self.fsdp_fullstate_save_policy
|
|
):
|
|
checkpoint = model.state_dict()
|
|
if self.is_main_node:
|
|
safe_save(checkpoint, full_path)
|
|
del checkpoint
|
|
else:
|
|
if self.is_main_node:
|
|
checkpoint = model.state_dict()
|
|
safe_save(checkpoint, full_path)
|
|
del checkpoint
|
|
|
|
def save_optimizer(self, optim, optim_id=None, full_path=None, fsdp_model=None):
|
|
if optim_id is not None and full_path is None:
|
|
full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{optim_id}.pt"
|
|
elif full_path is None and optim_id is None:
|
|
raise ValueError(
|
|
"This method expects either 'optim_id' or 'full_path' to be defined"
|
|
)
|
|
create_folder_if_necessary(full_path)
|
|
if fsdp_model is not None:
|
|
optim_statedict = FSDP.full_optim_state_dict(fsdp_model, optim)
|
|
if self.is_main_node:
|
|
safe_save(optim_statedict, full_path)
|
|
del optim_statedict
|
|
else:
|
|
if self.is_main_node:
|
|
checkpoint = optim.state_dict()
|
|
safe_save(checkpoint, full_path)
|
|
del checkpoint
|
|
|
|
# -----
|
|
|
|
def __init__(
|
|
self, config_file_path=None, config_dict=None, device="cpu", training=True
|
|
):
|
|
# Temporary setup, will be overriden by setup_ddp if required
|
|
self.device = device
|
|
self.process_id = 0
|
|
self.is_main_node = True
|
|
self.world_size = 1
|
|
# ----
|
|
|
|
self.config: self.Config = self.setup_config(
|
|
config_file_path, config_dict, training
|
|
)
|
|
self.info: self.Info = self.setup_info()
|
|
|
|
def __call__(self, single_gpu=False):
|
|
self.setup_ddp(
|
|
self.config.experiment_id, single_gpu=single_gpu
|
|
) # this will change the device to the CUDA rank
|
|
self.setup_wandb()
|
|
if self.config.allow_tf32:
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
|
|
if self.is_main_node:
|
|
print()
|
|
print("**STARTIG JOB WITH CONFIG:**")
|
|
print(yaml.dump(self.config.to_dict(), default_flow_style=False))
|
|
print("------------------------------------")
|
|
print()
|
|
print("**INFO:**")
|
|
print(yaml.dump(vars(self.info), default_flow_style=False))
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
# SETUP STUFF
|
|
extras = self.setup_extras_pre()
|
|
assert extras is not None, "setup_extras_pre() must return a DTO"
|
|
|
|
data = self.setup_data(extras)
|
|
assert data is not None, "setup_data() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**DATA:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: type(v).__name__ for k, v in data.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
models = self.setup_models(extras)
|
|
assert models is not None, "setup_models() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**MODELS:**")
|
|
print(
|
|
yaml.dump(
|
|
{
|
|
k: f"{type(v).__name__} - {f'trainable params {sum(p.numel() for p in v.parameters() if p.requires_grad)}' if isinstance(v, nn.Module) else 'Not a nn.Module'}"
|
|
for k, v in models.to_dict().items()
|
|
},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
optimizers = self.setup_optimizers(extras, models)
|
|
assert optimizers is not None, "setup_optimizers() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**OPTIMIZERS:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: type(v).__name__ for k, v in optimizers.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
schedulers = self.setup_schedulers(extras, models, optimizers)
|
|
assert schedulers is not None, "setup_schedulers() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**SCHEDULERS:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: type(v).__name__ for k, v in schedulers.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
post_extras = self.setup_extras_post(extras, models, optimizers, schedulers)
|
|
assert post_extras is not None, "setup_extras_post() must return a DTO"
|
|
extras = self.Extras.from_dict({**extras.to_dict(), **post_extras.to_dict()})
|
|
if self.is_main_node:
|
|
print("**EXTRAS:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: f"{v}" for k, v in extras.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
# -------
|
|
|
|
# TRAIN
|
|
if self.is_main_node:
|
|
print("**TRAINING STARTING...**")
|
|
self.train(data, extras, models, optimizers, schedulers)
|
|
|
|
if single_gpu is False:
|
|
barrier()
|
|
destroy_process_group()
|
|
if self.is_main_node:
|
|
print()
|
|
print("------------------------------------")
|
|
print()
|
|
print("**TRAINING COMPLETE**")
|
|
if self.config.wandb_project is not None:
|
|
wandb.alert(
|
|
title=f"Training {self.info.wandb_run_id} finished",
|
|
text=f"Training {self.info.wandb_run_id} finished",
|
|
)
|