import argparse import logging import os import random import math import copy import json import numpy as np import torch import torch.nn as nn import glob from tqdm.auto import tqdm, trange from torch.autograd import Variable from accelerate import Accelerator, DistributedDataParallelKwargs from accelerate.utils import InitProcessGroupKwargs from torch.utils.data import IterableDataset, DataLoader, Dataset import time import torch.distributed as dist import gc from datetime import timedelta from tokenizers import Tokenizer import wandb os.environ["WANDB_WATCH"] = "false" class BraLM(nn.Module): def __init__(self, hidden_size, use_ds=False, zero_freq_edges=None, vocab=None): super().__init__() self.hidden_size = hidden_size self.activation = nn.GELU() self.positions = nn.Parameter(torch.ones(1, 512, 1)) self.device = None # for fsdp self._tied_weights_keys = [] self.use_ds = use_ds self.zero_freq_edges = zero_freq_edges self.vocab = vocab def prepare_network(self, vocab): # Create index mappings for the flattened structure self.weight_indices = {} # Maps (s_idx, t_idx) to parameter index self.shared_param_idx = 0 # Current index for new parameters current_idx = 1 # Populate parameters and mappings for s_idx, s in enumerate(vocab.edge_dict): for t_idx, t in enumerate(vocab.edge_dict[s]): if self.zero_freq_edges is not None and t in self.zero_freq_edges[s]: # Use shared parameters self.weight_indices[(s_idx, t_idx)] = self.shared_param_idx else: self.weight_indices[(s_idx, t_idx)] = current_idx current_idx += 1 # Create new parameters self.weights = nn.Parameter(torch.randn(current_idx, self.hidden_size, self.hidden_size).uniform_(-0.5, 0.5)) self.biases = nn.Parameter(torch.randn(current_idx, 1, self.hidden_size).uniform_(-0.5, 0.5)) self.node_bias = nn.Parameter(torch.randn(len(vocab.edge_dict), 1, self.hidden_size).uniform_(-0.5, 0.5)) def to_device(self, device): self.weights.to(device) self.biases.to(device) self.positions.data = self.positions.data.to(device) self.device = device @staticmethod def _reshape12(x): return x.reshape(-1, x.size(-2), x.size(-1)) def get_positional_encoding(self, seq_len, d_model): position = torch.arange(0, seq_len).reshape(-1, 1) div_term = 10000.0 ** (torch.arange(0, d_model, 2) / d_model) position_encoding = torch.zeros(seq_len, d_model) position_encoding[:, 0::2] = torch.sin(position * div_term) position_encoding[:, 1::2] = torch.cos(position * div_term) return position_encoding.unsqueeze(0).to(self.device) # def get_initial_tensor(self, batch_size, max_norm=1.0): # # initialize energy_tensor # energy_tensor = torch.zeros(batch_size, 1, self.hidden_size).normal_(0, 1).to(self.device) # delta_norm = torch.norm(energy_tensor.view(energy_tensor.shape[0], -1), dim=-1, p="fro").detach() # clip_mask = (delta_norm > max_norm).to(energy_tensor) # clip_weights = max_norm / delta_norm * clip_mask + (1 - clip_mask) # energy_tensor = (energy_tensor * clip_weights.view(-1, 1, 1)).detach() #(bs, 1, hs) # return energy_tensor def get_initial_tensor(self, batch_size, d, pe): # initialize energy_tensor energy_tensor = torch.ones(batch_size, 1, self.hidden_size) / self.hidden_size #(bs, 1, hs) energy_tensor = energy_tensor.to(self.device) node_bias = self.node_bias[d[:, 0, 0]] energy_tensor = self.activation(energy_tensor + node_bias + Variable(pe[:,0], requires_grad=False)) return energy_tensor def forward(self, neighbor_ids): # neighbor_ids: (bs, sen_len, 1+k, 2) ; k is the number of negative samples batch_size = neighbor_ids.size(0) loss = 0 pe = self.get_positional_encoding(512, self.hidden_size) #(1, 512, hs) for i in range(neighbor_ids.size(1)): d = neighbor_ids[:, i] #(bs, 1+k, 2) if i == 0: # for the first token, initialize energy_tensor as an all-one tensor energy_tensor = self.get_initial_tensor(batch_size, d, pe) #(bs, 1, hs) else: energy_tensor = (energy_cache * self.positions[:, :i, :].softmax(1)).sum(1, keepdim=True) #(bs, 1, hs) :fix dim bug # Vectorized parameter lookup src_idx = d[..., 0] # (bs, 1+k) tgt_idx = d[..., 1] # (bs, 1+k) param_indices = torch.tensor([self.weight_indices.get((s.item(), t.item()), self.shared_param_idx) for s, t in zip(src_idx.reshape(-1), tgt_idx.reshape(-1))], device=self.device).reshape(batch_size, -1) # (bs, 1+k) # Batch gather operation w = self.weights[param_indices] # (bs, 1+k, hidden_size, hidden_size) b = self.biases[param_indices] # (bs, 1+k, 1, hidden_size) expand_energy_tensor = self._reshape12(energy_tensor.unsqueeze(1).repeat(1, w.size(1), 1, 1)) #(bs*(1+k), 1, hs) # for deepspeed fp16: expand_energy_tensor.half() if self.use_ds: expand_energy_tensor = expand_energy_tensor.half() nxt_energy_tensor = self.activation(expand_energy_tensor.bmm(self._reshape12(w))+self._reshape12(b)+Variable(pe[:,i+1], requires_grad=False)) #(bs*(1+k), 1, hs) output_tensor = nxt_energy_tensor.reshape(batch_size, -1, nxt_energy_tensor.size(-2), nxt_energy_tensor.size(-1)) #(bs, 1+k, 1, hs) if i == 0: energy_cache = output_tensor[:,0] #(bs, 1, hs) else: energy_cache = torch.cat([energy_cache, output_tensor[:,0]], dim=1) #(bs, i+1, hs) if 1: energy = output_tensor.norm(2, (-2, -1)) label = torch.LongTensor([0 for _ in range(batch_size)]).to(self.device) loss += nn.CrossEntropyLoss()(energy, label) return loss / neighbor_ids.size(1) def decode(self, start, vocab, max_new_tokens=16, do_sample=False, temperature=1): ret = [] pe = self.get_positional_encoding(512, self.hidden_size) for i, pair in enumerate(start): if i == 0: energy_tensor = self.get_initial_tensor(batch_size=1, d=torch.tensor([[pair]], device=self.device), pe=pe).squeeze(0) else: energy_tensor = (energy_cache * self.positions[:, :i, :].softmax(1)).sum(1, keepdim=True).squeeze(0) # Get parameter index for this edge param_idx = self.weight_indices.get((pair[0], pair[1]), self.shared_param_idx) # Get weights and biases using parameter index w = self.weights[param_idx].to(self.device) b = self.biases[param_idx].to(self.device) energy_tensor = self.activation(energy_tensor.mm(w) + b + pe.squeeze(0)[i]) if i == 0: energy_cache = energy_tensor.unsqueeze(0) # Add batch dimension else: energy_cache = torch.cat([energy_cache, energy_tensor.unsqueeze(0)], dim=1) ret += [pair] x = pair[1] prev_i = len(start) for i in range(max_new_tokens): candidates = vocab(vocab.get_neighbor_of_node(x, -1)) # Get parameter indices for all candidates param_indices = torch.tensor([self.weight_indices.get((x, t[1]), self.shared_param_idx) for t in candidates], device=self.device) # Get weights and biases for all candidates all_w = self.weights[param_indices].to(self.device) all_b = self.biases[param_indices].to(self.device) curr_i = prev_i + i energy_tensor = (energy_cache * self.positions[:, :curr_i, :].softmax(1)).sum(1, keepdim=True) expand_energy_tensor = energy_tensor.unsqueeze(1).repeat(1, all_w.size(0), 1, 1) expand_energy_tensor = self._reshape12(expand_energy_tensor) nxt_energy_tensor = self.activation(expand_energy_tensor.bmm(self._reshape12(all_w)) + self._reshape12(all_b) + pe[:,curr_i].unsqueeze(0)) output_tensor = nxt_energy_tensor.reshape(1, -1, nxt_energy_tensor.size(-2), nxt_energy_tensor.size(-1)) energy = output_tensor.norm(2, (-2,-1)).squeeze() probs = torch.softmax(energy, dim=-1) if temperature > 0: probs = probs / temperature if do_sample: index = torch.multinomial(probs, 1).item() else: index = probs.argmax(-1).item() y = candidates[index][-1] ret += [(x, y)] energy_tensor = output_tensor[0, index] x = y energy_cache = torch.cat([energy_cache, energy_tensor.unsqueeze(0)], dim=1) return ret class Vocab: def __init__(self, node_dict, nodeindex_dict, edge_dict, edge_decode_dict): self.node_dict = node_dict #{'node_p': index_p} ---- size: num_nodes self.nodeindex_dict = nodeindex_dict #{index_p: 'node_p'} ---- size: num_nodes self.edge_dict = edge_dict #{'node_p': {'node_q': (index_p, index_q), 'node_m': (index_p, index_m)},...} ---- size: num_nodes self.edge_decode_dict = edge_decode_dict #{(index_p, index_q): 'node_p->node_q'} ---- size: num_nodes*num_nodes def __call__(self, x): if isinstance(x, list): return [self.__call__(_) for _ in x] else: return self.fetch(x) def fetch(self, x): s, t = x.split("->") return self.edge_dict[s][t] if s in self.edge_dict and t in self.edge_dict[s] else self.edge_dict[""][""] @classmethod def from_node_dict(cls, dictname): node_dict = dict() nodeindex_dict = dict() edge_dict = dict() edge_decode_dict = dict() for s in dictname: node_dict[s] = dictname[s] nodeindex_dict[dictname[s]] = s # nodeindex_dict: {index_p: 'node_p'} edge_dict[s] = {} # edge_dict: {'node_p': {'node_q': (index_p, index_q), 'node_m': (index_p, index_m)}} for t in dictname: edge_dict[s][t] = (dictname[s], dictname[t]) edge_decode_dict[(dictname[s], dictname[t])] = "->".join([s, t]) return cls(node_dict, nodeindex_dict, edge_dict, edge_decode_dict) @classmethod def from_edge(cls, filename): edge_dict = dict() edge_dict[""] = {} edge_dict[""][""] = (0, 0) edge_decode_dict = dict() with open(filename) as f: for line in f: # line: node_p->node_q s, t = line.strip().split("->") if s not in edge_dict: i = len(edge_dict) j = 0 edge_dict[s] = dict() else: i = edge_dict[s][list(edge_dict[s].keys())[0]][0] j = len(edge_dict[s]) edge_dict[s][t] = (i, j) edge_decode_dict[(i, j)] = "->".join([s, t]) return cls(None, edge_dict, edge_decode_dict) def get_neighbor_of_edge(self, key, k, frequency_dict=None): s, t = key.split("->") # s, t: node _s = s if s in self.edge_dict else "" # if s in self.edge_dict: # ret = ["->".join([s, _t]) for _t in self.edge_dict[s].keys() if _t != t] # else: # ret = ["->".join([s, _t]) for _t in self.edge_dict[""].keys() if _t != t] # ret = ["->".join([s, _t]) for _t in self.edge_dict[s].keys() if _t != t] # select by word_frequency if frequency_dict: frequency_lst = list(frequency_dict[_s].keys()) # index = frequency_lst.index(t) # half = k // 2 # if index <= k: # t_lst = [x for i, x in enumerate(frequency_lst[:k+1]) if i != index] # else: # t_lst = frequency_lst[:half] + frequency_lst[index-half:index] t_lst = [x for i, x in enumerate(frequency_lst[:k+1]) if x != t][:k] ret = ["->".join([_s, _t]) for _t in t_lst] random.shuffle(ret) return ret # randomly select k negative samples else: ret = ["->".join([_s, _t]) for _t in self.edge_dict[_s].keys() if _t != t] random.shuffle(ret) return ret[:k] if k != -1 else ret def get_neighbor_of_node(self, key, k): #key :index s = self.nodeindex_dict[key] #node #_t: node ret = ["->".join([s, _t]) for _t in self.edge_dict[s].keys() if _t != s] # randomly select k negative samples random.shuffle(ret) return ret[:k] if k != -1 else ret def get_neighbor_of_edge_broadcast(self, key, edges, k=100): s, t = key.split("->") _ret = [_t for _t in self.edge_dict[s].keys() if _t != t] # all neighbors of s except t random.shuffle(_ret) ret = [] for edge in edges: s, t = edge.split("->") ret += [["->".join([s, _t]) for _t in _ret[:k]]] return ret @staticmethod def to_path(tokens): path = [] for left, right in zip(tokens[:-1], tokens[1:]): path.append("->".join([left, right])) return path def get_edge_of_node(self, key): return list(self.edge_dict[key].values()) def decode(self, x): return self.edge_decode_dict[x] logging.basicConfig(format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO) logger = logging.getLogger(__name__) def stdf(string): def _h(char): inside_code = ord(char) if inside_code == 0x3000: inside_code = 0x0020 else: inside_code -= 0xfee0 if inside_code < 0x0020 or inside_code > 0x7e: return char return chr(inside_code) return "".join([_h(char) for char in string]) class WikiDataset(Dataset): """ Processor for wiki data. """ def __init__(self, filename, vocab, max_seq_length, num_neg_samples, seed, buffer_size=100000, shuffle=True, use_frequency=False, use_bpe=False, bpe_tokenizer=None): super().__init__() self.vocab = vocab self.max_seq_length = max_seq_length self.num_neg_samples = num_neg_samples self.generator = np.random.default_rng(seed=seed) self.use_bpe = use_bpe self.bpe_tokenizer = bpe_tokenizer self.data = self.read(filename) if use_frequency: freq_file = 'word_frequency_en.json' if use_bpe else 'word_frequency.json' with open(freq_file, 'r') as f: self.frequency_dict = json.load(f) else: self.frequency_dict = None def read(self, filename): lines = [] with open(filename, "r", encoding="utf-8") as f: for line in f: if self.use_bpe: lines.append(line.strip()) else: src = list(line.strip()[:self.max_seq_length]) lines.append(src) return lines def __len__(self): return len(self.data) def __getitem__(self, idx): src = self.data[idx] return self.vectorize(src) def vectorize(self, src): if self.use_bpe: # For English with BPE bpe_tokens = self.bpe_tokenizer.encode(src).tokens # Truncate/pad pad_token = "[PAD]" if len(bpe_tokens) > self.max_seq_length: bpe_tokens = bpe_tokens[:self.max_seq_length] else: bpe_tokens.extend(pad_token for _ in range(self.max_seq_length - len(bpe_tokens))) tokens = bpe_tokens else: # For Chinese without BPE if len(src) > self.max_seq_length: src = src[:self.max_seq_length] else: src.extend("" for _ in range(self.max_seq_length-len(src))) tokens = src edges = self.vocab.to_path(tokens) edge_ids = self.vocab(edges) edge_ids = edge_ids[:self.max_seq_length] neighbor_ids = [self.vocab(self.vocab.get_neighbor_of_edge(e, self.num_neg_samples, self.frequency_dict)) for e in edges] new_neighbor_ids = [] for i, e_ids in enumerate(edge_ids): new_neighbor_ids.append([e_ids] + neighbor_ids[i]) return torch.LongTensor(new_neighbor_ids) def main(): parser = argparse.ArgumentParser() # Data config parser.add_argument("--data_dir", type=str, default="data/wiki", help="Directory to contain the input data for all tasks.") parser.add_argument("--output_dir", type=str, default="model/", help="Directory to output predictions and checkpoints.") parser.add_argument("--load_state_dict", type=str, default=None, help="Trained model weights to load for evaluation if needed.") # Training config parser.add_argument("--do_train", action="store_true", help="Whether to run training.") parser.add_argument("--do_eval", action="store_true", help="Whether to evaluate on the dev set.") parser.add_argument("--num_neg_samples", type=int, default=100, help="Number of negative samples.") parser.add_argument("--max_seq_length", type=int, default=128, help="Maximum total input sequence length after word-piece tokenization.") parser.add_argument("--train_batch_size", type=int, default=128, help="Total batch size for training.") parser.add_argument("--eval_batch_size", type=int, default=128, help="Total batch size for evaluation.") parser.add_argument("--learning_rate", type=float, default=5e-5, help="Initial learning rate for Adam.") parser.add_argument("--num_train_epochs", type=float, default=3.0, help="Total number of training epochs to perform.") parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides training epochs.") parser.add_argument("--weight_decay", type=float, default=0., help="L2 weight decay for training.") parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward pass.") parser.add_argument("--no_cuda", action="store_true", help="Whether not to use CUDA when available.") parser.add_argument("--fp16", action="store_true", help="Whether to use mixed precision.") parser.add_argument("--seed", type=int, default=42, help="Random seed for initialization.") parser.add_argument("--save_steps", type=int, default=500, help="How many steps to save the checkpoint once.") parser.add_argument("--hidden_size", type=int, default=32, help="Mask rate for masked-fine-tuning.") parser.add_argument("--local_rank", type=int) parser.add_argument("--initial_file_number", type=int, default=0, help="From which file to begin training.") parser.add_argument("--end_file_number", type=int, default=0, help="End file number for training.") parser.add_argument("--wiki_sorted_size", type=int, default=70, help="Total file numbers for sorted wikidata.") parser.add_argument("--run_name", type=str, default="plusb_pluspe_order", help="Run name for wandb.") parser.add_argument("--use_frequency", action="store_true", help="Whether to use word frequency.") parser.add_argument("--train_full", type=str, default=None, help="Path to train on full text.") parser.add_argument("--checkpoint_save_step", type=int, default=0, help="Interval to save checkpoint.(Only support when train_full is True)") parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="Path to checkpoint to resume training from") parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading.") parser.add_argument("--vocab_path", type=str, default="vocab_wiki_4k.json", help="Path to vocab file.") parser.add_argument("--use_ds", action="store_true", help="Whether to use deepspeed.") parser.add_argument("--sparse", action="store_true", help="Whether to use sparse.") parser.add_argument("--use_bpe", action="store_true", help="Whether to use BPE tokenizer for English.") parser.add_argument("--bpe_tokenizer_path", type=str, default="wiki_bpe_tokenizer_4000_bytelevel.json", help="Path to BPE tokenizer file.") args = parser.parse_args() device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") n_gpu = torch.cuda.device_count() logger.info("device: {}, n_gpu: {}, distributed training: {}, 16-bits training: {}".format( device, n_gpu, "-accelerate", args.fp16)) args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if n_gpu > 0: torch.cuda.manual_seed_all(args.seed) # if not os.path.exists(args.output_dir): # os.makedirs(args.output_dir) with open(args.vocab_path) as f: node_dict = json.load(f) vocab = Vocab.from_node_dict(node_dict) if args.sparse: with open('word_frequency.json', 'r') as f: freq_dict = json.load(f) zero_freq_edges = {} for s in freq_dict: zero_freq_edges[s] = [] for t in freq_dict[s]: if freq_dict[s][t] == 0: zero_freq_edges[s].append(t) else: zero_freq_edges = None def stat_cuda(epoch, cur_file_num, step, location): if accelerator.is_local_main_process: with open("cuda_stat.txt", "a") as f: if epoch is not None: f.write('epoch: %d, cur_file_num: %d, step: %d\n' % (epoch, cur_file_num, step)) f.write(f'--{location}\n') f.write('allocated: %dG, max allocated: %dG, cached: %dG, max cached: %dG\n' % ( torch.cuda.memory_allocated() / 1024 / 1024 / 1024, torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024, torch.cuda.memory_reserved() / 1024 / 1024 / 1024, torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024 )) if args.do_train: # training arguments os.environ["NCCL_DEBUG"] = "WARN" os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1" ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) init_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=1080000)) accelerator = Accelerator(kwargs_handlers=[ddp_kwargs, init_kwargs], cpu=args.no_cuda, mixed_precision="fp16" if args.fp16 else "no") device = accelerator.device # prepare model model = BraLM(args.hidden_size, args.use_ds, zero_freq_edges, vocab=vocab) model.prepare_network(vocab) # model.shared_weight.requires_grad = False # model.shared_bias.requires_grad = False # load model from checkpoint if args.load_state_dict: print(f"Loading model from checkpoint: {args.load_state_dict}") checkpoint = torch.load(args.load_state_dict, map_location="cpu") #model.load_state_dict(checkpoint["model_state_dict"]) model.load_old(checkpoint["model_state_dict"]) # Load checkpoint if specified wandb_id = None global_step = 0 if args.resume_from_checkpoint: print(f"Resuming from checkpoint: {args.resume_from_checkpoint}") checkpoint = torch.load(args.resume_from_checkpoint, map_location="cpu") model.load_state_dict(checkpoint["model_state_dict"]) #optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) start_epoch = checkpoint["epoch"] # + 1 global_step = checkpoint.get("global_step", 0) # Get saved global step wandb_id = checkpoint.get("wandb_id") else: start_epoch = 0 # if accelerator.is_local_main_process: # for name, param in model.named_parameters(): # print(name) model.to_device(device) if accelerator.is_local_main_process: print(f"start_epoch: {start_epoch}, global_step: {global_step}") # prepare optimizer no_decay = ["bias", "LayerNorm.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": args.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 } ] optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate) if args.resume_from_checkpoint: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) if accelerator.is_local_main_process: print(f"before prepare") #input('-' * 10) #stat_cuda(None, None, None, "before prepare") #print(f"{accelerator.device}, model: {model.weights.device}, tensor: {model.tensor.device}, pe: {model.positions.device}") if not args.use_ds: model, optimizer = accelerator.prepare(model, optimizer) # for deepspeed: # this line #stat_cuda(None, None, None, "after prepare") #print(f"{accelerator.device}, model: {model.module.weights.device}, tensor: {model.module.tensor.device}, pe: {model.module.positions.device}") if accelerator.is_local_main_process: print(f"after prepare") if args.do_train: if accelerator.is_local_main_process: # init wandb wandb.init( project="brain", name=args.run_name, id=wandb_id, # 如果有之前的run_id，使用它；否则会创建新的 resume="allow", # "allow"表示如果有id就恢复，没有就创建新的 config=vars(args) ) wandb.define_metric("custom_step") wandb.define_metric("batch_*", step_metric="custom_step") wandb.define_metric("epoch") wandb.define_metric("epoch_*", step_metric="epoch") print(f"Started wandb run with id: {wandb.run.id}") print(f"View run at: {wandb.run.get_url()}") if args.train_full: cur_file_num = args.train_full cur_filename = f"{cur_file_num}.txt" if args.use_bpe: with open(args.bpe_tokenizer_path, 'r') as f: bpe_tokenizer = json.load(f) else: bpe_tokenizer = None dataset = WikiDataset( os.path.join(args.data_dir, cur_filename), vocab, args.max_seq_length, args.num_neg_samples, seed=args.seed, shuffle=True, use_frequency=args.use_frequency, use_bpe=args.use_bpe, bpe_tokenizer=bpe_tokenizer ) train_dataloader = DataLoader(dataset, batch_size=args.train_batch_size, num_workers=args.num_workers, pin_memory=True) train_dataloader = accelerator.prepare(train_dataloader) elif args.resume_from_checkpoint: cur_file_num = checkpoint["cur_file_num"] if isinstance(cur_file_num, int) or cur_file_num.isdigit(): cur_file_num = int(cur_file_num) + 1 #start_epoch = start_epoch - 1 else: cur_file_num = args.initial_file_number if args.resume_from_checkpoint and global_step > 0: if args.train_full and global_step % len(train_dataloader) == 0: start_epoch = start_epoch + 1 if not args.train_full and cur_file_num > args.end_file_number: start_epoch = start_epoch + 1 cur_file_num = args.initial_file_number for epoch in trange(start_epoch, int(args.num_train_epochs), desc="Epoch"): # traverse all wiki files if epoch != start_epoch or args.train_full: cur_file_num = args.initial_file_number while cur_file_num <= args.wiki_sorted_size: if args.train_full: cur_file_num = args.train_full logger.info("***** Running training for wiki = %s *****", cur_file_num) logger.info(" Batch size = %d", args.train_batch_size * accelerator.num_processes) # prepare data if not args.train_full: cur_filename = f"{cur_file_num}.txt" if args.use_bpe: with open(args.bpe_tokenizer_path, 'r') as f: bpe_tokenizer = json.load(f) else: bpe_tokenizer = None dataset = WikiDataset( os.path.join(args.data_dir, cur_filename), vocab, args.max_seq_length, args.num_neg_samples, seed=args.seed, shuffle=True, use_frequency=args.use_frequency, use_bpe=args.use_bpe, bpe_tokenizer=bpe_tokenizer ) train_dataloader = DataLoader(dataset, batch_size=args.train_batch_size, num_workers=args.num_workers, pin_memory=True) if not args.use_ds: train_dataloader = accelerator.prepare(train_dataloader) else: model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader) # for deepspeed # training train_loss = 0 num_train_examples = 0 if accelerator.is_local_main_process: progress_bar = tqdm(train_dataloader, desc="Iteration") # start_time = time.time() #for _ in range(3): for step, batch in enumerate(train_dataloader, start=global_step % len(train_dataloader)): # batch: (bs, sen_len, 1+k, 2) batch_train_loss = 0 batch_num_train_examples = 0 #for ind in range(2, batch.size(1)): for ind in range(batch.size(1) - 1, batch.size(1)): # fix: only use the sen_len-1 # ind: 2, 3, ..., sen_len-1 # if accelerator.is_local_main_process: # end_time = time.time() # step_time = end_time - start_time # logger.info(f"Step training time: {step_time:.2f} seconds") model.train() neighbor_ids = batch[:, :ind] #(bs, ind, 1+k, 2) #stat_cuda(epoch, cur_file_num, global_step, "before forward") outputs = model(neighbor_ids) loss = outputs # if n_gpu > 1: # loss = loss.mean() if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps accelerator.backward(loss) if n_gpu > 1: dist.all_reduce(loss) loss = loss / dist.get_world_size() train_loss += loss.detach().item() batch_train_loss += loss.detach().item() num_train_examples += 1 batch_num_train_examples += 1 del outputs del loss del neighbor_ids gc.collect() # if step % 5 == 0: # torch.cuda.empty_cache() if (step + 1) % args.gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() ## modified ppl = math.exp(batch_train_loss / batch_num_train_examples) if accelerator.is_local_main_process: progress_bar.update(1) progress_bar.set_postfix(loss=batch_train_loss / batch_num_train_examples, perplexity=ppl) wandb.log({ "batch_loss": batch_train_loss / batch_num_train_examples, "batch_perplexity": math.exp(batch_train_loss / batch_num_train_examples), "batch_epoch": epoch, #"step": global_step, "custom_step": global_step })#, step=global_step) global_step += 1 # Save checkpoint every checkpoint_save_step steps at the end of each step if accelerator.is_local_main_process and args.checkpoint_save_step > 0 and global_step % args.checkpoint_save_step == 0: output_dir_f = f"{args.output_dir}/HS{args.hidden_size}/step_{global_step}/" if not os.path.exists(output_dir_f): os.makedirs(output_dir_f) output_model_file = os.path.join(output_dir_f, f"checkpoint_{global_step}.bin") model_to_save = model.module if hasattr(model, "module") else model checkpoint = { "model_state_dict": model_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "epoch": epoch, "global_step": global_step, "args": vars(args), "wandb_id": wandb.run.id } if not args.train_full: checkpoint["cur_file_num"] = cur_file_num print(f"Saving checkpoint to {output_model_file}") torch.save(checkpoint, output_model_file) print(f"Checkpoint saved to {output_model_file}") # save model for current training file if accelerator.is_local_main_process: epoch_avg_loss = train_loss / num_train_examples epoch_ppl = math.exp(epoch_avg_loss) wandb.log({ "epoch_loss": epoch_avg_loss, "epoch_perplexity": epoch_ppl, "epoch": epoch, })#, step=global_step) model_to_save = model.module if hasattr(model, "module") else model output_dir_f = f"{args.output_dir}/HS{args.hidden_size}/EPOCH{epoch}/" if not os.path.exists(output_dir_f): os.makedirs(output_dir_f) output_model_file = os.path.join(output_dir_f, "f{}_pytorch_model.bin".format(cur_file_num)) # only save the last model if args.train_full or cur_file_num == args.end_file_number: #torch.save(model_to_save.state_dict(), output_model_file) checkpoint = { "model_state_dict": model_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "epoch": epoch, "global_step": global_step, # Save global step "args": vars(args), "wandb_id": wandb.run.id # 保存当前运行的wandb_id } if not args.train_full: checkpoint["cur_file_num"] = cur_file_num print(f"Saving model to {output_model_file}") torch.save(checkpoint, output_model_file) print(f"Model saved to {output_model_file}") if args.train_full: break cur_file_num += 1 if cur_file_num > args.end_file_number: break if __name__ == "__main__": main()