chore: import upstream snapshot with attribution

nemo/core/optim/flash_optim.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from nemo.utils import logging
+
+
+def patch_flashoptim_uneven_shard_support(optimizer) -> None:
+    """Patch flashoptim to handle FSDP2 unevenly-sharded parameters in DCP.
+
+    FlashOptim <= 0.1.3 raises ``ValueError`` when saving optimizer state for
+    parameters whose shard dimension is not evenly divisible by the FSDP mesh
+    size. The root cause is that ``DTensor.from_local()`` is called without an
+    explicit ``shape``, so it infers ``global = local * mesh_size`` which is
+    wrong for padded (uneven) shards.
+
+    This patch replaces ``_wrap_state_as_dtensor`` on the optimizer class so
+    that ``shape=param.shape`` and ``stride=param.stride()`` are always passed,
+    which is correct for both even and uneven shards.
+    """
+
+    klass = type(optimizer)
+    if not hasattr(klass, "_wrap_state_as_dtensor"):
+        return
+    if getattr(klass, "_nemo_patched_uneven_shard", False):
+        return
+
+    @staticmethod
+    def _fixed_wrap_state_as_dtensor(state: dict[str, Any], param: torch.Tensor) -> None:  # noqa: UP006
+        if not hasattr(param, "device_mesh"):
+            return
+
+        from torch.distributed.tensor import DTensor
+
+        mesh = param.device_mesh
+        placements = param.placements
+
+        for key, val in state.items():
+            if isinstance(val, torch.Tensor) and not isinstance(val, DTensor) and val.dim() > 0:
+                state[key] = DTensor.from_local(
+                    val,
+                    mesh,
+                    placements,
+                    shape=param.shape,
+                    stride=param.stride(),
+                )
+
+    klass._wrap_state_as_dtensor = _fixed_wrap_state_as_dtensor
+    klass._nemo_patched_uneven_shard = True
+    logging.info("Patched flashoptim %s to support unevenly-sharded FSDP2 parameters in DCP.", klass.__name__)