[megatron] Accept dtype-string optimizer_config_kwargs (coerce exp_avg_dtype etc. to torch.dtype)

docs/content/docs/configuration/config.mdx

@@ -213,6 +213,23 @@ Some rules for configuring these parameters:
 `optimizer_config_kwargs.use_precision_aware_optimizer=true` can cause checkpointing to fail. See: https://github.com/nvidia/megatron-lm/issues/1820. We recommend leaving this setting to `false`.
 </Callout>
 
+`optimizer_config_kwargs` accepts string values for Megatron `*_dtype` fields:
+
+```yaml
+optimizer_config_kwargs:
+  use_precision_aware_optimizer: true
+  exp_avg_dtype: bf16
+  exp_avg_sq_dtype: fp8
+  main_params_dtype: fp32
+```
+
+Accepted names are case-insensitive: `fp32` (`float32`, `float`), `fp16` (`float16`, `half`), `bf16` (`bfloat16`), and `fp8` (`float8`, `uint8`). `fp8` maps to `torch.uint8`, matching TransformerEngine optimizer state storage.
+
+Field-specific checks:
+
+- `main_params_dtype` (master weights): `fp32`, `fp16`
+- `exp_avg_dtype` / `exp_avg_sq_dtype`: `fp32`, `fp16`, `bf16`, `fp8`
+
 ## Optimizer Configuration
 
 For both the critic and policy model, we provide a common optimizer configuration

docs/content/docs/examples/megatron.mdx

@@ -85,10 +85,22 @@ empty_cuda_cache: true
 
 These default values can be overridden by passing in the corresponding arguments to `trainer.policy.megatron_config` in the launch script.
 
