diff --git a/makani/utils/driver.py b/makani/utils/driver.py index 0c954ee6..771daac9 100644 --- a/makani/utils/driver.py +++ b/makani/utils/driver.py @@ -60,6 +60,30 @@ def _log_timers(self): for k,v in self.timers.items(): self.logger.info(f"{print_prefix}{k} [s]: {v:.2f}") + def _watch_model(self, log="all"): + """Register wandb gradient/parameter logging, unless the model will be torch.compiled. + + ``wandb.watch`` installs per-module/per-parameter hooks that compute histograms via + ``.item()``/``.tolist()``. Under ``torch.compile`` these hooks get traced into the graph: + every ``.item()`` is a graph break and the parameter ``name`` is specialized as a + constant, so dynamo recompiles once per parameter until it hits ``recompile_limit`` and + abandons compilation for that frame. The hook-based design is fundamentally incompatible + with a single compiled graph, so we skip watching when ``jit_mode`` enables compilation. + (If gradient/parameter histograms are needed under compile, log them out-of-graph in the + training loop after ``backward()`` instead.) + """ + if not self.log_to_wandb: + return + if self.params.get("jit_mode", "none") == "inductor": + if self.log_to_screen: + self.logger.info( + "Skipping wandb.watch under jit_mode=inductor: its hooks would be traced into " + "the compiled graph and trigger a recompile per parameter (recompile_limit). " + "Gradient/parameter histograms are disabled under torch.compile." + ) + return + wandb.watch(self.model, log=log) + def __init__(self, params: YParams = None, world_rank: Optional[int] = 0, device: Optional[str] = None): # define timer dict self.timers = {} diff --git a/makani/utils/training/autoencoder_trainer.py b/makani/utils/training/autoencoder_trainer.py index b21efa71..d30e419e 100644 --- a/makani/utils/training/autoencoder_trainer.py +++ b/makani/utils/training/autoencoder_trainer.py @@ -132,8 +132,7 @@ def __init__(self, params: Optional[YParams] = None, world_rank: Optional[int] = ddp_process_group = comm.get_group("data") # log gradients to wandb - if self.log_to_wandb: - wandb.watch(self.model, log="all") + self._watch_model(log="all") # print model if self.log_to_screen: diff --git a/makani/utils/training/deterministic_trainer.py b/makani/utils/training/deterministic_trainer.py index 12d787e1..4b8995d7 100644 --- a/makani/utils/training/deterministic_trainer.py +++ b/makani/utils/training/deterministic_trainer.py @@ -148,8 +148,7 @@ def __init__(self, params: Optional[YParams] = None, world_rank: Optional[int] = ddp_process_group = comm.get_group("data") # log gradients to wandb - if self.log_to_wandb: - wandb.watch(self.model, log="all") + self._watch_model(log="all") # print model if self.log_to_screen: diff --git a/makani/utils/training/ensemble_trainer.py b/makani/utils/training/ensemble_trainer.py index 8fa30d33..a54f15d2 100644 --- a/makani/utils/training/ensemble_trainer.py +++ b/makani/utils/training/ensemble_trainer.py @@ -148,9 +148,8 @@ def __init__(self, params: Optional[YParams] = None, world_rank: Optional[int] = if dist.is_initialized() and not self.params.disable_ddp: ddp_process_group = comm.get_group("data") - # log gradients to wandb - if self.log_to_wandb: - wandb.watch(self.model, log="all") + # log gradients to wandb (skipped under torch.compile; see Driver._watch_model) + self._watch_model(log="all") # print model if self.log_to_screen: diff --git a/makani/utils/training/stochastic_trainer.py b/makani/utils/training/stochastic_trainer.py index f418ab00..18f296df 100644 --- a/makani/utils/training/stochastic_trainer.py +++ b/makani/utils/training/stochastic_trainer.py @@ -138,8 +138,7 @@ def __init__(self, params: Optional[YParams] = None, world_rank: Optional[int] = ddp_process_group = comm.get_group("data") # log gradients to wandb - if self.log_to_wandb: - wandb.watch(self.model, log="all") + self._watch_model(log="all") # print model if self.log_to_screen: