Support FP8 accelerate config #39370
Conversation
| @@ -1256,6 +1256,14 @@ class AcceleratorConfig: | |||
| Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined | |||
| before calling `TrainingArguments`. If `True`, an `Accelerator` or `PartialState` | |||
| must be initialized. May lead to issues using sweeps or hyperparameter tuning. | |||
| mixed_precision (`str`, *optional*, defaults to `"no"`): | |||
There was a problem hiding this comment.
Can clarify how is it different from the args bf16 and fp16?
There was a problem hiding this comment.
@qgallouedec is this because in training args, that bf16 and fp16 are implicitly mixed precision, vs float16 and bfloat16 in the args?
There was a problem hiding this comment.
It looks like bf16 and fp16 are ultimately used to set ACCELERATE_MIXED_PRECISION, which could more cleanly be passed as mixed_precision, since the AcceleratorConfig exposes this. Digging into the training arguments code, we have:
transformers/src/transformers/training_args.py
Line 1861 in 3d8be20
# if training args is specified, it will override the one specified in the accelerate config
if self.half_precision_backend != "apex":
mixed_precision_dtype = os.environ.get("ACCELERATE_MIXED_PRECISION", "no")
if self.fp16:
mixed_precision_dtype = "fp16"
elif self.bf16:
mixed_precision_dtype = "bf16"
os.environ["ACCELERATE_MIXED_PRECISION"] = mixed_precision_dtypeWe set the fp16, bf16 from the training args, or from the ACCELERATE_MIXED_PRECISION if set, in that priority order. But, we could just use / pass the mixed_precision config as exposed by accelerate instead. fp8 is not set by the existing logic.
| if "fp8_config" in accelerator_config and accelerator_config["fp8_config"] is not None: | ||
| if "backend" in accelerator_config["fp8_config"]: | ||
| recipe_kwargs = MP_BACKEND_TO_KWARGS[accelerator_config["fp8_config"]["backend"]] | ||
| fp8_config = accelerator_config["fp8_config"].copy() | ||
|
|
||
| if fp8_config["backend"] == "AO": | ||
| from torchao.float8 import Float8LinearConfig | ||
|
|
||
| if "recipe_name" in fp8_config: | ||
| recipe_name = fp8_config["recipe_name"] | ||
| fp8_config["config"] = ( | ||
| Float8LinearConfig.from_recipe_name(recipe_name=recipe_name) | ||
| ) | ||
| fp8_config.pop("recipe_name") | ||
| elif "config" in accelerator_config["fp8_config"]: | ||
| config = fp8_config["config"] | ||
| kwargs = {k: v for k, v in config.items() if v is not None} | ||
| fp8_config["config"] = Float8LinearConfig( | ||
| **kwargs | ||
| ) | ||
|
|
||
| fp8_config.pop("backend") | ||
| kwargs_handlers = [recipe_kwargs(**fp8_config)] | ||
| args["kwargs_handlers"] = kwargs_handlers |
There was a problem hiding this comment.
Maybe it will be easier to allow users to pass directly kwargs_handler in accelerate_config ?
That would be great! For the moment we are patching things downstream in axolotl just to get things working. Having this eventually in an upstream release would be nice. |
I'll spend some time very soon on that PR to clean a bit trainer then ! I'll ping you when this is done to have it reviewed ! |
|
Sorry, I forgot to add the link of the PR I was talking about. Here you go #37259 |
What does this PR do?
Adds config options for mixed precision and fp8 config, which are supported in accelerate's Accelerator object.
Also parses the
configfield for the torchao ("AO") fp8 backend from dictionary values into the requiredFloat8LinearConfigobject.This requires also a simple gating change to accelerate, which is actually covered by an existing PR: https://github.com/huggingface/accelerate/pull/3677/files#diff-2d7515874eaecac2687c7fc1a9c720be53f802bf14b4c3dcebe14ad443d075dcR501-R505.
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
TODO
Note
I've tested this downstream in
axolotland find that it works. Performance benefits can be had for certain models when settingtorch_compile: true; it's not yet clear to me which models + hyperparameter settings. I'll add some quick numbers here to demonstrate this.