""" Utilities for checkpointing learning dynamics-related states (i.e. activations, weights, grads, etc.) We save the learning dynamics states in a subdirectory of the checkpointing directory. """ import os import re from typing import Dict, Optional import deepspeed import torch import torch.nn as nn import torch.optim as optim from datasets import Dataset from huggingface_hub import upload_folder from lightning.fabric import Fabric from lightning.fabric.strategies import DeepSpeedStrategy from lightning.fabric.utilities.rank_zero import rank_zero_only from torch.nn import functional as F from torch.utils.data import DataLoader from transformers import PreTrainedTokenizerBase from src.config import CheckpointingConfig from src.config.checkpointing_config import LearningDynamicsCheckpointingConfig from src.training.utils.initialization import initialize_model from src.training.utils.io import use_backoff # NOTE: DeepSpeed requires a dummy optimizer to be passed in to the setup function class DummyOptimizer(optim.Optimizer): def __init__(self, params): super().__init__(params, defaults={}) class CheckpointStateExtractor: """ Class to extract and save the states of a model at a given checkpoint step for learning dynamics research. """ def __init__( self, learning_dynamics_config: LearningDynamicsCheckpointingConfig, fabric: Fabric, model: nn.Module, ): self.learning_dynamics_config = learning_dynamics_config self.fabric = fabric self.model = model def extract_states(self, dataloader, compute_gradients: bool = False): """Extracts model states (activations, weights, and optionally gradients). Given a dataloader, this function will perform a forward pass of the model on each batch, and save the activations and weights at each layer. If compute_gradients is True, it will also compute the gradients of the model parameters. Args: dataloader: The dataloader containing the dataset to extract states from. compute_gradients: Whether to compute the gradients of the model parameters. Returns: A dictionary containing the activations, weights, and optionally gradients of the model. """ checkpoint_activations = {} checkpoint_weights = {} # NOTE: to extract activations and weights, we need to setup forward hooks on the layers # of the model that we are interested in. This is a good intro to forward hooks if you # are not familiar: https://web.stanford.edu/~nanbhas/blog/forward-hooks-pytorch/ forward_hooks = self._setup_forward_hooks( checkpoint_activations, checkpoint_weights, ) ######################################################## # # Forward Pass: Extract activations and weights; and compute gradients # ######################################################## for sub_batch in dataloader: _input_ids = torch.tensor(sub_batch["input_ids"], device=self.fabric.device) if compute_gradients: if "labels" in sub_batch: input_ids = _input_ids labels = torch.tensor( sub_batch["labels"], device=self.fabric.device ) else: input_ids = _input_ids[:, :-1] labels = _input_ids[:, 1:] else: input_ids = _input_ids labels = None if labels is None: # we can throw away the outputs, we are only interested in the hidden states with torch.no_grad(): _ = self.model(input_ids) else: # NOTE: if we are computing gradients, calling backwards will compute the gradients # of the model parameters. outputs, _ = self.model(input_ids) outputs = outputs.transpose(1, 2) loss = F.cross_entropy(outputs, labels) self.fabric.backward(loss, model=self.model) # cleanup forward hooks # NOTE this is not strictly necessary, since self.model is a deepcopy of the original model # but it is good practice to remove the hooks after the forward pass is complete. for hook in forward_hooks: hook.remove() ######################################################## # # Extract gradients from the target tensors of the model # ######################################################## layer_suffixes = self.learning_dynamics_config.layer_suffixes checkpoint_gradients = {} if compute_gradients: for name, param in self.model.named_parameters(): # only do this for the weight matrix of the layer_suffixes if ( any(layer_suffix in name for layer_suffix in layer_suffixes) and "weight" in name ): if isinstance(self.fabric.strategy, DeepSpeedStrategy): _grad = deepspeed.utils.safe_get_full_grad(param) else: _grad = param.grad assert _grad is not None, f"Gradient is None for layer: {name}" name = re.sub(r"\.weight", "", name) checkpoint_gradients[name] = _grad.detach().cpu() # zero out the gradients self.model.zero_grad() return checkpoint_activations, checkpoint_weights, checkpoint_gradients ######################################################## # # Setup forward hooks to save activations and weights at each layer # ######################################################## def _setup_forward_hooks(self, checkpoint_activations, checkpoint_weights): """Setup forward hooks for the model to save activations and weights at each layer. This function will setup forward hooks on the layers of the model that we are interested in. The forward hooks will save the activations and weights at each layer whenever the forward pass is performed. Args: checkpoint_activations: A dictionary to store the activations at each layer. checkpoint_weights: A dictionary to store the weights at each layer. Returns: A list of forward hooks. We do this so that we can remove the hooks after the forward pass is complete. """ forward_hooks = [] layer_suffixes = self.learning_dynamics_config.layer_suffixes for name, module in self.model.named_modules(): if any(layer_suffix in name for layer_suffix in layer_suffixes): _forward_hook = module.register_forward_hook( self._get_forward_hook( name, checkpoint_activations, checkpoint_weights ) ) forward_hooks.append(_forward_hook) return forward_hooks def _get_forward_hook( self, module_name, checkpoint_activations, checkpoint_weights ): """Get a forward hook for a given module. This function is called by the _setup_forward_hooks function to setup a forward hook for a given module. This functions is a closure that captures the module_name, checkpoint_activations, and checkpoint_weights. Args: module_name: The name of the module to setup a forward hook for. checkpoint_activations: A dictionary to store the activations at each layer. checkpoint_weights: A dictionary to store the weights at each layer. Returns: A forward hook for the given module. """ def _forward_hook(module, _, module_out): sequence_idx = self.learning_dynamics_config.sequence_idx local_activations = module_out[:, sequence_idx, :].detach() # Gather activations from all processes using fabric gathered_activations = self.fabric.all_gather(local_activations) # Reshape from [num_processes, batch_size, hidden_dim] to [total_batch_size, hidden_dim] # NOTE: transposing allows us to interleave the activations from each process so that # they are in the correct order. (i.e. activation N is from data sample N) gathered_activations = gathered_activations.transpose(0, 1).reshape( -1, gathered_activations.shape[-1] ) # check if there is already a key for the module name if module_name not in checkpoint_activations: # if there is no key, then we create a new key and store the hidden states checkpoint_activations[module_name] = ( gathered_activations.detach().cpu() ) # extract the weight matrix just once weight_matrix = module.weight.detach().cpu() checkpoint_weights[module_name] = weight_matrix else: # if there is already a key, then we concatenate the new hidden states to the existing ones checkpoint_activations[module_name] = torch.cat( ( checkpoint_activations[module_name], gathered_activations.detach().cpu(), ) ) return _forward_hook def compute_learning_dynamics_states( checkpointing_config: CheckpointingConfig, fabric: Fabric, model: nn.Module, dataset: Dataset, compute_gradients: bool = False, ) -> Dict[str, torch.Tensor]: """Computes the learning dynamics metrics for a given checkpoint step. Uses the CheckpointStateExtractor to extract the activations, weights, and optionally gradients of the model at a given checkpoint step. Args: checkpointing_config: The configuration object for checkpointing. fabric: The Fabric instance for distributed training. model: The model to extract states from. dataset: The dataset to extract states from. compute_gradients: Whether to compute the gradients of the model parameters. Returns: A dictionary containing the activations, weights, and optionally gradients of the model. """ # NOTE: Synchronizing processes for fabric dataloader setup fabric.barrier() model.to("cpu") # Offloading model to CPU # Setting up Dataloader for learning dynamics def _collate_fn(batch): return {"input_ids": [entry["input_ids"] for entry in batch]} batch_size = checkpointing_config.learning_dynamics.batch_size sub_batch_size = batch_size // fabric.world_size # NOTE: Make sure to set drop_last to False, otherwise the last batch will be dropped # and we will not have a complete set of activations for the last sample. Also, # we need to set shuffle to False, otherwise the activations will be shuffled across # processes and we will not be able to interleave them correctly. extractor_dataloader = DataLoader( dataset, batch_size=sub_batch_size, shuffle=False, collate_fn=_collate_fn, drop_last=False, ) extractor_dataloader = fabric.setup_dataloaders( extractor_dataloader, use_distributed_sampler=True ) # Create a new model instance with same parameters but zero gradients _model = initialize_model(model.config) _model.load_state_dict(model.state_dict()) if isinstance(fabric.strategy, DeepSpeedStrategy): _model, _ = fabric.setup(_model, DummyOptimizer(_model.parameters())) else: _model = fabric.setup(_model) _model.zero_grad() # setup forward hooks for the model to save activations and weights at each layer state_extractor = CheckpointStateExtractor( checkpointing_config.learning_dynamics, fabric, _model ) checkpoint_activations, checkpoint_weights, checkpoint_gradients = ( state_extractor.extract_states( extractor_dataloader, compute_gradients=compute_gradients ) ) del _model torch.cuda.empty_cache() # NOTE: Synchronizing processes for model setup fabric.barrier() model.to(fabric.device) # NOTE: Trimming down the activations to match the dataset size; # This is because the DataSampler might add extra samples to the dataset to make it evenly divisible # by the number of processes. We need to remove these extra samples. for layer_name, layer_activations in checkpoint_activations.items(): if len(layer_activations) > len(dataset): checkpoint_activations[layer_name] = layer_activations[: len(dataset)] elif len(layer_activations) < len(dataset): raise ValueError( f"Number of activations ({len(layer_activations)}) in layer {layer_name} does not match number of samples in dataset ({len(dataset)})" ) return { "activations": checkpoint_activations, "weights": checkpoint_weights, "gradients": checkpoint_gradients, } @rank_zero_only @use_backoff() def save_learning_dynamics_states( checkpointing_config: CheckpointingConfig, checkpoint_step: int, prefix: str, fabric: Fabric, learning_dynamics_states: Dict[str, torch.Tensor], learning_dynamics_dataset: Optional[Dataset] = None, tokenizer: Optional[PreTrainedTokenizerBase] = None, ) -> None: """Save the learning dynamics metrics to the checkpointing directory. By default only the learning dynamics states are saved. If the learning dynamics dataset is provided, it is also saved; if a tokenizer is provided, the dataset is also detokenized (i.e. a new column with the text is added to the dataset). The learning dynamics dataset is saved in the checkpointing directory as a HuggingFace dataset. Creates a versioned checkpoint directory with the following structure: {checkpointing_config.runs_dir}/ └── {checkpointing_config.run_name}/ └── {checkpointing_config.checkpoints_dir}/ ├── step_{checkpoint_step}/ │ └── {checkpointing_config.learning_dynamics_dir}/ # Learning Dynamics files │ ├── {prefix}_activations.pt │ ├── {prefix}_weights.pt │ └── {prefix}_gradients.pt │ └── {prefix}_data/ # if learning_dynamics_dataset is provided └── latest -> step_{checkpoint_step}/ NOTE: this function is only called on rank 0 Args: checkpointing_config: The configuration object for checkpointing. checkpoint_step: The checkpoint step at which the learning dynamics states were computed. prefix: The prefix for the learning dynamics states. fabric: The Fabric instance for distributed training. learning_dynamics_states: The learning dynamics states to save. learning_dynamics_dataset: The dataset containing learning dynamics data, including input IDs that need to be decoded. (optional) tokenizer: The tokenizer used to decode input IDs into text. (optional) """ runs_dir = checkpointing_config.runs_dir run_name = checkpointing_config.run_name checkpoints_dir = checkpointing_config.checkpoints_dir learning_dynamics_dir = checkpointing_config.learning_dynamics_dir run_path = os.path.join(runs_dir, run_name) root_checkpoint_path = os.path.join(run_path, checkpoints_dir) checkpoint_path = os.path.join(root_checkpoint_path, f"step_{checkpoint_step}") learning_dynamics_path = os.path.join(checkpoint_path, learning_dynamics_dir) os.makedirs(learning_dynamics_path, exist_ok=True) # save the learning dynamics states for key, value in learning_dynamics_states.items(): if value is not None and len(value) > 0: torch.save( value, os.path.join(learning_dynamics_path, f"{prefix}_{key}.pt") ) if learning_dynamics_dataset is not None: if tokenizer is not None: # go through dataset and decode the input ids; and add back into dataset detokenized_dataset = {"input_ids": [], "text": []} for entry in learning_dynamics_dataset: input_ids = entry["input_ids"] decoded_text = tokenizer.decode(input_ids, skip_special_tokens=True) detokenized_dataset["input_ids"].append(input_ids) detokenized_dataset["text"].append(decoded_text) learning_dynamics_dataset = Dataset.from_dict(detokenized_dataset) learning_dynamics_dataset_path = os.path.join( learning_dynamics_path, f"{prefix}_data" ) learning_dynamics_dataset.save_to_disk(learning_dynamics_dataset_path) if checkpointing_config.save_to_hf: # Upload the HF model upload_folder( folder_path=learning_dynamics_path, path_in_repo=learning_dynamics_dir, repo_id=checkpointing_config.hf_checkpoint.repo_id, commit_message=f"Saving Learning Dynamics Data ({prefix}) -- Step {checkpoint_step}", revision=checkpointing_config.run_name, token=os.getenv("HF_TOKEN"), )