Gpu support finished

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."""
+    """Setup process function with optimized CUDA handling."""
-    # 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)
-
     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!
+        if torch.distributed.is_initialized():
-        torch.distributed.destroy_process_group()
+            torch.distributed.destroy_process_group()
 
 
 def multi_gpu_example():
-    """Example of finetuning TimesFM using multiple GPUs."""
+    """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
-    # Define which GPUs to use
+    mp.set_start_method("spawn", force=True)
-    gpu_ids = [0]  # Just using one GPU
+
+    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
+    try:
-    # basic_example()  # Single GPU
+        # single_gpu_example()  # Single GPU
-    multi_gpu_example()  # Multi-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()

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:
+        with torch.no_grad():
-            x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
+            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 = self.model(x_context, x_padding.float(), freq)
-            predictions_mean = predictions[..., 0]
+                predictions_mean = predictions[..., 0]
-            last_patch_pred = predictions_mean[:, -1, :]
+                last_patch_pred = predictions_mean[:, -1, :]
 
-            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+                loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
-            total_loss += loss.item()
+
+                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)
+                if self.config.distributed:
-                history["val_loss"].append(val_loss)
+                    train_tensor = torch.tensor(train_loss, device=self.device)
-                history["learning_rate"].append(current_lr)
+                    val_tensor = torch.tensor(val_loss, device=self.device)
 
-                metrics = {
+                    world_size = dist.get_world_size()
-                    "train_loss": train_loss,
+                    train_losses = [torch.zeros_like(train_tensor, device=self.device) for _ in range(world_size)]
-                    "val_loss": val_loss,
+                    val_losses = [torch.zeros_like(val_tensor, device=self.device) for _ in range(world_size)]
-                    "learning_rate": current_lr,
-                    "epoch": epoch + 1,
-                }
 
-                if self.config.use_wandb:
+                    dist.all_gather(train_losses, train_tensor)
-                    wandb.log(metrics)
+                    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")