Use the string "off" instead of None to disable mixed precision

Author: cbalioglu

setup.py

@@ -58,7 +58,7 @@
         "numpy~=1.23",
         "packaging~=24.1",
         "psutil~=5.9",
-        "pyyaml~=6.0",
+        "ruamel.yaml~=0.18",
         "rich~=13.7",
         "sacrebleu~=2.4",
         "tiktoken~=0.7",

src/fairseq2/recipes/common/_model.py

@@ -230,7 +230,7 @@ def load(self, recipe_config: object, gangs: Gangs) -> Model:
 
         # Load the model.
         trainer_section = get_config_section(recipe_config, "trainer", TrainerSection)
-        if trainer_section.mixed_precision is None:
+        if trainer_section.mixed_precision == "off":
             dtype = trainer_section.dtype
         else:
             dtype = torch.float32
@@ -359,7 +359,7 @@ def load(self, recipe_config: object, gangs: Gangs) -> Model:
 
         # Load the model.
         trainer_section = get_config_section(recipe_config, "trainer", TrainerSection)
-        if trainer_section.mixed_precision is None:
+        if trainer_section.mixed_precision == "off":
             dtype = trainer_section.dtype
         else:
             dtype = torch.float32
@@ -498,7 +498,7 @@ def load(self, recipe_config: object, gangs: Gangs) -> Model:
 
         # Create the model.
         trainer_section = get_config_section(recipe_config, "trainer", TrainerSection)
-        if trainer_section.mixed_precision is None:
+        if trainer_section.mixed_precision == "off":
             dtype = trainer_section.dtype
         else:
             dtype = torch.float32

src/fairseq2/recipes/config.py

@@ -118,7 +118,7 @@ class TrainerSection:
 
     fsdp: FsdpSection = field(default_factory=lambda: FsdpSection())
 
-    mixed_precision: Literal["static", "dynamic"] | None = "static"
+    mixed_precision: Literal["static", "dynamic", "off"] = "static"
     """
     If 'none', the whole training will be run in `dtype`. If 'static', forward
     and backward passes will be run in `dtype`, but the optimizer step will be

src/fairseq2/recipes/wav2vec2/asr/_train.py

@@ -244,7 +244,7 @@ def load_wav2vec2_asr_trainer(
             config.pretrained_model.name,
             gangs,
             config.trainer.dtype,
-            mp=config.trainer.mixed_precision is not None,
+            mp=config.trainer.mixed_precision != "off",
         )
 
         pt_module = cast(Wav2Vec2Model, pt_model.module)

src/fairseq2/utils/yaml.py

@@ -10,9 +10,9 @@
 from pathlib import Path
 from typing import IO, TypeAlias, final
 
-import yaml
+from ruamel.yaml import YAML
+from ruamel.yaml.error import YAMLError
 from typing_extensions import override
-from yaml import YAMLError
 
 from fairseq2.utils.file import FileMode, FileSystem
 
@@ -32,9 +32,12 @@ def dump(self, obj: object, output: Path | IO[str]) -> None: ...
 
 @final
 class StandardYamlLoader(YamlLoader):
+    _yaml: YAML
     _file_system: FileSystem
 
     def __init__(self, file_system: FileSystem) -> None:
+        self._yaml = YAML(typ="safe", pure=True)
+
         self._file_system = file_system
 
     @override
@@ -47,16 +50,22 @@ def load(self, input_: Path | IO[str]) -> list[object]:
             finally:
                 fp.close()
 
-        itr = yaml.safe_load_all(input_)
+        itr = self._yaml.load_all(input_)
 
         return list(itr)
 
 
 @final
 class StandardYamlDumper(YamlDumper):
+    _yaml: YAML
     _file_system: FileSystem
 
     def __init__(self, file_system: FileSystem) -> None:
+        self._yaml = YAML(typ="safe", pure=True)
+
+        self._yaml.default_flow_style = False
+        self._yaml.sort_base_mapping_type_on_output = False  # type: ignore[assignment]
+
         self._file_system = file_system
 
     @override
@@ -69,8 +78,10 @@ def dump(self, obj: object, output: Path | IO[str]) -> None:
             finally:
                 fp.close()
         else:
-            yaml.safe_dump(obj, output, sort_keys=False)
+            self._yaml.dump(obj, output)
 
 
 def read_yaml(s: str) -> object:
-    return yaml.safe_load(s)
+    yaml = YAML(typ="safe", pure=True)
+
+    return yaml.load(s)