Clarify HBM memory estimation log

[Shuwen-Fang]

hbm memory estimation is estimated see https://b.corp.google.com/issues/530231381 comments for details

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[NuojCheng]

could you add one more line saying compiler could over-estimate the HBM usage?

[shuningjin]

Thanks for the change! Could you also update this doc with your finding?

maxtext/docs/guides/monitoring_and_debugging/understand_logs_and_metrics.md

Lines 157 to 176 in f13ccd1

	### 3.1. Memory analysis
	We first perform a "dry run" compilation of a training step to [analyze its memory requirement](https://github.com/AI-Hypercomputer/maxtext/blob/28e5097ac467ed8b1d17676d68aa5acc50f9d60d/src/MaxText/train.py#L380-L382). This static analysis is performed by the XLA compiler. The log outputs [memory sizes](https://github.com/AI-Hypercomputer/maxtext/blob/28e5097ac467ed8b1d17676d68aa5acc50f9d60d/src/MaxText/max_utils.py#L672-L690):
	```none
	Total memory size: 100.4 GB, Output size: 44.5 GB, Temp size: 55.9 GB, Argument size: 44.5 GB, Host temp size: 0.0 GB.
	```
	The most important number is `Total memory size: 100.4 GB`. This is the total High Bandwidth Memory (HBM) the TPU device needs to execute the program. Here is a breakdown:
	- `Argument size: 44.5 GB`: This is the memory needed to hold the inputs for your function. This typically includes the batch of data, parameter (master copy), and optimizer state (e.g., moment).
	- `Output size: 44.5 GB`: This is the space required to store the results of the computation, such as the updated model weights and updated optimizer states.
	- `Temp size: 55.9 GB`: This is the "scratch space" memory. It's used for all the intermediate values created during the forward and backward passes that are discarded once the step is complete. This includes activation (forward pass), gradient (backward pass), and parameter (working copy, if mixed precision).
	- Memory aliasing: You might notice that the sum of the parts is greater than `100.4 GB (total)`: `44.5 GB (Argument) + 44.5 GB (Output) + 55.9 GB (Temp) = 144.9 GB`. The difference is due to a compiler optimization called memory aliasing. The compiler reuses memory blocks. The true calculation is `Total = Argument + Output + Temp - Aliased`. In our case, the compiler identified `44.5 GB (144.9 GB - 100.4 GB)` of memory that could be safely reused. Most likely, it reuses memory for `Argument` and `Output`.
	In addition, it shows temporary memory used on the host CPU. In this case, `Host temp size: 0.0 GB`, indicating that all the significant memory allocation happens on the accelerator device.
	### 3.2. Memory snapshot
	The previous section is a forecast of memory usage for entire training step, based on static analysis of the compiled code from the XLA compiler. To see the actual memory usage, we now turn to a real-time snapshot from the JAX runtime, captured right after the training state is initialized.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!