def init_fsdp_model_from_weights( cls, model: FullyShardedDataParallel, checkpoint: Dict[str, Any], weights_path: List[str], ): """ Load the weights of the checkpoint to the FSDP model: Take into account the type of checkpoint to decide how to perform the load (sharded or consolidated load) """ if checkpoint["type"] == CheckpointItemType.slice_list.name: SlicedCheckpointLoader.init_model_weights(model, checkpoint) elif checkpoint["type"] == CheckpointItemType.consolidated.name: weights = cls._extract_weights(checkpoint, weights_path) model.load_state_dict(weights) else: weights = cls._extract_weights(checkpoint, weights_path) model.load_local_state_dict(weights)
def init_fsdp_model_from_weights( cls, model: FullyShardedDataParallel, checkpoint: Dict[str, Any], weights_path: List[str], strict: bool = True, head_index: int = -1, ): """ Load the weights of the checkpoint to the FSDP model: - Take into account the type of checkpoint to decide on how to perform the load (sharded or consolidated load) - Takes into account the head_index (-1 if trunk else >= 0) to find the appropriate weights for the head """ if checkpoint["type"] == CheckpointItemType.slice_list.name: # Hack for checkpoints consolidated with the "layers" format # instead of the new "classy_state_dict" format: in that case # the slices are directly saved under "layers" and do not take # into account the 'weights_path' variable if "classy_state_dict" not in checkpoint: weights = checkpoint["layers"] else: weights = cls._extract_weights(checkpoint, weights_path, head_index) if weights is not None: SlicedCheckpointLoader.load_slice_state_dict(model, weights, strict=strict) else: raise ValueError( f"Could not find weights path: {weights_path}") elif checkpoint["type"] == CheckpointItemType.consolidated.name: weights = cls._extract_weights(checkpoint, weights_path, head_index) if weights is not None: out = model.load_state_dict(weights, strict=False) cls._check_load_state_dict_out(out, strict=strict) elif strict: raise ValueError( f"Could not find weights path: {weights_path}") else: weights = cls._extract_weights(checkpoint, weights_path, head_index) if weights is not None: out = model.load_local_state_dict(weights, strict=False) cls._check_load_state_dict_out(out, strict=strict) elif strict: raise ValueError( f"Could not find weights path: {weights_path}")