File size: 9,867 Bytes
799d677 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 |
import glob
import os
import shutil
from typing import Optional
import torch.multiprocessing as mp
from nemo.core.config import hydra_runner
from nemo.utils import logging
from nemo.utils.exp_manager import exp_manager
from nemo.utils.get_rank import is_global_rank_zero
from nemo_aligner.algorithms.supervised import SupervisedTrainer
from nemo_aligner.models.nlp.gpt.gpt_sft_model import GPTSFTModel
from nemo_aligner.utils.distributed import Timer
from nemo_aligner.utils.train_script_utils import (
CustomLoggerWrapper,
add_custom_checkpoint_callback,
extract_optimizer_scheduler_from_ptl_model,
init_distributed,
init_peft,
init_using_ptl,
resolve_and_create_trainer,
retrieve_custom_trainer_state_dict,
)
from nemo_aligner.utils.utils import load_from_nemo
from omegaconf.omegaconf import OmegaConf, open_dict
from pytorch_lightning import Trainer
from dataset import LLMJPSFTDataset, build_dataloader, load_datasets
OmegaConf.register_new_resolver("multiply", lambda x, y: x * y, replace=True)
OmegaConf.register_new_resolver("int_div", lambda x, y: x // y, replace=True)
mp.set_start_method("spawn", force=True)
def get_latest_checkpoint(checkpoints_dir: str) -> Optional[str]:
if not os.path.exists(checkpoints_dir):
return None
checkpoint_dirs: list[str] = [
d for d in os.listdir(checkpoints_dir) if d.startswith("step=")
]
if not checkpoint_dirs:
return None
latest_checkpoint = max(checkpoint_dirs, key=lambda d: int(d.split("=")[1]))
return os.path.join(checkpoints_dir, latest_checkpoint)
def _modify_config(gpt_cfg, cfg, add_cfg_to_tree=False):
"""
This function modifies the original gpt pre-training config (gpt_cfg) with attributes from the finetuning config (cfg).
The `add_cfg_to_tree` arg adds `cfg` to the top of the yaml tree which is needed for all `hparams.yaml` files when passed as an arg to `load_from_checkpoint()`.
"""
OmegaConf.set_struct(gpt_cfg, True)
OmegaConf.resolve(cfg)
with open_dict(gpt_cfg):
gpt_cfg.megatron_amp_O2 = cfg.model.get("megatron_amp_O2", False)
gpt_cfg.micro_batch_size = cfg.mbs
gpt_cfg.global_batch_size = cfg.gbs
gpt_cfg.data = cfg.data
gpt_cfg.sequence_parallel = cfg.model.get("sequence_parallel", False)
gpt_cfg.activations_checkpoint_granularity = cfg.model.get(
"activations_checkpoint_granularity", None
)
gpt_cfg.activations_checkpoint_num_layers = cfg.model.get(
"activations_checkpoint_num_layers", None
)
gpt_cfg.activations_checkpoint_method = cfg.model.get(
"activations_checkpoint_method", None
)
gpt_cfg.activations_checkpoint_layers_per_pipeline = cfg.model.get(
"activations_checkpoint_layers_per_pipeline", None
)
gpt_cfg.peft = cfg.model.peft
gpt_cfg.optim = cfg.model.optim
gpt_cfg.precision = cfg.trainer.precision
gpt_cfg.restore_from_path = cfg.model.restore_from_path
gpt_cfg.resume_from_checkpoint = cfg.model.resume_from_checkpoint
gpt_cfg.save_nemo_on_validation_end = cfg.model.save_nemo_on_validation_end
gpt_cfg.gradient_as_bucket_view = cfg.model.gradient_as_bucket_view
gpt_cfg.hidden_dropout = cfg.model.get("hidden_dropout", 0.0)
gpt_cfg.attention_dropout = cfg.model.get("attention_dropout", 0.0)
gpt_cfg.ffn_dropout = cfg.model.ffn_dropout
gpt_cfg.use_flash_attention = cfg.model.get("use_flash_attention", False)
# if TP/PP size is -1, use default TP/PP size as original model
if cfg.model.get("tensor_model_parallel_size", 1) > 0:
gpt_cfg.tensor_model_parallel_size = cfg.model.get(
"tensor_model_parallel_size", 1
)
if cfg.model.get("pipeline_model_parallel_size", 1) > 0:
gpt_cfg.pipeline_model_parallel_size = cfg.model.get(
"pipeline_model_parallel_size", 1
)
gpt_cfg.pipeline_model_parallel_split_rank = cfg.model.get(
"pipeline_model_parallel_split_rank", 0
)
gpt_cfg.use_loss_mask = cfg.model.use_loss_mask
sft_cls = GPTSFTModel
gpt_cfg.target = f"{sft_cls.__module__}.{sft_cls.__name__}"
if cfg.model.get("use_flash_attention", None) is not None:
gpt_cfg.use_flash_attention = cfg.model.use_flash_attention
if cfg.model.get("seq_len_interpolation_factor", None) is not None:
gpt_cfg.seq_len_interpolation_factor = (
cfg.model.seq_len_interpolation_factor
)
# This is needed when modifying a hparam file directly to load `.ckpt` files.
# This is not needed to modify the cfg in `.nemo` files.
if add_cfg_to_tree:
OmegaConf.resolve(gpt_cfg)
gpt_cfg.cfg = gpt_cfg
return gpt_cfg
@hydra_runner(config_path="configs", config_name="sft")
def main(cfg):
if is_global_rank_zero():
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f"\n{OmegaConf.to_yaml(cfg)}")
if cfg.use_mpi:
global_rank = int(os.getenv("OMPI_COMM_WORLD_RANK", 0))
local_rank = int(os.getenv("OMPI_COMM_WORLD_LOCAL_RANK", 0))
world_size = int(os.getenv("OMPI_COMM_WORLD_SIZE", 1))
os.environ["RANK"] = str(global_rank)
os.environ["LOCAL_RANK"] = str(local_rank)
os.environ["WORLD_SIZE"] = str(world_size)
if cfg.use_slurm:
os.environ["SLURM_PROCID"] = str(global_rank)
os.environ["SLURM_LOCALID"] = str(local_rank)
os.environ["SLURM_NTASKS"] = str(world_size)
logging.info(
f"global_rank: {global_rank}, local_rank: {local_rank}, world_size: {world_size}"
)
trainer: Trainer = resolve_and_create_trainer(cfg, "sft")
log_dir = exp_manager(trainer, cfg.exp_manager)
logger = CustomLoggerWrapper(trainer.loggers)
# hydra interpolation does not work here as the interpolation key is lost when PTL saves hparams
with open_dict(cfg):
cfg.model.precision = cfg.trainer.precision
ptl_model, updated_cfg = load_from_nemo(
GPTSFTModel,
model_cfg=cfg,
trainer=trainer,
strict=True,
modify_config_fn=_modify_config,
restore_path=cfg.model.restore_from_path,
return_updated_cfg=True,
)
init_peft(ptl_model, updated_cfg)
latest_checkpoint: Optional[str] = get_latest_checkpoint(f"{log_dir}/checkpoints")
if latest_checkpoint is not None:
logging.info(f"Resuming from checkpoint: {latest_checkpoint}")
custom_trainer_state_dict = retrieve_custom_trainer_state_dict(trainer)
consumed_samples: int = custom_trainer_state_dict["consumed_samples"]
else:
logging.info("No checkpoint found. Starting from scratch.")
custom_trainer_state_dict = None
consumed_samples = 0
# save the updated config to the log directory
if is_global_rank_zero():
updated_config_path: str = f"{log_dir}/checkpoints/model_config.yaml"
os.makedirs(os.path.dirname(updated_config_path), exist_ok=True)
OmegaConf.save(updated_cfg, updated_config_path)
with open_dict(cfg):
# overwrite the model config with the config from the checkpoint
cfg.model.encoder_seq_length = ptl_model.cfg.encoder_seq_length
train_examples, dev_examples = load_datasets(cfg)
init_distributed(trainer, ptl_model, cfg.model.get("transformer_engine", False))
train_dataset = LLMJPSFTDataset(
loaded_examples=train_examples,
tokenizer=ptl_model.tokenizer,
use_loss_mask=cfg.model.use_loss_mask,
max_seq_length=cfg.model.max_seq_length,
)
val_dataset = LLMJPSFTDataset(
loaded_examples=dev_examples,
tokenizer=ptl_model.tokenizer,
use_loss_mask=cfg.model.use_loss_mask,
max_seq_length=cfg.model.max_seq_length,
)
train_dataloader = build_dataloader(
dataset=train_dataset,
consumed_samples=consumed_samples,
micro_batch_size=cfg.data.train_ds.micro_batch_size,
global_batch_size=cfg.data.train_ds.global_batch_size,
collate_fn=train_dataset.collate_fn,
seed=cfg.seed,
)
val_dataloader = build_dataloader(
dataset=val_dataset,
consumed_samples=0,
micro_batch_size=cfg.data.validation_ds.micro_batch_size,
global_batch_size=cfg.data.validation_ds.global_batch_size,
collate_fn=val_dataset.collate_fn,
)
init_using_ptl(trainer, ptl_model, train_dataloader, train_dataset)
optimizer, scheduler = extract_optimizer_scheduler_from_ptl_model(ptl_model)
ckpt_callback = add_custom_checkpoint_callback(trainer, ptl_model)
logger.log_hyperparams(OmegaConf.to_container(cfg))
timer = Timer(cfg.exp_manager.get("max_time_per_run"))
sft_trainer = SupervisedTrainer(
cfg=cfg.trainer.sft,
model=ptl_model,
optimizer=optimizer,
scheduler=scheduler,
train_dataloader=train_dataloader,
val_dataloader=val_dataloader,
test_dataloader=None,
logger=logger,
ckpt_callback=ckpt_callback,
run_timer=timer,
)
if custom_trainer_state_dict is not None:
sft_trainer.load_state_dict(custom_trainer_state_dict)
sft_trainer.fit()
# remove optimizer state files
for optimizer_state_file in glob.glob(
f"{log_dir}/checkpoints/step*/optimizer.state.*"
):
try:
shutil.rmtree(optimizer_state_file)
logging.info(f"Deleted directory: {optimizer_state_file}")
except OSError as e:
logging.error(f"Error: {optimizer_state_file} : {e.strerror}")
if __name__ == "__main__":
main()
|