fix fsdp error from unused register token module

[Jubeku]

Only create the auxiliary-token parameter when aux tokens are actually used.
With num_register_tokens == num_class_tokens == 0, a zero-size parameter still participates in autograd and yields an empty gradient, which makes FSDP2's reduce-scatter fail with "RuntimeError: _chunk_cat expects non-empty tensor".

Description

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[clessig]

@shmh40 : can you have a look? q_aux comes from develop-ssl.

[clessig]

Can we also add this:

diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py
index d67ceaed..9a52aede 100644
--- a/src/weathergen/model/model_interface.py
+++ b/src/weathergen/model/model_interface.py
@@ -94,7 +94,7 @@ def init_model_and_shard(
             find_unused_parameters=cf.get("ddp_find_unused_parameters", True),
             gradient_as_bucket_view=True,
             bucket_cap_mb=512,
-            static_graph=cf.get("ddp_static_graph", True),
+            static_graph=cf.get("ddp_static_graph", False),
         )
 
     elif with_ddp and with_fsdp:

It's related to the PR and causes silent degradation of model convergence. Potentially we should remove it at all as a config parameter.