ronantakizawa/codereview-bench huggingface dataset:
Repository: microsoft/DeepSpeed
File: tests/unit/runtime/zero/test_offload_states.py (Python)

Code before:
class TestOffloadStates(DistributedTest):
    # Need multiple gpus to test possible hanging
    world_size = 2

    def test_offload_states(self, included_state, pin_memory, non_blocking):
        hidden_dim = 1024

        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-6
                }
            },
            "zero_optimization": {
                "stage": 3,
            }
        }
        config_dict["bf16"] = {"enabled": True}

        with deepspeed.zero.Init(config_dict_or_path=config_dict):
            model = SimpleModel(hidden_dim, nlayers=4)

        params = list(model.parameters())
        param_groups = [{"params": params[::2], "weight_decay": 0.0}, {"params": params[1::2], "weight_decay": 0.1}]

        include = None if included_state is None else [included_state]
        run_model(model, param_groups, config_dict, hidden_dim, torch.bfloat16, include, pin_memory, non_blocking)


Review comment (suggestion):
I recommend improving readability by making things more explicit like so:

```python
  param_groups = [
      {"params": [p for n, p in model.named_parameters() if not 'bias' in n], "weight_decay" : 0.1}, 
      {"params": [p for n, p in model.named_parameters() if 'bias' in n], "weight_decay" : 0.0}
  ]
```

On a minor note, I also flipped the weight decay to weights instead of biases :). 

Code after:
class TestOffloadStates(DistributedTest):
    # Need multiple gpus to test possible hanging
    world_size = 2

    def test_offload_states(self, included_state, pin_memory, non_blocking):
        hidden_dim = 1024

        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-6
                }
            },
            "zero_optimization": {
                "stage": 3,
            }
        }
        config_dict["bf16"] = {"enabled": True}

        with deepspeed.zero.Init(config_dict_or_path=config_dict):
            model = SimpleModel(hidden_dim, nlayers=4)

        param_groups = [{
            "params": [p for n, p in model.named_parameters() if not 'bias' in n],
            "weight_decay": 0.1
        }, {
            "params": [p for n, p in model.named_parameters() if 'bias' in n],
            "weight_decay": 0.0
        }]
        include = None if included_state is None else [included_state]
        run_model(model, param_groups, config_dict, hidden_dim, torch.bfloat16, include, pin_memory, non_blocking)
Repository: microsoft/DeepSpeed
File: deepspeed/utils/nvtx.py (Python)

Code before:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

from deepspeed.accelerator import get_accelerator
import torch

enable_nvtx = True


def instrument_w_nvtx(func):
    """Decorator that records an NVTX range for the duration of the function call.
       Skips NVTX instrumentation when torch.compile is active to avoid graph breaks.
    """

    def wrapped_fn(*args, **kwargs):
        if enable_nvtx and not torch.is_compiling():
            get_accelerator().range_push(func.__qualname__)
        ret_val = func(*args, **kwargs)
        if enable_nvtx and not torch.is_compiling():
            get_accelerator().range_pop()
        return ret_val

    return wrapped_fn


Review comment (refactor):
Better to utilize DeepSpeed utility 
```suggestion
       from deepspeed.runtime.compiler import is_compiling
        if enable_nvtx and not is_compiling():
```

Code after:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

from deepspeed.accelerator import get_accelerator
from deepspeed.runtime.compiler import is_compiling

enable_nvtx = True


def instrument_w_nvtx(func):
    """Decorator that records an NVTX range for the duration of the function call.
       Skips NVTX instrumentation when torch.compile is active to avoid graph breaks.
    """

    def wrapped_fn(*args, **kwargs):
        if enable_nvtx and not is_compiling():
            get_accelerator().range_push(func.__qualname__)
        ret_val = func(*args, **kwargs)
        if enable_nvtx and not is_compiling():
            get_accelerator().range_pop()
        return ret_val

    return wrapped_fn
Repository: microsoft/DeepSpeed
File: deepspeed/runtime/checkpoint_engine/checkpoint_engine.py (Python)

Code before:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os


