freeze_vae Flag May Not Take Effect in Current Training Pipeline?

Hi Bagel Team,

While reviewing the training code, I noticed that the freeze_vae flag might not have the intended effect in the current implementation.

Specifically, in the training setup:

```python
if training_args.freeze_vae and training_args.visual_gen:
for param in vae_model.parameters():
param.requires_grad = False
```

This will correctly set requires_grad=False for VAE parameters, but from my reading of the code:

vae_model.parameters() is never passed to the optimizer(s) in the training pipeline;

only model (wrapped with fsdp_wrapper) is passed to the optimizer;

vae_model is loaded separately and does not seem to be included in fsdp_wrapper or any optimizer group.

This means that even if freeze_vae=False, VAE parameters will not be updated, because they are not in any optimizer parameter group in the first place.
So the freeze_vae flag might effectively be a no-op in the current setup.

Could you clarify:

Is this intentional (VAE is always frozen in the current release)?

If not intentional, should vae_model.parameters() be included in the optimizer (or a separate opt_vae) when freeze_vae=False?

Should vae_model also be wrapped in fsdp_wrapper for consistency with other modules if it’s meant to be trainable?

Thanks!

Start by tracing the training setup where vae_model is created, the optimizer parameter groups, and fsdp_wrapper usage. Verify whether vae_model.parameters() ever reaches an optimizer and whether freeze_vae is expected to control training. Done requires a maintainer decision on the intended behavior and corresponding tests or documented confirmation.

Written by the indexing model from the issue text.