Logging system metrics from all GPUs to Logger

[chunweilim]

Hi all,

I'm trying to log the system metrics from all GPUs to the configured Logger (now using wandb).

I have written a Callback, which has the following on_train_batch_start:

def on_train_batch_start(self, trainer: Trainer, pl_module: LightningModule, batch, batch_idx) -> None:
        current_device = torch.device(f"cuda:{trainer.local_rank}")
        global_rank = trainer.global_rank
        result = get_nvidia_gpu_stats(current_device) # This is from lightning.pytorch.accelerators.cuda
        new_keys_result = {}
        for key in result.keys():
            new_key = f"{key}_{global_rank}"
            new_keys_result[new_key] = result[key]
        pl_module.log_dict(new_keys_result)

Currently, I'm only seeing GPU of index 0 being logged.

1. Is there a way to log the result from all GPUs?
2. Is there also a way to group the result from all GPUs once logged? The following is the system metric when viewing from wandb. It logs only every 30 seconds, which is not what I want, hence I'm trying the above to log more frequently.

[tsdev]

I am having the exact same issue, any thoughts?

[chunweilim]

Hi @tsdev,

In the end, I went back to vanilla PyTorch for my experiments as I found out that the iteration times for training a model on Lightning vary quite a lot, and I have no good way of just measuring the iteration times for the forward pass and backward pass individually.

To accumulate the system metrics, I used torch.cuda.{max_memory_allocated, max_memory_reserved}. Then I accumulate the result in the rank 0 device before sending it to wandb for logging.

gpu_memory_max_allocated = torch.cuda.max_memory_allocated(device)
max_allocated_tensor = torch.tensor(gpu_memory_max_allocated, device=device)
torch.distributed.all_reduce(tensor=max_allocated_tensor, op=torch.distributed.ReduceOp.SUM)

[tsdev]

Thanks! I will try to gather the results. Still I see why your above method does not produce individual log entries per each device.