src/pixparse/task/task_cruller_finetune_RVLCDIP.py [190:298]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - interpolation=transforms.InterpolationMode.BICUBIC, antialias=True, ), # transforms.CenterCrop(448), # FIXME need better aspect preserving resize & pad transforms.Normalize( mean=self.img_mean, std=self.img_std, ), ] ) def train_setup( self, num_batches_per_interval: int, ): """ FIXME this interface needs refinement * currently, training duration is 'interval' based, where interval is either full dataset epoch, or sampled with replacement periods, intervals correspond to checkpoint / eval periods * LR schedule is updated per-step, so num_steps_per_interval is required to translate intervals -> total steps for scheduling * future should allow for step based durations (keeping interval as option), where train and warmup durations are specified in steps, checkpoint intervals in steps or time Args: num_batches_per_interval: Returns: """ _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) ) # FIXME currently thinking moving to device, setup DDP / FSDP makes sense # in setup here vs in __init__(). For __init__ need the model structure to # instantiate / setup tokenizer, other aspects. I don't think we need to init # weights / move to device until here. 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, ) 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 ) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - src/pixparse/task/task_cruller_finetune_docvqa.py [167:277]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - interpolation=transforms.InterpolationMode.BICUBIC, antialias=True, ), # transforms.CenterCrop(448), # FIXME need better aspect preserving resize & pad transforms.Normalize( mean=self.img_mean, std=self.img_std, ), ] ) def train_setup( self, num_batches_per_interval: int, ): """ FIXME this interface needs refinement * currently, training duration is 'interval' based, where interval is either full dataset epoch, or sampled with replacement periods, intervals correspond to checkpoint / eval periods * LR schedule is updated per-step, so num_steps_per_interval is required to translate intervals -> total steps for scheduling * future should allow for step based durations (keeping interval as option), where train and warmup durations are specified in steps, checkpoint intervals in steps or time Args: num_batches_per_interval: Returns: """ # FIXME currently thinking moving to device, setup DDP / FSDP makes sense # in setup here vs in __init__(). For __init__ need the model structure to # instantiate / setup tokenizer, other aspects. I don't think we need to init # weights / move to device until here # if self.resume: _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 ) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -