Wire total_weights into ML Diagnostics metrics - #474
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request simplifies metric reporting by removing the dependency on metric_types from google_cloud_mldiagnostics and directly mapping internal metrics to canonical string names. It also removes the gradient norm metric and updates several trainers (including base_wan, dreambooth, flux, sdxl, and stable_diffusion) to explicitly pass total_weights=num_model_parameters when recording scalar metrics. Corresponding updates were made to the documentation and tests. There are no review comments, so I have no feedback to provide.
- Wire total_weights into telemetry: Pass the existing `num_model_parameters` calculation into `record_scalar_metrics(..., total_weights=num_model_parameters)` across all trainers (Stable Diffusion, SDXL, Flux, Wan, and DreamBooth). This populates the total_weights card in Google Cloud ML Diagnostics. - Remove the redundant `metric_types` import and `if/else` branching. Standardize `_METRICS_TO_MANAGED` directly on canonical string literals, matching the SDK's internal representation.
11a98ed to
90c3ee5
Compare
Perseus14
left a comment
Perseus14
left a comment
There was a problem hiding this comment.
Thanks for the PR. Most of it looks good, added a couple of comments. PTAL!
| new_time - last_step_completion, | ||
| self.per_device_tflops, | ||
| unet_learning_rate_scheduler(step), | ||
| total_weights=num_model_parameters, |
There was a problem hiding this comment.
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.
| new_time - last_step_completion, | ||
| self.per_device_tflops, | ||
| learning_rate_scheduler(step), | ||
| total_weights=num_model_parameters, |
There was a problem hiding this comment.
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.
2 participants
num_model_parameterscalculation intorecord_scalar_metrics(..., total_weights=num_model_parameters)across all trainers (Stable Diffusion, SDXL, Flux, Wan, and DreamBooth). This populates the total_weights card in Google Cloud ML Diagnostics.metric_typesimport andif/elsebranching. Standardize_METRICS_TO_MANAGEDdirectly on canonical string literals, matching the SDK's internal representation.