import contextlib import math import os from collections.abc import Generator, Iterable from datetime import timedelta import torch import torch.distributed._functional_collectives as funcol import torch.distributed.distributed_c10d as c10d from torch import distributed as dist from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor import matplotlib.pyplot as plt from datetime import datetime import torch @torch.no_grad() def clip_grad_norm_( parameters: torch.Tensor | Iterable[torch.Tensor], max_norm: float, norm_type: float = 2.0, error_if_nonfinite: bool = False, foreach: bool | None = None, pp_mesh: DeviceMesh | None = None, ) -> torch.Tensor: """ Clip the gradient norm of an iterable of parameters. Gradient norm clipping requires computing the gradient norm over the entire model. `torch.nn.utils.clip_grad_norm_` only computes gradient norm along DP/FSDP/TP dimensions. We need to manually reduce the gradient norm across PP stages. See https://github.com/pytorch/torchtitan/issues/596 for details. Args: parameters: an iterable of Tensors or a single Tensor that will have gradients normalized max_norm (float): max norm of the gradients norm_type (float): type of the used p-norm. Can be ``'inf'`` for infinity norm. error_if_nonfinite (bool): if True, an error is thrown if the total norm of the gradients from :attr:`parameters` is ``nan``, ``inf``, or ``-inf``. Default: False (will switch to True in the future) foreach (bool): use the faster foreach-based implementation. If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently fall back to the slow implementation for other device types. Default: ``None`` pp_mesh: pipeline parallel device mesh. If not None, will reduce gradient norm across PP stages. Returns: Total norm of the parameter gradients (viewed as a single vector). """ grads = [p.grad for p in parameters if p.grad is not None] total_norm = torch.nn.utils.get_total_norm( grads, norm_type, error_if_nonfinite, foreach ) # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`. # We can simply reduce the DTensor to get the total norm in this tensor's process group # and then convert it to a local tensor. # NOTE: It has two purposes: # 1. to make sure the total norm is computed correctly when PP is used (see below) # 2. to return a reduced total_norm tensor whose .item() would return the correct value if isinstance(total_norm, DTensor): # Will reach here if any non-PP parallelism is used. # If only using PP, total_norm will be a local tensor. # Remove FT replicate dimension if it exists. total_norm = total_norm.full_tensor() total_norm **= norm_type dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=pp_mesh.get_group()) total_norm **= 1.0 / norm_type torch.nn.utils.clip_grads_with_norm_(parameters, max_norm, total_norm, foreach) return total_norm def reconstruct_full_mask(mask): """ Utilities for testing and visualizing attention masks in EEG transformer models. Functions: - reconstruct_full_mask: Converts sparse block attention mask to full dense mask - visualize_attention_mask: Creates and saves attention pattern visualization """ # Get the block structure (260x260 with 1s for active blocks) block_structure = mask.to_dense()[0, 0] # 260x260 sparse block pattern full_shape = mask.shape[-1] # Get the full sequence length (33180) full_mask = torch.zeros(full_shape, full_shape, device=block_structure.device, dtype=torch.bool) # Reconstruct full mask from active blocks only block_size = mask.BLOCK_SIZE[0] # 128 active_blocks = torch.where(block_structure == 1) # Get coordinates of active blocks for q_block, kv_block in zip(active_blocks[0], active_blocks[1]): # Fill in the 128x128 regions for active blocks q_start, q_end = q_block * block_size, min((q_block + 1) * block_size, full_shape) kv_start, kv_end = kv_block * block_size, min((kv_block + 1) * block_size, full_shape) # Use mask_mod to get actual within-block attention pattern (vectorized) q_indices = torch.arange(q_start, q_end, device=block_structure.device) kv_indices = torch.arange(kv_start, kv_end, device=block_structure.device) q_grid, kv_grid = torch.meshgrid(q_indices, kv_indices, indexing='ij') # Call mask_mod with flattened arrays for efficiency block_mask = mask.mask_mod(0, 0, q_grid.flatten(), kv_grid.flatten()) block_mask = block_mask.reshape(q_end - q_start, kv_end - kv_start) full_mask[q_start:q_end, kv_start:kv_end] = block_mask return full_mask def visualize_attention_mask(mask, sample_size=5000, title_suffix=""): """ Plot the attention mask. Attentino mask needs to be constructed using reconstruct_full_mask() """ if mask is not None: # Reconstruct full mask from block structure full_mask = reconstruct_full_mask(mask) mask_2d = full_mask.cpu().numpy() # Create binary attention pattern plot plt.figure(figsize=(10, 10)) # Show sample or full mask depending on size display_mask = mask_2d[:sample_size, :sample_size] if mask_2d.shape[0] > sample_size else mask_2d plt.imshow(display_mask, cmap='Blues', aspect='equal') plt.xlabel('Key Position') plt.ylabel('Query Position') plt.title(f'Attention Mask') plt.colorbar(label='Attention Allowed') # Generate filename with timestamp and config values # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") save_path = f"figures/attention_mask/mask_{title_suffix}.png" # Create directory if it doesn't exist os.makedirs(os.path.dirname(save_path), exist_ok=True) plt.savefig(save_path, dpi=300, bbox_inches='tight') plt.close() print(f"Attention mask saved to {save_path}") def plot_random_samples_in_grid(data, num_samples=100, grid_rows=10, grid_cols=10, save_path='figures/enc_out_samples_grid.png', title='100 Random Samples from encoder output'): """ Plot 100 random samples from xxx in a 10x10 grid and save as PNG. """ random_indices = torch.randperm(data.shape[0])[:num_samples].cpu().numpy() fig, axes = plt.subplots(grid_rows, grid_cols, figsize=(20, 20)) fig.suptitle(title, fontsize=16) for idx, ax in enumerate(axes.flat): sample_idx = random_indices[idx] sample = data[sample_idx, :].float().detach().cpu().numpy() ax.plot(sample) ax.set_title(f'S{sample_idx}', fontsize=6) ax.tick_params(labelsize=4) ax.grid(True, alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close()