Skip to content

Wire total_weights into ML Diagnostics metrics #474

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
richaguptaa17 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Wire total_weights into ML Diagnostics metrics #474

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions docs/metrics.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ MaxDiffusion integrates with Google Cloud ML Diagnostics to provide real-time te

### Predefined Metrics

MaxDiffusion automatically translates internal scalar keys to canonical `MetricType` enums expected by the Control Plane UI:
MaxDiffusion automatically translates internal scalar keys to canonical metric names expected by the Control Plane UI:

- **Loss** (`loss`): Training loss value per step (mapped from `learning/loss`).
- **Learning Rate** (`learning_rate`): Current optimizer learning rate (mapped from `learning/current_learning_rate`).
Expand Down Expand Up @@ -80,12 +80,16 @@ Inside the trainer's `training_loop()`:
```python
from maxdiffusion import train_utils

# Record standard step metrics (and any custom metrics in train_metric["scalar"]):
# Optional: Add any custom metrics directly to the scalar dictionary
train_metric["scalar"]["custom/my_metric"] = my_metric_value

# Record standard step metrics:
train_utils.record_scalar_metrics(
train_metric,
step_time_delta,
self.per_device_tflops,
learning_rate_scheduler(step),
total_weights=num_model_parameters,
)

if self.config.write_metrics:
Expand Down
10 changes: 6 additions & 4 deletions src/maxdiffusion/tests/metrics_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,11 +146,13 @@ def test_write_metrics_mld_dispatch_master(self, mock_process_index, mock_mld_me
mock_mld_metrics.record_metrics.assert_called_once()
records = mock_mld_metrics.record_metrics.call_args[0][0]

# Verify records contain translated names and float values
# Verify records contain translated string names and float values
record_dict = {r["metric_name"]: r["value"] for r in records}
self.assertAlmostEqual(record_dict[train_utils._METRICS_TO_MANAGED["learning/loss"]], 0.42, places=4)
self.assertAlmostEqual(record_dict[train_utils._METRICS_TO_MANAGED["learning/current_learning_rate"]], 0.0001, places=6)
self.assertAlmostEqual(record_dict[train_utils._METRICS_TO_MANAGED["learning/total_weights"]], 1000000.0, places=1)
self.assertAlmostEqual(record_dict["loss"], 0.42, places=4)
self.assertAlmostEqual(record_dict["learning_rate"], 0.0001, places=6)
self.assertAlmostEqual(record_dict["total_weights"], 1000000.0, places=1)
self.assertAlmostEqual(record_dict["step_time"], 1.0, places=4)
self.assertAlmostEqual(record_dict["tflops"], 50.0, places=4)
self.assertAlmostEqual(record_dict["custom/accuracy"], 0.95, places=4)

@patch("maxdiffusion.train_utils.mld_metrics")
Expand Down
32 changes: 10 additions & 22 deletions src/maxdiffusion/train_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,30 +77,18 @@ def _validate_gcs_bucket_name(bucket_name, config_var):


try:
from google_cloud_mldiagnostics import metrics as mld_metrics, metric_types
from google_cloud_mldiagnostics import metrics as mld_metrics
except ImportError:
mld_metrics = None
metric_types = None


if metric_types is not None:
_METRICS_TO_MANAGED = {
"learning/loss": metric_types.MetricType.LOSS,
"learning/current_learning_rate": metric_types.MetricType.LEARNING_RATE,
"learning/grad_norm": metric_types.MetricType.GRADIENT_NORM,
"learning/total_weights": metric_types.MetricType.TOTAL_WEIGHTS,
"perf/step_time_seconds": metric_types.MetricType.STEP_TIME,
"perf/per_device_tflops_per_sec": metric_types.MetricType.TFLOPS,
}
else:
_METRICS_TO_MANAGED = {
"learning/loss": "loss",
"learning/current_learning_rate": "learning_rate",
"learning/grad_norm": "gradient_norm",
"learning/total_weights": "total_weights",
"perf/step_time_seconds": "step_time",
"perf/per_device_tflops_per_sec": "tflops",
}

_METRICS_TO_MANAGED = {
"learning/loss": "loss",
"learning/current_learning_rate": "learning_rate",
"learning/grad_norm": "gradient_norm",
"learning/total_weights": "total_weights",
"perf/step_time_seconds": "step_time",
"perf/per_device_tflops_per_sec": "tflops",
}


def record_scalar_metrics(metrics, step_time_delta, per_device_tflops, lr, total_weights=None):
Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/trainers/base_wan_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,11 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
self._profiler.stop()

train_utils.record_scalar_metrics(
train_metric, last_step_completion - start_step_time, per_device_tflops, learning_rate_scheduler(step)
train_metric,
last_step_completion - start_step_time,
per_device_tflops,
learning_rate_scheduler(step),
total_weights=num_model_parameters,
)
if self.config.write_metrics:
train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)
Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/trainers/dreambooth_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,11 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
new_time = datetime.datetime.now()

train_utils.record_scalar_metrics(
train_metric, new_time - last_step_completion, self.per_device_tflops, learning_rate_scheduler(step)
train_metric,
new_time - last_step_completion,
self.per_device_tflops,
learning_rate_scheduler(step),
total_weights=num_model_parameters,

@Perseus14 Perseus14 Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The same counting issue applies here: num_model_parameters includes only the UNet, while DreamBooth also updates the text encoder when train_text_encoder=True. Please conditionally include text_encoder_state.params so total_weights reflects all trainable parameters.

)
if self.config.write_metrics:
train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)
Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/trainers/flux_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -441,7 +441,11 @@ def training_loop(
new_time = datetime.datetime.now()

record_scalar_metrics(
train_metric, new_time - last_step_completion, self.per_device_tflops, unet_learning_rate_scheduler(step)
train_metric,
new_time - last_step_completion,
self.per_device_tflops,
unet_learning_rate_scheduler(step),
total_weights=num_model_parameters,
)
if self.config.write_metrics:
write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)
Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/trainers/sdxl_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,11 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
difference_in_ms = time_difference.total_seconds() * 1000
max_logging.log(f"Step time {difference_in_ms}ms")
record_scalar_metrics(
train_metric, last_step_completion - start_step_time, self.per_device_tflops, unet_learning_rate_scheduler(step)
train_metric,
last_step_completion - start_step_time,
self.per_device_tflops,
unet_learning_rate_scheduler(step),
total_weights=num_model_parameters,
)
if self.config.write_metrics:
write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)
Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/trainers/stable_diffusion_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,11 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
new_time = datetime.datetime.now()

train_utils.record_scalar_metrics(
train_metric, new_time - last_step_completion, self.per_device_tflops, unet_learning_rate_scheduler(step)
train_metric,
new_time - last_step_completion,
self.per_device_tflops,
unet_learning_rate_scheduler(step),
total_weights=num_model_parameters,

@Perseus14 Perseus14 Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

num_model_parameters currently counts only unet_state.params, but when train_text_encoder=True, this trainer also updates text_encoder_state.params. Publishing the UNet-only count as total_weights underreports the documented total trainable parameter count.

Please include the text-encoder parameters when that flag is enabled, keeping the calculation outside the training loop. Please also cover both flag values in a focused test; the current metrics tests supply the count directly and cannot catch this omission.

)
if self.config.write_metrics:
train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)
Expand Down
Loading