Wan checkpointing

[coolkp]

Fixed issue of >flash_block_sizes in config or cli, causes saving the checkpoint to crash.
Changed behavior

• Parse objects of dataclasses correctly
• Skip with warning incompatible fields (this is harmless since checkpointing is for checkpointing mainly model config and weights and restore args for checkpoint only use that.)
  Added test for verifying.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid