Support block_size and bf16_stochastic_round keyword arguments for torchao optimizers (#42972)

feat(trainer): support block_size and bf16_stochastic_round for torchao optimizers

Author: klae01

src/transformers/trainer.py

@@ -1671,6 +1671,12 @@ def optimizer_hook(param):
                 optimizer_cls = AdamW8bit
             else:
                 raise ValueError("Invalid optimizer")
+            optimizer_kwargs.update(
+                {
+                    "block_size": optim_args.get("block_size", 256),
+                    "bf16_stochastic_round": strtobool(optim_args.get("bf16_stochastic_round", "False")),
+                }
+            )
             optimizer_kwargs.update(adam_kwargs)
         elif args.optim in [
             OptimizerNames.SCHEDULE_FREE_RADAM,