freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Gpu support finished

Browse Source
This commit is contained in:
misha-chertushkin
2025-01-21 15:52:08 +00:00
parent a559718c66
commit 955d6cda22
2 changed files with 108 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/finetuning_example.py
+32 -31
View File
@@ -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()
notebooks/finetuning_torch.py
+76 -31
View File
@@ -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")