[llama4] Change expert_bias and tokens_per_expert to non-persistent buffer (#1403)

As titled. 

Tested on llama4 debugging model (dp=8, ep=2):
<img width="1188" height="226" alt="Screenshot 2025-07-15 at 8 05 12 PM"
src="https://github.com/user-attachments/assets/24a1bf87-b038-481e-b40b-96e2123c96fc"
/>

Author: wwwjn

torchtitan/experiments/llama4/model/moe.py

@@ -249,12 +249,10 @@ def __init__(self, model_args: TransformerModelArgs):
             self.register_buffer(
                 "expert_bias",
                 torch.zeros(num_experts, dtype=torch.float32),
-                persistent=True,
             )
             self.register_buffer(
                 "tokens_per_expert",
                 torch.zeros(num_experts, dtype=torch.float32),
-                persistent=True,
             )
         else:
             self.expert_bias = None

torchtitan/models/deepseek_v3/model/moe.py

@@ -290,12 +290,10 @@ def __init__(self, model_args: DeepSeekV3ModelArgs):
             self.register_buffer(
                 "expert_bias",
                 torch.zeros(num_experts, dtype=torch.float32),
-                persistent=True,
             )
             self.register_buffer(
                 "tokens_per_expert",
                 torch.zeros(num_experts, dtype=torch.float32),
-                persistent=True,
             )
         else:
             self.expert_bias = None