How to measure training energy in Hugging Face Trainer and PyTorch Lightning
On this page
InstallHugging Face TrainerPyTorch LightningWhat you getSeveral GPUs and machinesCompare two training setups fairlyKeep every runIf you already use CodeCarbonWhat it does not countA training run's energy is the number behind its cost and its carbon, and it is easy to guess and hard to know. These two callbacks read it from the hardware's own counters (NVIDIA's energy counter, AMD and Intel GPU sensors, RAPL for the CPU and memory, powermetrics on a Mac) for the whole of training, and put it in the same logs as your loss.
Install
pip install -U landauer-gap
joules doctor # what this machine can measure, with a live reading from each meter
joules itself has no dependencies; each callback needs only the framework you already use.
Hugging Face Trainer
from transformers import Trainer
from joules.integrations.huggingface import attach
trainer = Trainer(model=model, args=args, train_dataset=train_ds)
energy = attach(trainer)
trainer.train()
print(energy.result)
attach() puts the callback first in the Trainer's list, so the energy reaches W&B, TensorBoard, MLflow and every other report_to logger. You can also pass JoulesCallback() in Trainer(callbacks=[...]); the Trainer then calls it after those loggers, so the energy still prints and lands in .result, but not in their charts.
PyTorch Lightning
import lightning as L
from joules.integrations.lightning import JoulesCallback
energy = JoulesCallback()
trainer = L.Trainer(max_epochs=3, callbacks=[energy])
trainer.fit(model, datamodule)
print(energy.result)
It works with lightning and with the older pytorch_lightning package, and logs through self.log, so every Lightning logger gets it.
What you get
In the logs, next to the loss, every time the trainer logs (and at each epoch's end with Lightning):
energy/joules: energy used since training startedenergy/kwh: the same in kWhenergy/avg_watts: average power so farenergy/joules_per_step: energy per optimizer step so far, the number to compare two configurations by
When training ends, one line:
joules: training used … = … over …, … W average, … per step; cost … at 0.1/kWh, CO₂ … g
and energy.result holds the numbers: energy, kWh, duration, average power, joules per step, steps, epochs, cost, CO₂, each epoch as a phase (count, median energy and time, spread) and, on NVIDIA, what the GPU did (its clock against the maximum and the share of time it was throttled). price=, grid= (or ci= in g CO₂/kWh) and pue= set the cost and carbon, as in joules run.
Several GPUs and machines
Only the local main process of each machine measures, and it measures every GPU it can see, so a machine's GPUs are counted once however many processes share them. gpus=[0, 1] limits it to some. With Lightning on several machines, each machine's callback has its own .result, and the logs carry the first machine's numbers.
Compare two training setups fairly
To know whether a change (mixed precision, a larger batch, a new data loader) saves energy, run both on the same machine, alternating, and look at the medians, not at one run each:
joules compare --names bf16,fp32 --runs 5 --warmup 1 \
"python train.py --bf16 --max-steps 200" "python train.py --max-steps 200"
joules compare alternates the two commands so heat and clock drift hit both, reports the median energy and time of each with a p-value, and warns when one side ran throttled more than the other. Compare by joules per step, or keep the number of steps the same.
Keep every run
energy = attach(trainer, track="llama-finetune")
saves each training run to Tracking in your WattOps account when it ends (run joules login once first): energy, time, cost, CO₂, GPU health and the epochs, never code, data or hyperparameters. The line it prints says whether it was saved.
If you already use CodeCarbon
transformers ships a CodeCarbonCallback, switched on through report_to, and both read the hardware. CodeCarbon's focus is the CO₂ of a run, with the location's carbon intensity looked up for you. What this callback adds: each epoch as a phase, joules per step in the same logs as the loss, on NVIDIA the GPU's clock and throttling, joules compare for a fair A/B of two training configurations, and the open leaderboard. The two can run together.
What it does not count
The meters measure devices: GPU boards, CPU packages and memory. The power supply's losses, fans, storage, networking and the building's cooling are not in the number; pue= adds an overhead for them, and joules calibrate against a wall-plug meter (in the command list) tells you what value fits your machine. On a machine with no energy counters the callback says so in one line and training runs as before.
Every command is in the joules guide. To see every job's energy across a GPU fleet, live, with the GPUs that sit idle, try WattOps.