Repository: transformers
Issue: Trainer: TensorBoardCallback not working for "on_save" and "on_save_end" events
Labels: 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.