Update wan configs for training

- Ensure `adam_weight_decay` is a float.
- Add `tensorboard_dir` parameter for logging.

Co-authored-by: martinarroyo <martinarroyo@google.com>

Author: ninatu

src/maxdiffusion/configs/base_wan_14b.yml

@@ -145,6 +145,7 @@ diffusion_scheduler_config: {
 # Output directory
 # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
 base_output_directory: ""
+tensorboard_dir: ""
 
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
@@ -300,7 +301,7 @@ save_optimizer: False
 adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
 adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients.
 adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
-adam_weight_decay: 0 # AdamW Weight decay
+adam_weight_decay: 0.0 # AdamW Weight decay
 max_grad_norm: 1.0
 
 enable_profiler: False

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -122,6 +122,7 @@ diffusion_scheduler_config: {
 # Output directory
 # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
 base_output_directory: ""
+tensorboard_dir: ""
 
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
@@ -256,7 +257,7 @@ save_optimizer: False
 adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
 adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients.
 adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
-adam_weight_decay: 0 # AdamW Weight decay
+adam_weight_decay: 0.0 # AdamW Weight decay
 max_grad_norm: 1.0
 
 enable_profiler: False

src/maxdiffusion/configs/base_wan_27b.yml

@@ -133,6 +133,7 @@ diffusion_scheduler_config: {
 # Output directory
 # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
 base_output_directory: ""
+tensorboard_dir: ""
 
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
@@ -267,7 +268,7 @@ save_optimizer: False
 adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
 adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients.
 adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
-adam_weight_decay: 0 # AdamW Weight decay
+adam_weight_decay: 0.0 # AdamW Weight decay
 max_grad_norm: 1.0
 
 enable_profiler: False

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -128,6 +128,7 @@ diffusion_scheduler_config: {
 # Output directory
 # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
 base_output_directory: ""
+tensorboard_dir: ""
 
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
@@ -262,7 +263,7 @@ save_optimizer: False
 adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
 adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients.
 adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
-adam_weight_decay: 0 # AdamW Weight decay
+adam_weight_decay: 0.0 # AdamW Weight decay
 max_grad_norm: 1.0
 
 enable_profiler: False

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -129,6 +129,7 @@ diffusion_scheduler_config: {
 # Output directory
 # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
 base_output_directory: ""
+tensorboard_dir: ""
 
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
@@ -263,7 +264,7 @@ save_optimizer: False
 adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
 adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients.
 adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
-adam_weight_decay: 0 # AdamW Weight decay
+adam_weight_decay: 0.0 # AdamW Weight decay
 max_grad_norm: 1.0
 
 enable_profiler: False