How to Log Training Loss at Step Zero in Hugging Face Trainer or SFT Trainer?
Author: brando90
Pull request: no
Labels: name: Feature request

### Feature request

log train loss on start

----

’m using the Hugging Face `Trainer` (or `SFTTrainer`) for fine-tuning, and I want to log the training loss at step 0 (before any training steps are executed). I know there’s an `eval_on_start` option for evaluation, but I couldn't find a direct equivalent for training loss logging at the beginning of training.

Is there a way to log the initial training loss at step zero (before any updates) using `Trainer` or `SFTTrainer`? Ideally, I'd like something similar to `eval_on_start`.

Here’s what I’ve tried so far:

#### Solution 1: Custom Callback

I implemented a custom callback to log the training loss at the start of training:


```python
from transformers import TrainerCallback

class TrainOnStartCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, logs=None, **kwargs):
        # Log training loss at step 0
        logs = logs or {}
        logs["train/loss"] = None  # Replace None with an initial value if available
        logs["train/global_step"] = 0
        self.log(logs)

    def log(self, logs):
        print(f"Logging at start: {logs}")
        wandb.log(logs)

# Adding the callback to the Trainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    args=training_args,
    optimizers=(optimizer, scheduler),
    callbacks=[TrainOnStartCallback()],
)
```
This works but feels a bit overkill. It logs metrics at the start of training before any steps.

#### Solution 2: Manual Logging

Alternatively, I manually log the training loss before starting training:

```python
wandb.log({"train/loss": None, "train/global_step": 0})
trainer.train()
```

### Question:

Are there any built-in features in `Trainer` or `SFTTrainer` to log training loss at step zero? Or is a custom callback or manual logging the best solution here? If so, are there better ways to implement this functionality? similar to the `eval_on_start` but `train_on_start`?

cross: https://discuss.huggingface.co/t/how-to-log-training-loss-at-step-zero-in-hugging-face-trainer-or-sft-trainer/128188

### Motivation

Crucial sanity check

### Your contribution

yes, happy to implement this. 

Comments:
https://stackoverflow.com/questions/79232257/how-to-log-training-loss-at-step-zero-in-hugging-face-trainer-or-sft-trainer

https://discuss.huggingface.co/t/how-to-log-training-loss-at-step-zero-in-hugging-face-trainer-or-sft-trainer/128188

I think the loss at Step One is the same as Stage Zero because LR is 0 according to https://github.com/huggingface/transformers/issues/34754. You can using loss at Stage One.

cc @SunMarc @muellerzr 

@techkang 

what about: `logging_first_step`? `        logging_first_step (`bool`, *optional*, defaults to `False`):
            Whether to log the first `global_step` or not.`

> I think the loss at Step One is the same as Stage Zero because LR is 0 according to #34754. You can using loss at Stage One.

I'm not using warm up. I don't always use it. 

@brando90 If you don't using warm up, I think you can set `logging_first_step=True` and `warmup_steps=1`.