+`optimizer_config_kwargs` can set optimizer-state dtypes from YAML:
+
+```yaml
+optimizer_config_kwargs:
+  use_precision_aware_optimizer: true
+  exp_avg_dtype: bf16
+  exp_avg_sq_dtype: fp8
+  main_params_dtype: fp32
+```
+
+See the [Megatron configuration guide](../configuration/config#megatron-configuration) for accepted aliases and per-field checks.
+
 ## Parallelism Resources
 
 Understanding and configuring parallelism strategies for large models can be challenging.
 Some helpful resources for understanding and tuning large scale parallelism strategies can be found at the [Huggingface Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=finding_the_best_training_configuration), 
 the [The Mesh Parallelism Zoo](https://blog.ezyang.com/2025/08/the-parallelism-mesh-zoo/), and the [Visualizing 6-D Parallelism](https://main-horse.github.io/posts/visualizing-6d).
 
-Below, we show a diagram displaying how all 5 parallelism strategies - tensor, pipeline, context, expert, and data parallelism - can be utilized in SkyRL, as well as how dispatching data across these parallel groups works.
+Below, we show a diagram displaying how all 5 parallelism strategies - tensor, pipeline, context, expert, and data parallelism - can be utilized in SkyRL, as well as how dispatching data across these parallel groups works.

skyrl/backends/skyrl_train/distributed/megatron/optimizer.py

@@ -26,6 +26,9 @@
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from omegaconf import DictConfig
 
+from skyrl.backends.skyrl_train.distributed.megatron.optimizer_dtype import (
+    coerce_optimizer_dtype_kwargs,
+)
 from skyrl.train.config import OptimizerConfig as SkyRLOptimizerConfig
 
 
@@ -45,7 +48,8 @@ def init_megatron_optim_config(
         "params_dtype": torch.bfloat16,
         "use_distributed_optimizer": True,
     }
-    optim_args.update(optimizer_config_kwargs)
+    # YAML dtype overrides arrive as strings; Megatron expects torch.dtype.
+    optim_args.update(coerce_optimizer_dtype_kwargs(optimizer_config_kwargs))
 
     config = OptimizerConfig(**optim_args)
     return config

skyrl/backends/skyrl_train/distributed/megatron/optimizer_dtype.py

@@ -0,0 +1,64 @@
+"""Torch-only coercion for Megatron optimizer dtype kwargs."""
+
+from typing import Any, Dict, Set
+
+import torch
+
+# Megatron short names plus common YAML spellings. TE stores FP8 optimizer state
+# as uint8, matching Megatron-LM's dtype map.
+_DTYPE_NAME_TO_TORCH: Dict[str, torch.dtype] = {
+    "fp32": torch.float32,
+    "float32": torch.float32,
+    "float": torch.float32,
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+    "fp16": torch.float16,
+    "float16": torch.float16,
+    "half": torch.float16,
+    "fp8": torch.uint8,
+    "float8": torch.uint8,
+    "uint8": torch.uint8,
+}
+
+# Only TE FusedAdam-backed fields get field-specific checks. ``main_grads_dtype``
+# is not forwarded at the pinned megatron-core rev, so it is coerced only and
+# left to ``OptimizerConfig.__post_init__``.
+_LEGAL_FIELD_DTYPES: Dict[str, Set[torch.dtype]] = {
+    "main_params_dtype": {torch.float32, torch.float16},
+    "exp_avg_dtype": {torch.float32, torch.bfloat16, torch.float16, torch.uint8},
+    "exp_avg_sq_dtype": {torch.float32, torch.bfloat16, torch.float16, torch.uint8},
+}
+
+
+def coerce_optimizer_dtype_kwargs(optimizer_config_kwargs: Dict[str, Any] | None) -> Dict[str, Any]:
+    """Return kwargs with recognized ``*_dtype`` strings converted to ``torch.dtype``."""
+    if optimizer_config_kwargs is None:
+        return {}
+
+    coerced: Dict[str, Any] = {}
+    for key, value in optimizer_config_kwargs.items():
+        if not key.endswith("_dtype"):
+            coerced[key] = value
+            continue
+
+        if isinstance(value, torch.dtype):
+            dtype = value
+        elif isinstance(value, str):
+            name = value.strip().lower()
+            if name not in _DTYPE_NAME_TO_TORCH:
+                raise ValueError(
+                    f"Unrecognized dtype name {value!r} for optimizer kwarg {key!r}. "
+                    f"Expected one of {sorted(_DTYPE_NAME_TO_TORCH)} or a torch.dtype."
+                )
+            dtype = _DTYPE_NAME_TO_TORCH[name]
+        else:
+            # Let Megatron validate non-string, non-dtype values.
+            coerced[key] = value
+            continue
+
+        legal = _LEGAL_FIELD_DTYPES.get(key)
+        if legal is not None and dtype not in legal:
+            legal_names = sorted({n for n, d in _DTYPE_NAME_TO_TORCH.items() if d in legal})
+            raise ValueError(f"Illegal dtype {dtype} for optimizer kwarg {key!r}; legal values are {legal_names}.")
+        coerced[key] = dtype
+    return coerced

tests/backends/skyrl_train/distributed/test_optimizer_dtype_coercion.py

@@ -0,0 +1,57 @@
+"""Tests for Megatron optimizer dtype coercion."""
+
+import importlib
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+
+@pytest.fixture(scope="module")
+def coerce_optimizer_dtype_kwargs():
+    module_name = "skyrl.backends.skyrl_train.distributed.megatron.optimizer_dtype"
+    with patch.dict(sys.modules, {"megatron": None}):
+        sys.modules.pop(module_name, None)
+        module = importlib.import_module(module_name)
+    return module.coerce_optimizer_dtype_kwargs
+
+
+def test_coerces_dtype_strings_and_preserves_other_kwargs(coerce_optimizer_dtype_kwargs):
+    kwargs = {
+        "exp_avg_dtype": "  BF16 ",
+        "exp_avg_sq_dtype": "fp8",
+        "main_params_dtype": "fp16",
+        "params_dtype": "float32",
+        "main_grads_dtype": "bfloat16",
+        "use_precision_aware_optimizer": True,
+    }
+
+    out = coerce_optimizer_dtype_kwargs(kwargs)
+
+    assert out == {
+        "exp_avg_dtype": torch.bfloat16,
+        "exp_avg_sq_dtype": torch.uint8,
+        "main_params_dtype": torch.float16,
+        "params_dtype": torch.float32,
+        "main_grads_dtype": torch.bfloat16,
+        "use_precision_aware_optimizer": True,
+    }
+    assert kwargs["exp_avg_dtype"] == "  BF16 "
+
+
+@pytest.mark.parametrize(
+    "kwargs,match",
+    [
+        ({"exp_avg_dtype": "bf17"}, "Unrecognized dtype name"),
+        ({"main_params_dtype": "bf16"}, "main_params_dtype"),
+    ],
+)
+def test_rejects_unknown_or_field_illegal_dtype_names(coerce_optimizer_dtype_kwargs, kwargs, match):
+    with pytest.raises(ValueError, match=match):
+        coerce_optimizer_dtype_kwargs(kwargs)
+
+
+def test_none_mapping_is_empty_but_none_field_values_pass_through(coerce_optimizer_dtype_kwargs):
+    assert coerce_optimizer_dtype_kwargs(None) == {}
+    assert coerce_optimizer_dtype_kwargs({"main_grads_dtype": None}) == {"main_grads_dtype": None}