Soul-AILab/SoulX-Singer
185
1# coding: utf-82__author__ = 'Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/'3 4import argparse5import numpy as np6import torch7import torch.nn as nn8from ml_collections import ConfigDict9from torch.optim import Adam, AdamW, SGD, RAdam, RMSprop10from tqdm.auto import tqdm11from typing import Dict, List, Tuple, Any, Union, Optional12import loralib as lora13from .muon import SingleDeviceMuonWithAuxAdam14import torch.distributed as dist15 16def demix(17 config: ConfigDict,18 model: torch.nn.Module,19 mix: torch.Tensor,20 device: torch.device,21 model_type: str,22 pbar: bool = False23) -> Union[Dict[str, np.ndarray], np.ndarray]:24 """25 Perform audio source separation with a given model.26 27 Supports both Demucs-specific and generic processing modes, including28 overlapping chunk-based inference with optional progress bar display.29 Handles padding, fading, and batching to reduce artifacts during separation.30 31 Args:32 config (ConfigDict): Configuration object with audio and inference33 parameters (chunk size, overlap, batch size, etc.).34 model (torch.nn.Module): Source separation model for inference.35 mix (torch.Tensor): Input audio tensor of shape (channels, time).36 device (torch.device): Device on which to run inference (CPU or CUDA).37 model_type (str): Type of model (e.g., 'htdemucs', 'mdx23c') that38 determines processing mode.39 pbar (bool, optional): If True, show a progress bar during chunk40 processing. Defaults to False.41 42 Returns:43 Union[Dict[str, np.ndarray], np.ndarray]:44 - Dictionary mapping instrument names to separated waveforms if45 multiple instruments are predicted.46 - NumPy array of separated audio if only a single instrument is47 present (Demucs mode).48 """49 50 should_print = not dist.is_initialized() or dist.get_rank() == 051 52 mix = torch.tensor(mix, dtype=torch.float32)53 54 if model_type == 'htdemucs':55 mode = 'demucs'56 else:57 mode = 'generic'58 # Define processing parameters based on the mode59 if mode == 'demucs':60 chunk_size = config.training.samplerate * config.training.segment61 num_instruments = len(config.training.instruments)62 num_overlap = config.inference.num_overlap63 step = chunk_size // num_overlap64 else:65 if 'chunk_size' in config.inference:66 chunk_size = config.inference.chunk_size67 else:68 chunk_size = config.audio.chunk_size69 num_instruments = len(prefer_target_instrument(config))70 num_overlap = config.inference.num_overlap71 72 fade_size = chunk_size // 1073 step = chunk_size // num_overlap74 border = chunk_size - step75 length_init = mix.shape[-1]76 windowing_array = _getWindowingArray(chunk_size, fade_size)77 # Add padding for generic mode to handle edge artifacts78 if length_init > 2 * border and border > 0:79 mix = nn.functional.pad(mix, (border, border), mode="reflect")80 81 batch_size = config.inference.batch_size82 83 use_amp = getattr(config.training, 'use_amp', True)84 85 with torch.cuda.amp.autocast(enabled=use_amp):86 with torch.inference_mode():87 # Initialize result and counter tensors88 req_shape = (num_instruments,) + mix.shape89 result = torch.zeros(req_shape, dtype=torch.float32)90 counter = torch.zeros(req_shape, dtype=torch.float32)91 92 i = 093 batch_data = []94 batch_locations = []95 if pbar and should_print:96 progress_bar = tqdm(97 total=mix.shape[1], desc="Processing audio chunks", leave=False98 )99 else:100 progress_bar = None101 102 while i < mix.shape[1]:103 # Extract chunk and apply padding if necessary104 part = mix[:, i:i + chunk_size].to(device)105 chunk_len = part.shape[-1]106 if mode == "generic" and chunk_len > chunk_size // 2:107 pad_mode = "reflect"108 else:109 pad_mode = "constant"110 part = nn.functional.pad(part, (0, chunk_size - chunk_len), mode=pad_mode, value=0)111 112 batch_data.append(part)113 batch_locations.append((i, chunk_len))114 i += step115 116 # Process batch if it's full or the end is reached117 if len(batch_data) >= batch_size or i >= mix.shape[1]:118 arr = torch.stack(batch_data, dim=0)119 x = model(arr)120 121 if mode == "generic":122 window = windowing_array.clone() # using clone() fixes the clicks at chunk edges when using batch_size=1123 if i - step == 0: # First audio chunk, no fadein124 window[:fade_size] = 1125 elif i >= mix.shape[1]: # Last audio chunk, no fadeout126 window[-fade_size:] = 1127 128 for j, (start, seg_len) in enumerate(batch_locations):129 if mode == "generic":130 result[..., start:start + seg_len] += x[j, ..., :seg_len].cpu() * window[..., :seg_len]131 counter[..., start:start + seg_len] += window[..., :seg_len]132 else:133 result[..., start:start + seg_len] += x[j, ..., :seg_len].cpu()134 counter[..., start:start + seg_len] += 1.0135 136 batch_data.clear()137 batch_locations.clear()138 139 if progress_bar:140 progress_bar.update(step)141 142 if progress_bar:143 progress_bar.close()144 145 146 """147 # mix: B, 2, T148 # req_shape = (num_instruments,) + mix.shape149 req_shape = (num_instruments,) + mix.shape150 result = torch.zeros(req_shape, dtype=torch.float32)151 counter = torch.zeros(req_shape, dtype=torch.float32)152 153 # prev_i = 0154 i = 0155 batch_data = []156 batch_locations = []157 158 while i < mix.shape[-1]:159 part = mix[:, :, i:i + chunk_size].to(device)160 chunk_len = part.shape[-1]161 if mode == "generic" and chunk_len > chunk_size // 2:162 pad_mode = "reflect"163 else:164 pad_mode = "constant"165 part = nn.functional.pad(part, (0, chunk_size - chunk_len), mode=pad_mode, value=0)166 # batch_locations.append((i, chunk_len))167 # prev_i = i168 batch_location = i, i + chunk_len169 i += step170 171 # print(part.shape)172 x = model(part)173 x = x.transpose(0, 1)174 # print(x.shape)175 176 if mode == "generic":177 window = windowing_array.clone() # using clone() fixes the clicks at chunk edges when using batch_size=1178 if i - step == 0: # First audio chunk, no fadein179 window[:fade_size] = 1180 elif i >= mix.shape[1]: # Last audio chunk, no fadeout181 window[-fade_size:] = 1182 183 # for j, (start, seg_len) in enumerate(batch_locations):184 # l = chunk_len if chunk_len < chunk_size else chunk_size185 # print(l, x.shape, result.shape, counter.shape, window.shape)186 # print(result[..., batch_location[0]: batch_location[1]].shape, x[..., :chunk_len].cpu().shape, window[..., :chunk_len].shape)187 if mode == "generic":188 result[..., batch_location[0]: batch_location[1]] += x[..., :chunk_len].cpu() * window[..., :chunk_len]189 counter[..., batch_location[0]: batch_location[1]] += window[..., :chunk_len]190 else:191 result[..., batch_location[0]: batch_location[1]] += x[..., :chunk_len].cpu()192 counter[..., batch_location[0]: batch_location[1]] += 1.0193 194 batch_data.clear()195 batch_locations.clear()196 """197 # Compute final estimated sources198 estimated_sources = result / counter199 estimated_sources = estimated_sources.cpu().numpy()200 np.nan_to_num(estimated_sources, copy=False, nan=0.0)201 202 # Remove padding for generic mode203 if mode == "generic":204 if length_init > 2 * border and border > 0:205 estimated_sources = estimated_sources[..., border:-border]206 207 # Return the result as a dictionary or a single array208 if mode == "demucs":209 instruments = config.training.instruments210 else:211 instruments = prefer_target_instrument(config)212 213 ret_data = {k: v for k, v in zip(instruments, estimated_sources)}214 215 if mode == "demucs" and num_instruments <= 1:216 return estimated_sources217 else:218 return ret_data219 220 221def initialize_model_and_device(model: torch.nn.Module, device_ids: List[int]) -> Tuple[Union[torch.device, str], torch.nn.Module]:222 """223 Move a model to the correct computation device and wrap with DataParallel if needed.224 225 Selects GPU(s) if CUDA is available; otherwise defaults to CPU. If multiple226 GPU IDs are provided, wraps the model with `nn.DataParallel` for multi-GPU227 execution.228 229 Args:230 model (torch.nn.Module): PyTorch model to be initialized.231 device_ids (List[int]): List of GPU device IDs to use. If length > 1,232 the model will be wrapped with DataParallel.233 234 Returns:235 Tuple[Union[torch.device, str], torch.nn.Module]: A tuple containing:236 - The computation device (`torch.device` or "cpu").237 - The model moved to that device (wrapped in DataParallel if applicable).238 """239 240 if torch.cuda.is_available():241 if len(device_ids) <= 1:242 device = torch.device(f'cuda:{device_ids[0]}')243 model = model.to(device)244 else:245 device = torch.device(f'cuda:{device_ids[0]}')246 model = nn.DataParallel(model, device_ids=device_ids).to(device)247 else:248 device = 'cpu'249 model = model.to(device)250 print("CUDA is not available. Running on CPU.")251 252 return device, model253 254 255def get_optimizer(config: ConfigDict, model: torch.nn.Module) -> torch.optim.Optimizer:256 """257 Create and configure an optimizer for training.258 259 Selects the optimizer type based on `config.training.optimizer` and applies260 the corresponding parameters, including support for advanced optimizers261 such as Muon, Prodigy, and 8-bit AdamW. Handles parameter group separation262 for specialized optimizers (e.g., Muon vs. Adam parameters).263 264 Args:265 config (ConfigDict): Training configuration containing optimizer type,266 learning rate, and optional optimizer-specific parameters.267 model (torch.nn.Module): Model whose parameters will be optimized.268 269 Returns:270 torch.optim.Optimizer: Initialized optimizer ready for training.271 272 Raises:273 ValueError: If required optimizer configuration is missing (e.g., for Muon).274 SystemExit: If an unknown optimizer name is encountered.275 """276 277 should_print = not dist.is_initialized() or dist.get_rank() == 0278 optim_params = dict()279 if 'optimizer' in config:280 optim_params = dict(config['optimizer'])281 if config.training.optimizer != 'muon' and should_print:282 print(f'Optimizer params from config:\n{optim_params}')283 284 name_optimizer = getattr(config.training, 'optimizer',285 'No optimizer in config')286 287 if name_optimizer == 'adam':288 optimizer = Adam(model.parameters(), lr=config.training.lr, **optim_params)289 elif name_optimizer == 'adamw':290 optimizer = AdamW(model.parameters(), lr=config.training.lr, **optim_params)291 elif name_optimizer == 'radam':292 optimizer = RAdam(model.parameters(), lr=config.training.lr, **optim_params)293 elif name_optimizer == 'rmsprop':294 optimizer = RMSprop(model.parameters(), lr=config.training.lr, **optim_params)295 elif name_optimizer == 'prodigy':296 from prodigyopt import Prodigy297 # you can choose weight decay value based on your problem, 0 by default298 # We recommend using lr=1.0 (default) for all networks.299 optimizer = Prodigy(model.parameters(), lr=config.training.lr, **optim_params)300 elif name_optimizer == 'adamw8bit':301 import bitsandbytes as bnb302 optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=config.training.lr, **optim_params)303 elif name_optimizer == 'muon':304 if should_print:305 print("Using Muon optimizer (Single-Device) with AdamW for auxiliary parameters.")306 307 muon_params = [p for p in model.parameters() if p.ndim >= 2]308 adam_params = [p for p in model.parameters() if p.ndim < 2]309 310 if not hasattr(config, 'optimizer') or 'muon_group' not in config.optimizer or 'adam_group' not in config.optimizer:311 raise ValueError("For the 'muon' optimizer, the config must have an 'optimizer' section "312 "with 'muon_group' and 'adam_group' dictionaries.")313 314 muon_group_config = dict(config.optimizer.muon_group)315 adam_group_config = dict(config.optimizer.adam_group)316 317 if should_print:318 print(f"Muon group params: {muon_group_config}")319 print(f"Adam group params: {adam_group_config}")320 321 param_groups = [322 dict(params=muon_params, use_muon=True, **muon_group_config),323 dict(params=adam_params, use_muon=False, **adam_group_config),324 ]325 optimizer = SingleDeviceMuonWithAuxAdam(param_groups)326 elif name_optimizer == 'sgd':327 if should_print:328 print('Use SGD optimizer')329 optimizer = SGD(model.parameters(), lr=config.training.lr, **optim_params)330 else:331 if should_print:332 print(f'Unknown optimizer: {name_optimizer}')333 exit()334 return optimizer335 336 337def normalize_batch(x: torch.Tensor, y: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:338 """339 Apply mean-variance normalization to a pair of tensors.340 341 Computes the mean and standard deviation from `x` and normalizes both `x`342 and `y` using those statistics. This ensures the two tensors are scaled343 consistently.344 345 Args:346 x (torch.Tensor): Input tensor used to compute normalization statistics.347 y (torch.Tensor): Input tensor normalized using the same statistics as `x`.348 349 Returns:350 Tuple[torch.Tensor, torch.Tensor]: Normalized tensors `(x, y)`.351 """352 353 mean = x.mean()354 std = x.std()355 if std != 0:356 x = (x - mean) / std357 y = (y - mean) / std358 return x, y359 360 361def apply_tta(362 config,363 model: torch.nn.Module,364 mix: torch.Tensor,365 waveforms_orig: Dict[str, torch.Tensor],366 device: torch.device,367 model_type: str368) -> Dict[str, torch.Tensor]:369 """370 Enhance source separation results using Test-Time Augmentation (TTA).371 372 Applies augmentations such as channel reversal and polarity inversion to373 the input mixture, reprocesses with the model, and combines the results374 with the original predictions by averaging.375 376 Args:377 config: Configuration object with model and inference parameters.378 model (torch.nn.Module): Trained source separation model.379 mix (torch.Tensor): Input mixture tensor of shape (channels, time).380 waveforms_orig (Dict[str, torch.Tensor]): Dictionary of separated381 sources before augmentation.382 device (torch.device): Computation device (CPU or CUDA).383 model_type (str): Model type identifier used for demixing.384 385 Returns:386 Dict[str, torch.Tensor]: Dictionary of separated sources after applying TTA.387 """388 389 # Create augmentations: channel inversion and polarity inversion390 track_proc_list = [mix[::-1].copy(), -1.0 * mix.copy()]391 392 # Process each augmented mixture393 for i, augmented_mix in enumerate(track_proc_list):394 waveforms = demix(config, model, augmented_mix, device, model_type=model_type)395 for el in waveforms:396 if i == 0:397 waveforms_orig[el] += waveforms[el][::-1].copy()398 else:399 waveforms_orig[el] -= waveforms[el]400 401 # Average the results across augmentations402 for el in waveforms_orig:403 waveforms_orig[el] /= len(track_proc_list) + 1404 405 return waveforms_orig406 407 408def _getWindowingArray(window_size: int, fade_size: int) -> torch.Tensor:409 """410 Generate a windowing array with a linear fade-in at the beginning and a fade-out at the end.411 412 This function creates a window of size `window_size` where the first `fade_size` elements413 linearly increase from 0 to 1 (fade-in) and the last `fade_size` elements linearly decrease414 from 1 to 0 (fade-out). The middle part of the window is filled with ones.415 416 Parameters:417 ----------418 window_size : int419 The total size of the window.420 fade_size : int421 The size of the fade-in and fade-out regions.422 423 Returns:424 -------425 torch.Tensor426 A tensor of shape (window_size,) containing the generated windowing array.427 428 Example:429 -------430 If `window_size=10` and `fade_size=3`, the output will be:431 tensor([0.0000, 0.5000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.5000, 0.0000])432 """433 434 fadein = torch.linspace(0, 1, fade_size)435 fadeout = torch.linspace(1, 0, fade_size)436 437 window = torch.ones(window_size)438 window[-fade_size:] = fadeout439 window[:fade_size] = fadein440 return window441 442 443def prefer_target_instrument(config: ConfigDict) -> List[str]:444 """445 Return the list of target instruments based on the configuration.446 If a specific target instrument is specified in the configuration,447 it returns a list with that instrument. Otherwise, it returns the list of instruments.448 449 Parameters:450 ----------451 config : ConfigDict452 Configuration object containing the list of instruments or the target instrument.453 454 Returns:455 -------456 List[str]457 A list of target instruments.458 """459 if getattr(config.training, 'target_instrument', None):460 return [config.training.target_instrument]461 else:462 return config.training.instruments463 464 465def load_not_compatible_weights(model: torch.nn.Module, old_model: dict, verbose: bool = False) -> None:466 """467 Load a possibly incompatible state dict into `model` with best-effort matching.468 469 Accepts either a raw state_dict or a checkpoint dict with weights under "state" or "state_dict".470 For each param/buffer in `model`: if the name exists and shapes match → copy;471 if ndim matches but shapes differ → zero-pad/crop the source to fit the target;472 if the name is missing or ndim differs → skip. Optional logging on rank 0 when `verbose=True`.473 474 Args:475 model: Target PyTorch module.476 old_model: Source weights (state_dict or checkpoint dict).477 verbose: Print brief load decisions.478 479 Returns:480 None481 """482 483 should_print = verbose and (not dist.is_initialized() or dist.get_rank() == 0)484 485 new_model = model.state_dict()486 487 if 'state' in old_model:488 # Fix for htdemucs weights loading489 old_model = old_model['state']490 if 'state_dict' in old_model:491 # Fix for apollo weights loading492 old_model = old_model['state_dict']493 if 'model_state_dict' in old_model:494 # Fix for full_check_point495 old_model = old_model['model_state_dict']496 497 for el in new_model:498 if el in old_model:499 if should_print:500 print(f'Match found for {el}!')501 if new_model[el].shape == old_model[el].shape:502 if should_print:503 print('Action: Just copy weights!')504 new_model[el] = old_model[el]505 else:506 if len(new_model[el].shape) != len(old_model[el].shape) and should_print:507 print('Action: Different dimension! Too lazy to write the code... Skip it')508 else:509 if should_print:510 print(f'Shape is different: {tuple(new_model[el].shape)} != {tuple(old_model[el].shape)}')511 ln = len(new_model[el].shape)512 max_shape = []513 slices_old = []514 slices_new = []515 for i in range(ln):516 max_shape.append(max(new_model[el].shape[i], old_model[el].shape[i]))517 slices_old.append(slice(0, old_model[el].shape[i]))518 slices_new.append(slice(0, new_model[el].shape[i]))519 # print(max_shape)520 # print(slices_old, slices_new)521 slices_old = tuple(slices_old)522 slices_new = tuple(slices_new)523 max_matrix = np.zeros(max_shape, dtype=np.float32)524 for i in range(ln):525 max_matrix[slices_old] = old_model[el].cpu().numpy()526 max_matrix = torch.from_numpy(max_matrix)527 new_model[el] = max_matrix[slices_new]528 else:529 if should_print:530 print(f'Match not found for {el}!')531 model.load_state_dict(532 new_model533 )534 535 536def load_lora_weights(model: torch.nn.Module, lora_path: str, device: str = 'cpu') -> None:537 """538 Load LoRA weights into a model.539 This function updates the given model with LoRA-specific weights from the specified checkpoint file.540 It does not require the checkpoint to match the model's full state dictionary, as only LoRA layers are updated.541 542 Parameters:543 ----------544 model : Module545 The PyTorch model into which the LoRA weights will be loaded.546 lora_path : str547 Path to the LoRA checkpoint file.548 device : str, optional549 The device to load the weights onto, by default 'cpu'. Common values are 'cpu' or 'cuda'.550 551 Returns:552 -------553 None554 The model is updated in place.555 """556 lora_state_dict = torch.load(lora_path, map_location=device)557 model.load_state_dict(lora_state_dict, strict=False)558 559 560def load_start_checkpoint(args: argparse.Namespace,561 model: torch.nn.Module,562 old_model: None,563 type_: str = 'train') -> None:564 """565 Load an initial checkpoint into `model`.566 567 For `type_ == "train"`, performs a tolerant load using `old_model` (a state dict or a568 checkpoint dict) via `load_not_compatible_weights`, allowing partial shape mismatches.569 For other modes, loads a strict state dict from `args.start_check_point`, with special570 handling for HTDemucs/Apollo checkpoints (keys under "state"/"state_dict"). If571 `args.lora_checkpoint` is set, LoRA weights are applied after the base load.572 573 Args:574 args: Namespace with at least `start_check_point`, `model_type`, and optionally `lora_checkpoint`.575 model: Target PyTorch module to receive weights.576 old_model: Source weights for tolerant loading in train mode (state dict or checkpoint dict).577 type_: Loading strategy; "train" uses tolerant loading, otherwise strict loading from path.578 579 Returns:580 None581 """582 should_print = not dist.is_initialized() or dist.get_rank() == 0583 584 if should_print:585 print(f'Start from checkpoint: {args.start_check_point}')586 if type_ in ['train']:587 if 1:588 load_not_compatible_weights(model, old_model, verbose=False)589 else:590 model.load_state_dict(torch.load(args.start_check_point))591 else:592 device='cpu'593 if args.model_type in ['htdemucs', 'apollo']:594 state_dict = torch.load(args.start_check_point, map_location=device, weights_only=False)595 # Fix for htdemucs pretrained models596 if 'state' in state_dict:597 state_dict = state_dict['state']598 # Fix for apollo pretrained models599 if 'state_dict' in state_dict:600 state_dict = state_dict['state_dict']601 else:602 state_dict = torch.load(args.start_check_point, map_location=device, weights_only=True)603 model.load_state_dict(state_dict)604 605 if args.lora_checkpoint:606 if should_print:607 print(f"Loading LoRA weights from: {args.lora_checkpoint}")608 load_lora_weights(model, args.lora_checkpoint)609 610 611def bind_lora_to_model(config: Dict[str, Any], model: nn.Module) -> nn.Module:612 """613 Replaces specific layers in the model with LoRA-extended versions.614 615 Parameters:616 ----------617 config : Dict[str, Any]618 Configuration containing parameters for LoRA. It should include a 'lora' key with parameters for `MergedLinear`.619 model : nn.Module620 The original model in which the layers will be replaced.621 622 Returns:623 -------624 nn.Module625 The modified model with the replaced layers.626 """627 628 if 'lora' not in config:629 raise ValueError("Configuration must contain the 'lora' key with parameters for LoRA.")630 631 replaced_layers = 0 # Counter for replaced layers632 should_print = not dist.is_initialized() or dist.get_rank() == 0633 634 for name, module in model.named_modules():635 hierarchy = name.split('.')636 layer_name = hierarchy[-1]637 638 # Check if this is the target layer to replace (and layer_name == 'to_qkv')639 if isinstance(module, nn.Linear):640 try:641 # Get the parent module642 parent_module = model643 for submodule_name in hierarchy[:-1]:644 parent_module = getattr(parent_module, submodule_name)645 646 # Replace the module with LoRA-enabled layer647 setattr(648 parent_module,649 layer_name,650 lora.MergedLinear(651 in_features=module.in_features,652 out_features=module.out_features,653 bias=module.bias is not None,654 **config['lora']655 )656 )657 replaced_layers += 1 # Increment the counter658 659 except Exception as e:660 if should_print:661 print(f"Error replacing layer {name}: {e}")662 663 if replaced_layers == 0 and should_print:664 print("Warning: No layers were replaced. Check the model structure and configuration.")665 elif should_print:666 print(f"Number of layers replaced with LoRA: {replaced_layers}")667 668 return model669 670 671def save_weights(672 store_path: str,673 model: nn.Module,674 device_ids: List[int],675 optimizer: torch.optim.Optimizer,676 epoch: int,677 all_time_all_metrics,678 best_metric: float,679 scheduler: Optional[torch.optim.lr_scheduler.ReduceLROnPlateau] = None,680 train_lora: bool = False681) -> None:682 """683 Save a training checkpoint containing model weights, optimizer/scheduler states, and metadata.684 685 Behavior:686 - In Distributed Data Parallel (DDP), only rank 0 writes the file to avoid conflicts.687 - If `train_lora` is True, saves only LoRA adapter weights (`lora_state_dict`); otherwise saves the full model.688 - Uses `model.module.state_dict()` when the model is wrapped by DDP/DataParallel.689 - Stores `epoch` and `best_metric` alongside optimizer/scheduler states.690 691 Args:692 store_path: Destination file path for the checkpoint (will be overwritten).693 model: The model whose weights are being saved (may be wrapped by DDP/DataParallel).694 device_ids: List of GPU device IDs used during training (used to detect DP wrapping in non-DDP runs).695 optimizer: Optimizer whose state will be saved.696 epoch: Current training epoch to record in the checkpoint.697 all_time_all_metrics:698 best_metric: Best validation metric achieved so far.699 scheduler: Optional learning rate scheduler; its state is saved if provided.700 train_lora: If True, save only LoRA adapter weights instead of the full model.701 702 Returns:703 None704 """705 706 checkpoint: Dict[str, Any] = {707 "epoch": epoch,708 "optimizer_name": optimizer.__class__.__name__,709 "optimizer_state_dict": optimizer.state_dict(),710 "scheduler_state_dict": scheduler.state_dict() if scheduler else None,711 "best_metric": best_metric,712 "all_metrics": all_time_all_metrics713 }714 715 # Save model weights716 if train_lora:717 checkpoint["model_state_dict"] = lora.lora_state_dict(model)718 else:719 if dist.is_initialized():720 # In DDP, use .module721 checkpoint["model_state_dict"] = model.module.state_dict()722 else:723 checkpoint["model_state_dict"] = (724 model.state_dict() if len(device_ids) <= 1 else model.module.state_dict()725 )726 727 # Save only on rank 0 (or if not using DDP)728 if not dist.is_initialized() or dist.get_rank() == 0:729 torch.save(checkpoint, store_path)730 731 732def save_last_weights(733 args: argparse.Namespace,734 model: nn.Module,735 device_ids: List[int],736 optimizer: torch.optim.Optimizer,737 epoch: int,738 all_time_all_metrics,739 best_metric: float,740 scheduler: Optional[torch.optim.lr_scheduler.ReduceLROnPlateau] = None,741) -> None:742 """743 Save the latest training checkpoint for continuation or recovery.744 745 The checkpoint is always written to:746 {args.results_path}/last_{args.model_type}.ckpt747 748 This wraps `save_weights` and ensures the latest model/optimizer/scheduler749 states are recorded, along with the current epoch and best metric. In DDP,750 only rank 0 performs the save. Supports both standard and LoRA training.751 752 Args:753 all_time_all_metrics:754 args: Training arguments. Must define `results_path`, `model_type`,755 and `train_lora`.756 model: Model instance (may be wrapped by DDP/DataParallel).757 device_ids: List of GPU IDs used for training.758 optimizer: Optimizer whose state will be saved.759 epoch: Current training epoch.760 best_metric: Current best validation metric.761 scheduler: Optional learning rate scheduler to save state for.762 763 Returns:764 None765 """766 store_path = f"{args.results_path}/last_{args.model_type}.ckpt"767 save_weights(768 store_path,769 model,770 device_ids,771 optimizer,772 epoch,773 all_time_all_metrics,774 best_metric,775 scheduler,776 args.train_lora,777 )778