From 955d6cda2297c656b96e280fc5a935ea30d7f61b Mon Sep 17 00:00:00 2001 From: misha-chertushkin Date: Tue, 21 Jan 2025 15:52:08 +0000 Subject: [PATCH] Gpu support finished --- notebooks/finetuning_example.py | 63 ++++++++++--------- notebooks/finetuning_torch.py | 107 +++++++++++++++++++++++--------- 2 files changed, 108 insertions(+), 62 deletions(-) diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py index c42cf88..f1f8a74 100644 --- a/notebooks/finetuning_example.py +++ b/notebooks/finetuning_example.py @@ -107,7 +107,6 @@ def get_model(load_weights: bool = False): checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt") loaded_checkpoint = torch.load(checkpoint_path, weights_only=True) model.load_state_dict(loaded_checkpoint) - model = model.to(device) return model, hparams, tfm._model_config @@ -219,54 +218,55 @@ def single_gpu_example(): def setup_process(rank, world_size, model, config, train_dataset, val_dataset, return_dict): - """Initialize the distributed process.""" - # Set up the process group - os.environ["MASTER_ADDR"] = "localhost" - os.environ["MASTER_PORT"] = "12355" - - # Initialize the process group - torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=rank) - - # Set the device for this process - torch.cuda.set_device(rank) - + """Setup process function with optimized CUDA handling.""" try: + if torch.cuda.is_available(): + torch.cuda.set_device(rank) + + os.environ["MASTER_ADDR"] = config.master_addr + os.environ["MASTER_PORT"] = config.master_port + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl", world_size=world_size, rank=rank) + finetuner = TimesFMFinetuner(model, config, rank=rank) + results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset) - if rank == 0: # Only store results and plot from the main process + if rank == 0: return_dict["results"] = results plot_predictions( model=model, val_dataset=val_dataset, save_path="timesfm_predictions.png", ) + + except Exception as e: + print(f"Error in process {rank}: {str(e)}") + raise e finally: - # Cleanup - important! - torch.distributed.destroy_process_group() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() def multi_gpu_example(): - """Example of finetuning TimesFM using multiple GPUs.""" - # Define which GPUs to use - gpu_ids = [0] # Just using one GPU + """Example of finetuning TimesFM using multiple GPUs with optimized spawn.""" + mp.set_start_method("spawn", force=True) + + gpu_ids = [0, 1] world_size = len(gpu_ids) - # Initialize model and config model, hparams, tfm_config = get_model(load_weights=True) + + # Create config config = FinetuningConfig( batch_size=256, num_epochs=5, learning_rate=1e-4, - use_wandb=False, + use_wandb=True, distributed=True, gpu_ids=gpu_ids, ) - - # Get datasets train_dataset, val_dataset = get_data(128, tfm_config.horizon_len) - - # Create a multiprocessing manager to share results between processes manager = mp.Manager() return_dict = manager.dict() @@ -278,16 +278,17 @@ def multi_gpu_example(): join=True, ) - # Get results from the main process results = return_dict.get("results", None) print("\nFinetuning completed!") - if results: - print(f"Training history: {len(results['history']['train_loss'])} epochs") - return results if __name__ == "__main__": - # Use either single GPU or multi-GPU example - # basic_example() # Single GPU - multi_gpu_example() # Multi-GPU + try: + # single_gpu_example() # Single GPU + multi_gpu_example() # Multi-GPU + except Exception as e: + print(f"Training failed: {str(e)}") + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py index a10d99e..29215ef 100644 --- a/notebooks/finetuning_torch.py +++ b/notebooks/finetuning_torch.py @@ -39,7 +39,6 @@ from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, Dataset import wandb -from timesfm import TimesFm @dataclass @@ -63,14 +62,14 @@ class FinetuningConfig: gpu_ids: List[int] = field(default_factory=lambda: [0]) # List of GPU IDs to use distributed: bool = False - master_port: str = "12355" + master_port: str = "12358" master_addr: str = "localhost" class TimesFMFinetuner: def __init__( self, - model: TimesFm, + model, config: FinetuningConfig, rank: int = 0, loss_fn: Optional[callable] = None, @@ -87,7 +86,7 @@ class TimesFMFinetuner: self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2)) - if config.use_wandb and rank == 0: # Only initialize wandb on main process + if config.use_wandb and rank == 0: self._setup_wandb() def _setup_distributed(self, rank): @@ -100,7 +99,11 @@ class TimesFMFinetuner: def _setup_wandb(self) -> None: """Initialize Weights & Biases logging.""" - wandb.init(project=self.config.wandb_project, config=self.config.__dict__) + + def _setup_wandb(self) -> None: + """Initialize Weights & Biases logging only on the main process.""" + if self.rank == 0: # Only initialize on main process + wandb.init(project=self.config.wandb_project, config=self.config.__dict__) def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader: """Create a dataloader from a dataset.""" @@ -115,15 +118,11 @@ class TimesFMFinetuner: dataset, batch_size=self.config.batch_size, shuffle=(name == "train" and not self.config.distributed), - num_workers=mp.cpu_count() // len(self.config.gpu_ids), - pin_memory=self.device.type == "cuda", - persistent_workers=True, - prefetch_factor=2, sampler=sampler, ) def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float: - """Train for one epoch.""" + """Train for one epoch with loss debugging.""" self.model.train() total_loss = 0.0 n_batches = len(train_loader) @@ -136,6 +135,10 @@ class TimesFMFinetuner: last_patch_pred = predictions_mean[:, -1, :] loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1)) + if self.config.distributed: + losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())] + dist.all_gather(losses, loss) + optimizer.zero_grad() loss.backward() optimizer.step() @@ -144,21 +147,26 @@ class TimesFMFinetuner: return total_loss / n_batches - @torch.no_grad() def _validate(self, val_loader: DataLoader) -> float: - """Perform validation.""" + """Perform validation with loss debugging.""" self.model.eval() total_loss = 0.0 - for batch in val_loader: - x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch] + with torch.no_grad(): + for batch in val_loader: + x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch] - predictions = self.model(x_context, x_padding.float(), freq) - predictions_mean = predictions[..., 0] - last_patch_pred = predictions_mean[:, -1, :] + predictions = self.model(x_context, x_padding.float(), freq) + predictions_mean = predictions[..., 0] + last_patch_pred = predictions_mean[:, -1, :] - loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1)) - total_loss += loss.item() + loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1)) + + if self.config.distributed: + losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())] + dist.all_gather(losses, loss) + + total_loss += loss.item() return total_loss / len(val_loader) @@ -202,21 +210,58 @@ class TimesFMFinetuner: current_lr = optimizer.param_groups[0]["lr"] - history["train_loss"].append(train_loss) - history["val_loss"].append(val_loss) - history["learning_rate"].append(current_lr) + if self.config.distributed: + train_tensor = torch.tensor(train_loss, device=self.device) + val_tensor = torch.tensor(val_loss, device=self.device) - metrics = { - "train_loss": train_loss, - "val_loss": val_loss, - "learning_rate": current_lr, - "epoch": epoch + 1, - } + world_size = dist.get_world_size() + train_losses = [torch.zeros_like(train_tensor, device=self.device) for _ in range(world_size)] + val_losses = [torch.zeros_like(val_tensor, device=self.device) for _ in range(world_size)] - if self.config.use_wandb: - wandb.log(metrics) + dist.all_gather(train_losses, train_tensor) + dist.all_gather(val_losses, val_tensor) - print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}") + if self.rank == 0 and self.config.use_wandb: + train_losses = [t.cpu().item() for t in train_losses] + val_losses = [t.cpu().item() for t in val_losses] + + for gpu_idx, (t_loss, v_loss) in enumerate(zip(train_losses, val_losses)): + wandb.log( + { + f"train_loss_gpu_{gpu_idx}": t_loss, + f"val_loss_gpu_{gpu_idx}": v_loss, + }, + commit=False, + ) + + wandb.log( + { + "train_loss": train_loss, + "val_loss": val_loss, + "learning_rate": current_lr, + "epoch": epoch + 1, + } + ) + history["train_loss"].append(train_loss) + history["val_loss"].append(val_loss) + history["learning_rate"].append(current_lr) + + else: + if self.config.use_wandb: + wandb.log( + { + "train_loss": train_loss, + "val_loss": val_loss, + "learning_rate": current_lr, + "epoch": epoch + 1, + } + ) + history["train_loss"].append(train_loss) + history["val_loss"].append(val_loss) + history["learning_rate"].append(current_lr) + + if self.rank == 0: + print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}") except KeyboardInterrupt: self.logger.info("Training interrupted by user")