class CheckpointEngine(object):

    # init checkpoint engine for save/load
    def __init__(self, config_params=None):
        pass

    def create(self, tag):
        # create checkpoint on give tag for save/load.
        pass

    def makedirs(self, path, exist_ok=False):
        os.makedirs(path, exist_ok=exist_ok)

    def save(self, state_dict, path: str):
        pass

    def load(self, path: str, map_location=None):
        pass

    def commit(self, tag):
        # to tell checkpoint services if all files are ready.
        pass

    def wait(self):
        # To wait in asynchronous checkpoint engines (e.g. DataStates-LLM) for the previous snapshot to finish
        pass

    def preserves_storage_sharing(self):
        # Check if the checkpoint engine preserves storage sharing
        # (set to false if cloning is required to get actual tensor sizes)
        return False


Review comment (suggestion):
Thanks for adding this API. But I think the meaning is inverted here in the sense that `preserves_storage_sharing` is what leads to checkpoint blot and requires cloning to fix. Please see the following torch docs. I think it also helpful to add the doc link here. 
https://pytorch.org/docs/stable/notes/serialization.html#saving-and-loading-tensors-preserves-views

Further reading of the doc on my part makes me feel that `preserves_tensor_views()` might be a descriptive name. I am curious to know your thoughts. Thanks! 

Code after:

    @abc.abstractmethod
    def load(self, path: str, map_location=None):
        ...

    @abc.abstractmethod
    def commit(self, info: CheckpointCommitInfo):
        # to tell checkpoint services if all files are ready.
        ...

    def is_data_parallel_writer(self, dp_rank):
        return dp_rank == 0

    def is_decoupled(self):
        return False

    def set_commit_info(self, info: CheckpointCommitInfo):
        pass

    def get_commit_info(self):
        return None

    def cleanup(self):
        pass

    def preserves_storage_sharing(self):
        return True
Repository: microsoft/DeepSpeed
File: deepspeed/ops/fp_quantizer/quantize.py (Python)

Code before:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
import abc
from abc import ABC

import gc
from deepspeed.ops.op_builder import FPQuantizerBuilder
from deepspeed.accelerator import get_accelerator

fp_quant_module = None


class Quantizer(ABC):
    """
    Abstract Quantizer class that implmenents quantize/dequantize methods.

    Arguments:
        group_size (int, optional): number of values or elements that are grouped
            together for the quantization process.
    """

    def __init__(self, group_size=512) -> None:
        self.group_size = group_size

    @abc.abstractmethod
    def quantize(self,
                 input,
                 q_bits=8,
                 q_mantisa_bits=3,
                 stochastic_mode=False,
                 return_meta_tensor=False) -> torch.Tensor:
        ...

    @abc.abstractmethod
    def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor:
        ...


class FP_Quantize(Quantizer):


Review comment (refactor):
```suggestion
    Abstract Quantizer class that implements quantize/dequantize methods.
```

Code after:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
import abc
from abc import ABC

import gc
from deepspeed.ops.op_builder import FPQuantizerBuilder
from deepspeed.accelerator import get_accelerator

fp_quant_module = None


class Quantizer(ABC):
    """
    Abstract Quantizer class that implements quantize/dequantize methods.

    Arguments:
        group_size (int, optional): number of values or elements that are grouped
            together for the quantization process.
    """

    def __init__(self, group_size=512) -> None:
        self.group_size = group_size

    @abc.abstractmethod
    def quantize(self,
                 input,
                 q_bits=8,
                 q_mantisa_bits=3,
                 stochastic_mode=False,
                 return_meta_tensor=False) -> torch.Tensor:
        ...

    @abc.abstractmethod
    def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor:
        ...


class FP_Quantize(Quantizer):
Repository: microsoft/DeepSpeed
File: op_builder/cpu/comm.py (Python)

Code before:
    def extra_ldflags(self):
        ccl_root_path = os.environ.get("CCL_ROOT")
        if ccl_root_path is None:
            raise ValueError(
                "Didn't find CCL_ROOT, install oneCCL from https://github.com/oneapi-src/oneCCL and source its environment variable"
            )
        else:
            return ['-lccl', f'-L{ccl_root_path}/lib']


class ShareMemCommBuilder(CPUOpBuilder):
    BUILD_VAR = "DS_BUILD_SHM_COMM"
    NAME = "deepspeed_shm_comm"
    machine = platform.machine().lower()

    def __init__(self, name=None):
        name = self.NAME if name is None else name
        super().__init__(name=name)

    def absolute_name(self):
        return f'deepspeed.ops.comm.{self.NAME}_op'

    def sources(self):
        src_files = ['csrc/cpu/comm/shm_interface.cpp', 'csrc/cpu/comm/shm.cpp']
        if self.machine == 'riscv64':
            src_files.append('csrc/cpu/comm/riscv64/shm_rvv.cpp')
        else:
            src_files.append('csrc/cpu/comm/x86_64/shm_avx512.cpp')
        return src_files

    def include_paths(self):
        includes = ['csrc/cpu/includes']
        return includes

    def cxx_args(self):
        return ['-O2', '-fopenmp']

    def is_compatible(self, verbose=False):
        # TODO: add soft compatibility check for private binary release.
        #  a soft check, as in we know it can be trivially changed.
        return super().is_compatible(verbose)


