tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix total/trainable tokens log (#3344)

Browse Source 
* fix total/trainable tokens log

* fix total/trainable tokens log
This commit is contained in:
VED
2026-01-06 19:55:17 +05:30
committed by GitHub
parent 8aab807e67
commit 7bf6f70e96
2 changed files with 4 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/axolotl/core/trainers/base.py
View File

@@ -660,11 +660,10 @@ class AxolotlTrainer(
logs["tokens/train_per_sec_per_gpu"] = round( logs["tokens/train_per_sec_per_gpu"] = round(
self.state.last_tokens_per_second.item() / self.args.logging_steps, 2 self.state.last_tokens_per_second.item() / self.args.logging_steps, 2
) )
if ( if "total" in self.state.tokens:
hasattr(self.state, "total_tokens") logs["tokens/total"] = int(self.state.tokens["total"].item())
and self.state.total_tokens is not None if "trainable" in self.state.tokens:
): logs["tokens/trainable"] = int(self.state.tokens["trainable"].item())
logs["total_tokens"] = int(self.state.total_tokens.item())
del self._stored_metrics[train_eval] del self._stored_metrics[train_eval]

6
src/axolotl/utils/callbacks/tokens_per_second.py
View File

@@ -101,9 +101,3 @@ class TokensPerSecondCallback(TrainerCallback):
# Clear per-step tokens after logging # Clear per-step tokens after logging
if tokens and "trainable_tokens" in tokens: if tokens and "trainable_tokens" in tokens:
tokens["trainable_tokens"] = torch.zeros_like(tokens["trainable_tokens"]) tokens["trainable_tokens"] = torch.zeros_like(tokens["trainable_tokens"])
if tokens and "total" in tokens:
logs["tokens/total"] = tokens["total"].item()
if tokens and "trainable" in tokens:
logs["tokens/trainable"] = tokens["trainable"].item()