(
model,
args,
trainloader,
valloader,
logger=None,
)
| 31 | |
| 32 | |
| 33 | def block_ap( |
| 34 | model, |
| 35 | args, |
| 36 | trainloader, |
| 37 | valloader, |
| 38 | logger=None, |
| 39 | ): |
| 40 | logger.info("Starting ...") |
| 41 | if args.off_load_to_disk: |
| 42 | logger.info("offload the training dataset to disk, saving CPU memory, but may slowdown the training due to additional I/O...") |
| 43 | |
| 44 | dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 45 | use_cache = model.config.use_cache |
| 46 | model.config.use_cache = False |
| 47 | |
| 48 | # step 1: move embedding layer and first layer to target device, only suppress llama models now |
| 49 | layers = model.model.layers |
| 50 | model.model.embed_tokens = model.model.embed_tokens.to(dev) |
| 51 | model.model.norm = model.model.norm.to(dev) |
| 52 | if hasattr(model.model, 'rotary_emb'): |
| 53 | # for llama-3.1 |
| 54 | model.model.rotary_emb = model.model.rotary_emb.to(dev) |
| 55 | layers[0] = layers[0].to(dev) |
| 56 | dtype = torch.float16 |
| 57 | |
| 58 | # step 2: init dataset |
| 59 | flag = time.time() |
| 60 | if args.off_load_to_disk: |
| 61 | fp_train_cache_path = f'{args.cache_dir}/{flag}/block_training_fp_train' |
| 62 | fp_val_cache_path = f'{args.cache_dir}/{flag}/block_training_fp_val' |
| 63 | quant_train_cache_path = f'{args.cache_dir}/{flag}/block_training_quant_train' |
| 64 | quant_val_cache_path = f'{args.cache_dir}/{flag}/block_training_quant_val' |
| 65 | for path in [fp_train_cache_path,fp_val_cache_path,quant_train_cache_path,quant_val_cache_path]: |
| 66 | if os.path.exists(path): |
| 67 | shutil.rmtree(path) |
| 68 | else: |
| 69 | fp_train_cache_path = None |
| 70 | fp_val_cache_path = None |
| 71 | quant_train_cache_path = None |
| 72 | quant_val_cache_path = None |
| 73 | fp_train_inps = BlockTrainDataset(args.train_size, args.training_seqlen, |
| 74 | model.config.hidden_size, args.batch_size, dtype, cache_path=fp_train_cache_path,off_load_to_disk=args.off_load_to_disk) |
| 75 | fp_val_inps = BlockTrainDataset(args.val_size, args.training_seqlen, |
| 76 | model.config.hidden_size, args.batch_size, dtype, cache_path=fp_val_cache_path,off_load_to_disk=args.off_load_to_disk) |
| 77 | |
| 78 | # step 3: catch the input of thefirst layer |
| 79 | class Catcher(nn.Module): |
| 80 | def __init__(self, module, dataset): |
| 81 | super().__init__() |
| 82 | self.module = module |
| 83 | self.dataset = dataset |
| 84 | self.index = 0 |
| 85 | self.attention_mask = None |
| 86 | self.position_ids = None |
| 87 | |
| 88 | def forward(self, inp, **kwargs): |
| 89 | self.dataset.update_data(self.index, inp.squeeze(0).to('cpu')) |
| 90 | self.index += 1 |
no test coverage detected