Review comment (suggestion):
Since directory (`riscv64`, `x86_64`) already distinguish different architecture, suggest to use same file name under the directory for symmetry.  An example looks like the following.
```
# a line of code assert self.machine in ['riscv64', 'x86_64']
src_files.append(f'csrc/cpu/comm/{self.machine}/shm_impl.cpp')
```


Code after:
        ccl_root_path = os.environ.get("CCL_ROOT")
        if ccl_root_path is None:
            raise ValueError(
                "Didn't find CCL_ROOT, install oneCCL from https://github.com/oneapi-src/oneCCL and source its environment variable"
            )
        else:
            return ['-lccl', f'-L{ccl_root_path}/lib']


class ShareMemCommBuilder(CPUOpBuilder):
    BUILD_VAR = "DS_BUILD_SHM_COMM"
    NAME = "deepspeed_shm_comm"

    def __init__(self, name=None):
        name = self.NAME if name is None else name
        super().__init__(name=name)

    def absolute_name(self):
        return f'deepspeed.ops.comm.{self.NAME}_op'

    def sources(self):
        return ['csrc/cpu/comm/shm_interface.cpp', 'csrc/cpu/comm/shm.cpp']

    def include_paths(self):
        includes = ['csrc/cpu/includes']
        return includes

    def cxx_args(self):
        return ['-O2', '-fopenmp']

    def is_compatible(self, verbose=False):
        # TODO: add soft compatibility check for private binary release.
        #  a soft check, as in we know it can be trivially changed.
        return super().is_compatible(verbose)
Repository: microsoft/DeepSpeed
File: tests/unit/runtime/zero/test_zero_autocast.py (Python)

Code before:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

from copy import deepcopy

import pytest

import torch
from torch.nn.parallel import DistributedDataParallel as DDP

from unit.common import DistributedTest, enable_determinism, reduce_boolean_flags
from unit.simple_model import SimpleModel
from unit.util import bf16_required_version_check

import deepspeed
from deepspeed.accelerator import get_accelerator
from deepspeed.runtime.zero import GatheredParameters
from deepspeed.runtime.torch_autocast import PARAM_COMM_DTYPE_ATTR_NAME, get_comm_dtype
from deepspeed.utils import logger

RTOL = 0.1
ATOL = 0.0


def cls_to_qualname(cls):
    return f"{cls.__module__}.{cls.__name__}"


class SimpleModelWithLayerNorm(torch.nn.Module):

    def __init__(self, hidden_dim, nlayers=1):
        super(SimpleModelWithLayerNorm, self).__init__()
        self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for i in range(nlayers)])
        self.norm = torch.nn.LayerNorm(hidden_dim)
        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()

    def forward(self, x, y):
        x = self.linears[0](x)
        x = self.norm(x)
        return self.cross_entropy_loss(x, y)


def step_amp(enabled, baseline_model, baseline_optimizer, target_engine, dtype, enable_autocast_outside,
             baseline_scaler, step, x, y, rtol, atol, expect_match):

Review comment (bug):
I don't think `logger` is the right tool for tests since how do you even control the log levels. 

why not just `assert`? if the local loss doesn't match that's already a reason to fail.

Code after:
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

from copy import deepcopy

import pytest

import torch
from torch.nn.parallel import DistributedDataParallel as DDP

from unit.common import DistributedTest, enable_determinism, reduce_boolean_flags
from unit.simple_model import SimpleModel
from unit.util import bf16_required_version_check

import deepspeed
from deepspeed.accelerator import get_accelerator
from deepspeed.runtime.zero import GatheredParameters
from deepspeed.runtime.torch_autocast import PARAM_COMM_DTYPE_ATTR_NAME, get_comm_dtype

RTOL = 0.1
ATOL = 0.0


def cls_to_qualname(cls):
    return f"{cls.__module__}.{cls.__name__}"


