src/pixparse/task/task_cruller_finetune_CORD.py [316:409]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - _logger.info(f"Resuming from existing checkpoint. ") self.state_dict = {k.replace("module.", ""): v for k, v in self.state_dict.items()} self.model.load_state_dict(self.state_dict) self.newly_added_num = self.tokenizer.trunk.add_special_tokens( {"additional_special_tokens": sorted(set(self.special_tokens_finetune))} ) self.vocab_size = len(self.tokenizer.trunk) # We resize token embeddings after initializing if self.newly_added_num > 0: self.model.text_decoder.trunk.resize_token_embeddings( len(self.tokenizer.trunk) ) device = self.device_env.device self.model.to(device) if self.device_env.world_size > 1: # NOTE: the plan is to add option for FSDP w/ HYBRID_SHARD strategy to extend # model size capacity beyond DDP w/o overloading HF cluster NCCL throughput. self.model = torch.nn.parallel.DistributedDataParallel( self.model, device_ids=[device], static_graph=True, ) self.has_no_sync = hasattr(self.model, "no_sync") opt_kwargs = {} if self.cfg.opt.betas is not None: opt_kwargs["betas"] = self.cfg.opt.betas if self.cfg.opt.momentum is not None: opt_kwargs["momentum"] = self.cfg.opt.momentum self.optimizer = create_optimizer_v2( self.model, self.cfg.opt.optimizer, lr=self.cfg.opt.learning_rate, eps=self.cfg.opt.eps, layer_decay=self.cfg.opt.layer_decay, **opt_kwargs, ) #self.optimizer = torch.optim.Adam(self.model.parameters(), lr = self.cfg.opt.learning_rate) if self.cfg.amp: self.scaler = timm.utils.NativeScaler() self.autocast = partial( torch.autocast, device_type=device.type, dtype=self.amp_dtype ) else: self.scaler = None self.autocast = nullcontext # FIXME will need two paths here to support interval vs step based durations # in either case LR is always stepped with each optimizer update (train step) self.num_steps_per_interval = ( num_batches_per_interval // self.cfg.opt.grad_accum_steps ) self.scheduler, num_scheduled_epochs = create_scheduler_v2( self.optimizer, self.cfg.opt.scheduler, warmup_lr=self.cfg.opt.warmup_learning_rate, warmup_epochs=self.num_warmup_intervals, num_epochs=self.num_intervals, step_on_epochs=False, # sched is stepped on updates updates_per_epoch=self.num_steps_per_interval, ) self.scheduler.step_update(0) def text_input_to_target(self, text_input, ignore_id=-100): target = text_input.clone() # model doesn't need to predict pad token target[target == self.tokenizer.trunk.pad_token_id] = ignore_id # model doesn't need to predict prompt (for VQA) prompt_end_token_id = self.tokenizer.trunk.convert_tokens_to_ids( self.prompt_end_token ) slice_id = torch.nonzero(target == prompt_end_token_id).sum() + 1 target[:slice_id] = ignore_id return target def collate_fn(self, batch): """ basic collator for PIL images, as returned by rvlcdip dataloader (among others) """ tokenizer_fn = lambda x: self.tokenizer.trunk( x, # FIXME move this batcher/tokenizer elsewhere add_special_tokens=False, return_tensors="pt", max_length=512, padding="max_length", truncation=True, ).input_ids[0] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - src/pixparse/task/task_cruller_finetune_docvqa.py [202:292]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - _logger.info(f"Resuming from existing checkpoint. ") self.state_dict = {k.replace("module.", ""): v for k, v in self.state_dict.items()} self.model.load_state_dict(self.state_dict) self.newly_added_num = self.tokenizer.trunk.add_special_tokens( {"additional_special_tokens": sorted(set(self.special_tokens_finetune))} ) self.vocab_size = len(self.tokenizer.trunk) # We resize token embeddings after initializing if self.newly_added_num > 0: self.model.text_decoder.trunk.resize_token_embeddings( len(self.tokenizer.trunk) ) device = self.device_env.device self.model.to(device) if self.device_env.world_size > 1: # NOTE: the plan is to add option for FSDP w/ HYBRID_SHARD strategy to extend # model size capacity beyond DDP w/o overloading HF cluster NCCL throughput. self.model = torch.nn.parallel.DistributedDataParallel( self.model, device_ids=[device], static_graph=True, ) self.has_no_sync = hasattr(self.model, "no_sync") opt_kwargs = {} if self.cfg.opt.betas is not None: opt_kwargs["betas"] = self.cfg.opt.betas if self.cfg.opt.momentum is not None: opt_kwargs["momentum"] = self.cfg.opt.momentum self.optimizer = create_optimizer_v2( self.model, self.cfg.opt.optimizer, lr=self.cfg.opt.learning_rate, eps=self.cfg.opt.eps, layer_decay=self.cfg.opt.layer_decay, **opt_kwargs, ) #self.optimizer = torch.optim.Adam(self.model.parameters(), lr = self.cfg.opt.learning_rate) if self.cfg.amp: self.scaler = timm.utils.NativeScaler() self.autocast = partial( torch.autocast, device_type=device.type, dtype=self.amp_dtype ) else: self.scaler = None self.autocast = nullcontext # FIXME will need two paths here to support interval vs step based durations # in either case LR is always stepped with each optimizer update (train step) self.num_steps_per_interval = ( num_batches_per_interval // self.cfg.opt.grad_accum_steps ) self.scheduler, num_scheduled_epochs = create_scheduler_v2( self.optimizer, self.cfg.opt.scheduler, warmup_lr=self.cfg.opt.warmup_learning_rate, warmup_epochs=self.num_warmup_intervals, num_epochs=self.num_intervals, step_on_epochs=False, # sched is stepped on updates updates_per_epoch=self.num_steps_per_interval, ) self.scheduler.step_update(0) def text_input_to_target(self, text_input, ignore_id=-100): target = text_input.clone() # model doesn't need to predict pad token target[target == self.tokenizer.trunk.pad_token_id] = ignore_id # model doesn't need to predict prompt (for VQA) prompt_end_token_id = self.tokenizer.trunk.convert_tokens_to_ids( self.prompt_end_token ) slice_id = torch.nonzero(target == prompt_end_token_id).sum() + 1 target[:slice_id] = ignore_id return target def collate_fn(self, batch): tokenizer_fn = lambda x: self.tokenizer.trunk( x, add_special_tokens=False, return_tensors="pt", max_length=512, padding="max_length", truncation=True, ).input_ids[0] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -