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)