class SimpleModelWithLayerNorm(torch.nn.Module):

    def __init__(self, hidden_dim, nlayers=1):
        super(SimpleModelWithLayerNorm, self).__init__()
        self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for i in range(nlayers)])
        self.norm = torch.nn.LayerNorm(hidden_dim)
        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()

    def forward(self, x, y):
        x = self.linears[0](x)
        x = self.norm(x)
        return self.cross_entropy_loss(x, y)


def step_amp(enabled, baseline_model, baseline_optimizer, target_engine, dtype, enable_autocast_outside,
             baseline_scaler, step, x, y, rtol, atol, expect_match):
    device_type = get_accelerator().device_name()
Repository: microsoft/PowerToys
File: src/modules/launcher/Plugins/Microsoft.PowerToys.Run.Plugin.Service/Helpers/ServiceHelper.cs (C#)

Code before:
            }
        }

        private static string GetLocalizedErrorMessage(Action action)
        {
            if (action == Action.Start)
            {
                return Resources.wox_plugin_service_start_error_notification;
            }
            else if (action == Action.Stop)
            {
                return Resources.wox_plugin_service_stop_error_notification;
            }
            else if (action == Action.Restart)
            {
                return Resources.wox_plugin_service_restart_error_notification;
            }
            else
            {
                return string.Empty;
            }
        }

        private static bool IsDelayedStart(string serviceName)
        {
           return (int)Registry.LocalMachine.OpenSubKey(@"System\CurrentControlSet\Services\" + serviceName, false)?.GetValue("DelayedAutostart", 0, RegistryValueOptions.None) == 1;
        }
    }
}


Review comment (bug):
Can we add a try-catch here?
Don't want the plugin to crash if for some reason the subkey related to the service doesn't exists.

Code after:
            }
        }

        private static string GetLocalizedErrorMessage(Action action)
        {
            if (action == Action.Start)
            {
                return Resources.wox_plugin_service_start_error_notification;
            }
            else if (action == Action.Stop)
            {
                return Resources.wox_plugin_service_stop_error_notification;
            }
            else if (action == Action.Restart)
            {
                return Resources.wox_plugin_service_restart_error_notification;
            }
            else
            {
                return string.Empty;
            }
        }

        private static bool IsDelayedStart(string serviceName)
        {
           return (int?)Registry.LocalMachine.OpenSubKey(@"System\CurrentControlSet\Services\" + serviceName, false)?.GetValue("DelayedAutostart", 0, RegistryValueOptions.None) == 1;
        }
    }
}
Repository: microsoft/PowerToys
File: src/common/ManagedCommon/ThemeHelpers.cs (C#)

Code before:
            return (AppTheme)value;
        }

        public static Theme GetCurrentTheme()
        {
            string theme = (string)Registry.GetValue($"{HKeyRoot}\\{HkeyWindowsTheme}", HValueCurrentTheme, string.Empty);
            theme = theme.Split('\\').Last().Split('.').First().ToString();

            switch (theme)
            {
                case "hc1":
                    return Theme.HighContrastOne;
                case "hc2":
                    return Theme.HighContrastTwo;
                case "hcwhite":
                    return Theme.HighContrastWhite;
                case "hcblack":
                    return Theme.HighContrastBlack;
                default:
                    return Theme.System;
            }
        }

        public static bool SupportsImmersiveDarkMode()
        {
            return Environment.OSVersion.Version.Major >= 10 && Environment.OSVersion.Version.Build >= 18985;
        }

        public static void SetImmersiveDarkMode(IntPtr window, bool enabled)
        {
            if (SupportsImmersiveDarkMode())
            {
                int useImmersiveDarkMode = enabled ? 1 : 0;
                _ = DwmSetWindowAttribute(window, 20, ref useImmersiveDarkMode, sizeof(int));
            }
        }
    }
}


Review comment (bug):
It is generally better to simply try to call the API and see if it fails. If it fails, don't pursue it. :smile:

Code after:
﻿// Copyright (c) Microsoft Corporation
// The Microsoft Corporation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System;
using System.Linq;
using System.Runtime.InteropServices;
using Microsoft.Win32;

namespace ManagedCommon
{
    // Based on https://stackoverflow.com/a/62811758/5001796
    public static class ThemeHelpers
    {
        [DllImport("dwmapi.dll")]
        private static extern int DwmSetWindowAttribute(IntPtr hwnd, int attr, ref int attrValue, int attrSize);

