Trainer: TensorBoardCallback not working for "on_save" and "on_save_end" events
Author: vecorro
Pull request: no
Labels: name: bug

### System Info

transformers              4.47.1
torch                     2.5.0+cu121
Ubuntu 22.04 LTS

### Who can help?

Hi, @muellerz and @SunMarc 

I'm trying to get storage I/O related to the Trainer's checkpoint-saving operations. For that matter I implemented the following class:

```
import time
import psutil
from transformers.integrations import TensorBoardCallback

class DiskIOMonitoringCallback(TensorBoardCallback):
    def __init__(self, tb_writer=None):
        super().__init__(tb_writer=tb_writer)
        self.start_time = None
        self.initial_disk_io = None

    def _compute_disk_io_metrics(self):
        """Compute disk I/O metrics."""
        final_disk_io = psutil.disk_io_counters()
        if self.initial_disk_io:
            read_bytes = final_disk_io.read_bytes - self.initial_disk_io.read_bytes
            write_bytes = final_disk_io.write_bytes - self.initial_disk_io.write_bytes
        else:
            read_bytes, write_bytes = 0, 0
        return read_bytes, write_bytes

    def _init_summary_writer(self, args, log_dir=None):
        """Ensure TensorBoard writer is initialized."""
        log_dir = log_dir or args.logging_dir  # Use provided log_dir or fallback to args.logging_dir
        if self.tb_writer is None:
            from torch.utils.tensorboard import SummaryWriter
            self.tb_writer = SummaryWriter(log_dir=log_dir)

    def on_save(self, args, state, control, **kwargs):
        """Hook triggered before saving a checkpoint."""
        if self.tb_writer is None:
            self._init_summary_writer(args)

        # Record start time and initial disk I/O stats
        self.start_time = time.time()
        self.initial_disk_io = psutil.disk_io_counters()

    def on_save_end(self, args, state, control, **kwargs):
        """Hook triggered after saving a checkpoint."""
        # Calculate save duration
        save_duration = time.time() - self.start_time

        # Compute disk I/O metrics
        read_bytes, write_bytes = self._compute_disk_io_metrics()

        # Log metrics to TensorBoard
        if self.tb_writer:
            self.tb_writer.add_scalar("DiskIO/Save Duration (s)", save_duration, state.global_step)
            self.tb_writer.add_scalar("DiskIO/Read Bytes", read_bytes, state.global_step)
            self.tb_writer.add_scalar("DiskIO/Write Bytes", write_bytes, state.global_step)

        # Print metrics for debugging purposes
        print(f"Checkpoint saved in {save_duration:.2f}s. Read: {read_bytes} bytes, Write: {write_bytes} bytes.")

```

My trainer session is invoked line this:

```
training_args = TrainingArguments(
    output_dir="./results",  
    optim="adamw_torch",  
    num_train_epochs=6,
    per_device_train_batch_size=64, 
    gradient_accumulation_steps=8,  
    learning_rate=3e-5,
    weight_decay=0,
    warmup_steps=100,  
    lr_scheduler_type="cosine",  
    gradient_checkpointing=True,  
    dataloader_num_workers=8,  
    bf16=True,  
    logging_steps=10,  
    report_to="tensorboard", 
    save_strategy="epoch",
    save_total_limit=2,  
)


trainer = Trainer(
    model=model, 
    args=training_args, 
    train_dataset=train_dataset, 
    eval_dataset=None,
    data_collator=data_collator,
    callbacks=[
        GpuCpuMonitoringCallback(),
        SystemMonitoringCallback(),
        DiskIOMonitoringCallback(),
    ],
)

trainer.train()
```

The `GpuCpuMonitoringCallback` and `SystemMonitoringCallback` work properly, but I'm now getting data from `DiskIOMonitoringCallback` despite multiple implementation changes. I'm missing something, or something might not be working at the callbacks layer.

I really appreciate any help you can provide. 


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I provided the code for reproduction in the description. 

### Expected behavior

I expect Tensorboard to provide a card to show I/O disk data upon checkpoint saving operations. 

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> The GpuCpuMonitoringCallback and SystemMonitoringCallback work properly, but I'm now getting data from DiskIOMonitoringCallback despite multiple implementation changes. I'm missing something, or something might not be working at the callbacks layer.

Hey @vecorro, sorry for the delay. I'm not sure about the issue that you are experiencing. Is the issue that the callbacks you added are not called ? 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.