This patch most notably fixes an issue with an optional dependency (optax), which resulted in parsing errors with poetry. It contains the following fixes:
- fix optax dep issue
 - remove offload_state_dict from kwargs
 - Fix bnb fsdp loading for pre-quantized checkpoint (#41415)
 - Fix tests fsdp (#41422)
 - Fix trainer for py3.9 (#41359)