        internal const string HKeyRoot = "HKEY_CURRENT_USER";
        internal const string HkeyWindowsTheme = @"SOFTWARE\Microsoft\Windows\CurrentVersion\Themes";
        internal const string HkeyWindowsPersonalizeTheme = $@"{HkeyWindowsTheme}\Personalize";
        internal const string HValueAppTheme = "AppsUseLightTheme";
        internal const int DWMWAImmersiveDarkMode = 20;

        // based on https://stackoverflow.com/questions/51334674/how-to-detect-windows-10-light-dark-mode-in-win32-application
        public static AppTheme GetAppTheme()
        {
            int value = (int)Registry.GetValue($"{HKeyRoot}\\{HkeyWindowsPersonalizeTheme}", HValueAppTheme, 1);
            return (AppTheme)value;
        }

        public static void SetImmersiveDarkMode(IntPtr window, bool enabled)
        {
            int useImmersiveDarkMode = enabled ? 1 : 0;
            _ = DwmSetWindowAttribute(window, DWMWAImmersiveDarkMode, ref useImmersiveDarkMode, sizeof(int));
        }
    }
}
Repository: microsoft/PowerToys
File: src/common/Themes/theme_helpers.cpp (C++)

Code before:
// Port Based on https://stackoverflow.com/a/62811758/5001796
#include "theme_helpers.h"
#include "dwmapi.h"
#include <windows.h>
#include <iostream>
#include <vector>

#define STATUS_SUCCESS 0x00000000
#define DWMWA_USE_IMMERSIVE_DARK_MODE 20
#define HKEY_WINDOWS_THEME L"Software\\Microsoft\\Windows\\CurrentVersion\\Themes\\Personalize"

typedef NTSTATUS(WINAPI* RtlGetVersionPtr)(PRTL_OSVERSIONINFOW);

namespace
{
    RTL_OSVERSIONINFOW GetRealOSVersion()
    {
        HMODULE hMod = ::GetModuleHandleW(L"ntdll.dll");
        if (hMod)
        {
            RtlGetVersionPtr fxPtr = (RtlGetVersionPtr)::GetProcAddress(hMod, "RtlGetVersion");
            if (fxPtr != nullptr)
            {
                RTL_OSVERSIONINFOW info = { 0 };
                info.dwOSVersionInfoSize = sizeof(info);
                if (STATUS_SUCCESS == fxPtr(&info))
                {
                    return info;
                }
            }
        }
        RTL_OSVERSIONINFOW info = { 0 };
        return info;
    }
}

// based on https://stackoverflow.com/questions/51334674/how-to-detect-windows-10-light-dark-mode-in-win32-application
AppTheme ThemeHelpers::GetAppTheme()
{
    // The value is expected to be a REG_DWORD, which is a signed 32-bit little-endian
    auto buffer = std::vector<char>(4);

Review comment (suggestion):
Prefer VerifyVersionInfoW and _make sure your app manifest says you run on Windows 10_. There's no need to use internal ntdll APIs.

Code after:
// Port Based on https://stackoverflow.com/a/62811758/5001796
#include "theme_helpers.h"
#include "dwmapi.h"
#include <windows.h>
#include <vector>

#define DWMWA_USE_IMMERSIVE_DARK_MODE 20
#define HKEY_WINDOWS_THEME L"Software\\Microsoft\\Windows\\CurrentVersion\\Themes\\Personalize"

// based on https://stackoverflow.com/questions/51334674/how-to-detect-windows-10-light-dark-mode-in-win32-application
AppTheme ThemeHelpers::GetAppTheme()
{
    // The value is expected to be a REG_DWORD, which is a signed 32-bit little-endian
    auto buffer = std::vector<char>(4);
    auto cbData = static_cast<DWORD>(buffer.size() * sizeof(char));
    auto res = RegGetValueW(
        HKEY_CURRENT_USER,
        HKEY_WINDOWS_THEME,
        L"AppsUseLightTheme",
        RRF_RT_REG_DWORD, // expected value type
        nullptr,
        buffer.data(),
        &cbData);

    if (res != ERROR_SUCCESS)
    {
        return AppTheme::Light;
    }

    // convert bytes written to our buffer to an int, assuming little-endian
    auto i = int(buffer[3] << 24 |
                 buffer[2] << 16 |
                 buffer[1] << 8 |
                 buffer[0]);

    return AppTheme(i);
}

void ThemeHelpers::SetImmersiveDarkMode(HWND window, bool enabled)
{
    int useImmersiveDarkMode = enabled ? 1 : 0;