ronantakizawa/codereview-bench huggingface dataset:
Repository: microsoft/DeepSpeed
File: deepspeed/utils/config.py (Python)

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

# DeepSpeed Team

from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject
from deepspeed.utils.constants import *


class DeepSpeedThroughputTimerConfig(DeepSpeedConfigObject):

    def __init__(self, param_dict):
        super(DeepSpeedThroughputTimerConfig, self).__init__()

        self.enabled = None
        self.synchronized = None

        timers_dict = {}
        if param_dict and TIMERS in param_dict:
            if TIMERS_THROUGHPUT in param_dict[TIMERS]:
                timers_dict = param_dict[TIMERS][TIMERS_THROUGHPUT]

        self._initialize(timers_dict)

    def _initialize(self, param_dict):
        self.enabled = get_scalar_param(param_dict, TIMERS_THROUGHPUT_ENABLED, TIMERS_THROUGHPUT_ENABLED_DEFAULT)
        self.synchronized = get_scalar_param(param_dict, TIMERS_THROUGHPUT_SYNCHRONIZED,
                                             TIMERS_THROUGHPUT_SYNCHRONIZED_DEFAULT)


Review comment (suggestion):
This should use the `DeepSpeedConfigModel` rather than `DeepSpeedConfigObject` to take advantage of config validation via pydantic: https://github.com/microsoft/DeepSpeed/blob/4d4c99d9e51afa01a8f4fc7d59a0bc09128abf3e/deepspeed/runtime/config_utils.py#L16

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

# DeepSpeed Team

from deepspeed.runtime.config_utils import DeepSpeedConfigModel

#########################################
# Timers
#########################################
# Timers. By default, timers are enabled.
# Users can configure in ds_config.json as below example:
TIMERS_FORMAT = '''
Timers should be enabled as:
"timers": {
  "throughput": {
    "enabled": true,
    "synchronized": true
  }
}
'''

TIMERS = "timers"
TIMERS_THROUGHPUT = "throughput"


def get_timers_config(param_dict):
    if param_dict and TIMERS in param_dict and TIMERS_THROUGHPUT in param_dict[TIMERS]:
        timers_config_dict = param_dict[TIMERS][TIMERS_THROUGHPUT]
    else:
        timers_config_dict = {}
    return DeepSpeedThroughputTimerConfig(**timers_config_dict)


class DeepSpeedThroughputTimerConfig(DeepSpeedConfigModel):
Repository: microsoft/DeepSpeed
File: accelerator/cpu_accelerator.py (Python)

Code before:
            from op_builder import __deepspeed__  # noqa: F401 # type: ignore
            from op_builder.cpu import CCLCommBuilder, ShareMemCommBuilder, FusedAdamBuilder, CPUAdamBuilder, NotImplementedBuilder
        except ImportError:
            from deepspeed.ops.op_builder.cpu import CCLCommBuilder, ShareMemCommBuilder, FusedAdamBuilder, CPUAdamBuilder, NotImplementedBuilder

        if class_name == "CCLCommBuilder":
            return CCLCommBuilder
        elif class_name == "ShareMemCommBuilder":
            return ShareMemCommBuilder
        elif class_name == "FusedAdamBuilder":
            return FusedAdamBuilder
        elif class_name == "CPUAdamBuilder":
            return CPUAdamBuilder
        else:
            # return a NotImplementedBuilder to avoid get NoneType[Name] in unit tests
            return NotImplementedBuilder

    def build_extension(self):
        from torch.utils.cpp_extension import BuildExtension
        return BuildExtension

    def export_envs(self):
        return []

    def visible_devices_envs(self):
        return ['CUDA_VISIBLE_DEVICES']

    def set_visible_devices_envs(self, current_env, local_accelerator_ids):
        for env in self.visible_devices_envs():
            current_env[env] = ",".join(map(str, local_accelerator_ids))


Review comment (suggestion):
```suggestion
    # TODO: cpu's visible envs is confirmed, keep as CUDA_VISIBLE_DEVICES
    def visible_devices_envs(self):
```

Code after:
            from op_builder import __deepspeed__  # noqa: F401 # type: ignore
            from op_builder.cpu import CCLCommBuilder, ShareMemCommBuilder, FusedAdamBuilder, CPUAdamBuilder, NotImplementedBuilder
        except ImportError:
            from deepspeed.ops.op_builder.cpu import CCLCommBuilder, ShareMemCommBuilder, FusedAdamBuilder, CPUAdamBuilder, NotImplementedBuilder

        if class_name == "CCLCommBuilder":
            return CCLCommBuilder
        elif class_name == "ShareMemCommBuilder":
            return ShareMemCommBuilder
        elif class_name == "FusedAdamBuilder":
            return FusedAdamBuilder
        elif class_name == "CPUAdamBuilder":
            return CPUAdamBuilder
        else:
            # return a NotImplementedBuilder to avoid get NoneType[Name] in unit tests
            return NotImplementedBuilder

    def build_extension(self):
        from torch.utils.cpp_extension import BuildExtension
        return BuildExtension

    def export_envs(self):
        return []

    # TODO: cpu's visible envs is confirmed, keep as CUDA_VISIBLE_DEVICES
    def visible_devices_envs(self):
        return ['CUDA_VISIBLE_DEVICES']

    def set_visible_devices_envs(self, current_env, local_accelerator_ids):
        for env in self.visible_devices_envs():
            current_env[env] = ",".join(map(str, local_accelerator_ids))
Repository: microsoft/DeepSpeed
File: accelerator/mps_accelerator.py (Python)

Code before:
        except ImportError:
            return "deepspeed.ops.op_builder"

    # create an instance of op builder, specified by class_name
    def create_op_builder(self, op_name):
        builder_class = self.get_op_builder(op_name)
        if builder_class is not None:
            return builder_class()
        return None

    # return an op builder class, specified by class_name
    def get_op_builder(self, class_name):
        from deepspeed.ops.op_builder.cpu import NotImplementedBuilder

        return NotImplementedBuilder

    def build_extension(self):
        from torch.utils.cpp_extension import BuildExtension

        return BuildExtension

    def export_envs(self):
        return []

    def visible_devices_envs(self):
        # TODO: could not find visible devices env for mps
        return ['CUDA_VISIBLE_DEVICES']

    def set_visible_devices_envs(self, current_env, local_accelerator_ids):
        for env in self.visible_devices_envs():
            current_env[env] = ",".join(map(str, local_accelerator_ids))


Review comment (suggestion):
```suggestion
    # TODO: mpu's visible envs is confirmed, keep as CUDA_VISIBLE_DEVICES
    def visible_devices_envs(self):
```

Code after:
        except ImportError:
            return "deepspeed.ops.op_builder"

    # create an instance of op builder, specified by class_name
    def create_op_builder(self, op_name):
        builder_class = self.get_op_builder(op_name)
        if builder_class is not None:
            return builder_class()
        return None

    # return an op builder class, specified by class_name
    def get_op_builder(self, class_name):
        from deepspeed.ops.op_builder.cpu import NotImplementedBuilder

        return NotImplementedBuilder

    def build_extension(self):
        from torch.utils.cpp_extension import BuildExtension

        return BuildExtension

    def export_envs(self):
        return []

    # TODO: mpu's visible envs is confirmed, keep as CUDA_VISIBLE_DEVICES
    def visible_devices_envs(self):
        # TODO: could not find visible devices env for mps
        return ['CUDA_VISIBLE_DEVICES']

    def set_visible_devices_envs(self, current_env, local_accelerator_ids):
        for env in self.visible_devices_envs():
            current_env[env] = ",".join(map(str, local_accelerator_ids))
Repository: microsoft/DeepSpeed
File: accelerator/xpu_accelerator.py (Python)

Code before:
    # Graph operations
    def create_graph(self):
        return None

    def capture_to_graph(self, graph, pool=None, stream=None):
        from deepspeed.runtime.utils import noop_context
        return noop_context()

    def replay_graph(self, graph):
        return

    # Data types
    def is_bf16_supported(self):
        return True

    def is_fp16_supported(self):
        return True

    def supported_dtypes(self):
        return [torch.float, torch.half, torch.bfloat16]

    # Tensor operations

    @property
    def BFloat16Tensor(self):
        return functools.partial(torch.tensor, dtype=torch.bfloat16, device='xpu')

    @property
    def ByteTensor(self):
        return functools.partial(torch.tensor, dtype=torch.uint8, device='xpu')

    @property
    def DoubleTensor(self):
        return functools.partial(torch.tensor, dtype=torch.double, device='xpu')

    @property
    def FloatTensor(self):
        return functools.partial(torch.tensor, dtype=torch.float, device='xpu')

    @property
    def HalfTensor(self):
        return functools.partial(torch.tensor, dtype=torch.half, device='xpu')

    @property
    def IntTensor(self):
        return functools.partial(torch.tensor, dtype=torch.int, device='xpu')

    @property
    def LongTensor(self):
        return functools.partial(torch.tensor, dtype=torch.long, device='xpu')


Review comment (suggestion):
```suggestion
        return functools.partial(torch.tensor, dtype=torch.bfloat16, device=self.name)
```

Code after:
    # Graph operations
    def create_graph(self):
        return None

    def capture_to_graph(self, graph, pool=None, stream=None):
        from deepspeed.runtime.utils import noop_context
        return noop_context()

    def replay_graph(self, graph):
        return

    # Data types
    def is_bf16_supported(self):
        return True

    def is_fp16_supported(self):
        return True

    def supported_dtypes(self):
        return [torch.float, torch.half, torch.bfloat16]

    # Tensor operations

    @property
    def BFloat16Tensor(self):
        return functools.partial(torch.tensor, dtype=torch.bfloat16, device=self._name)

    @property
    def ByteTensor(self):
        return functools.partial(torch.tensor, dtype=torch.uint8, device=self._name)

    @property
    def DoubleTensor(self):
        return functools.partial(torch.tensor, dtype=torch.double, device=self._name)

    @property
    def FloatTensor(self):
        return functools.partial(torch.tensor, dtype=torch.float, device=self._name)

    @property
    def HalfTensor(self):
        return functools.partial(torch.tensor, dtype=torch.half, device=self._name)

    @property
    def IntTensor(self):
        return functools.partial(torch.tensor, dtype=torch.int, device=self._name)

    @property
    def LongTensor(self):
        return functools.partial(torch.tensor, dtype=torch.long, device=self._name)
Repository: microsoft/DeepSpeed
File: deepspeed/runtime/engine.py (Python)

Code before:
            self.checkpoint_engine.save(state_dict, path)

        self.checkpoint_engine.commit(tag)

        return True

    def empty_partition_cache(self):
        """
        Release GPU memory consumed by offloaded model parameters.
        """
        if hasattr(self.optimizer, 'empty_partition_cache'):
            self.optimizer.empty_partition_cache()
            gc.collect()
            get_accelerator().empty_cache()

    def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwargs={}) -> None:
        """Compile the module using the specified backend and kwargs.
        If a compiler_fn is set, it will be used instead of torch.compile().
        """
        if not is_compile_supported():
            raise RuntimeError("compile is not supported in your version of PyTorch.")

        if self.is_compiled:
            return

        self.module = self.module.compile(backend=backend, **compile_kwargs)
        self._is_compiled = True

    @property
    def is_compiled(self) -> bool:
        return self._is_compiled


Review comment (suggestion):
```suggestion
        self.module.compile(backend=backend, **compile_kwargs)
```
nn.Module.compile returns None

Code after:

        if dist.get_rank() == 0:
            self.checkpoint_engine.makedirs(save_dir, exist_ok=True)
            logger.info(f"Saving model weights to {path}, tag: {tag}")
            self.checkpoint_engine.save(state_dict, path)

        self.checkpoint_engine.commit(tag)

        return True

    def empty_partition_cache(self):
        """
        Release GPU memory consumed by offloaded model parameters.
        """
        if hasattr(self.optimizer, 'empty_partition_cache'):
            self.optimizer.empty_partition_cache()
            gc.collect()
            get_accelerator().empty_cache()

    def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwargs={}) -> None:
        """Compile the module using the specified backend and kwargs.
        If a compiler_fn is set, it will be used instead of torch.compile().
        """
        if not is_compile_supported():
            raise RuntimeError("compile is not supported in your version of PyTorch.")

        if self.is_compiled:
            return

        self.module.compile(backend=backend, **compile_kwargs)
        self._is_compiled = True

    @property
    def is_compiled(self) -> bool:
        return self._is_compiled
Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zero/offload_config.py (Python)

Code before:
    For tile-based optimizer step processing, overlap read of next tile with
    computation of current tile. Used in ZeRO-Infinity.
    """

    pipeline_write: bool = False
    """
    For tile-based optimizer step processing, overlap write of previous tile
    with computation of current tile.
    """

    fast_init: bool = False
    """ Enable fast optimizer initialization when offloading to NVMe. """

    ratio: float = Field(1.0, ge=0.0, le=1.0)
    """ Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3."""

    @model_validator(mode="after")
    def set_pipeline(self):
        pipeline = self.pipeline_read or self.pipeline_write
        self.__dict__["pipeline"] = pipeline
        return self


class OffloadStateTypeEnum(str, Enum):
    """ Enum for internal buffer types """
    opt_states = "opt_states"
    hp_params = "hp_params"
    lp_params = "lp_params"
    lp_grads = "lp_grads"
    contiguous_grad_buffer = "contiguous_grad_buffer"


Review comment (suggestion):
```suggestion
    optim_states = "optim_states"
```

Code after:
    For tile-based optimizer step processing, overlap read of next tile with
    computation of current tile. Used in ZeRO-Infinity.
    """

    pipeline_write: bool = False
    """
    For tile-based optimizer step processing, overlap write of previous tile
    with computation of current tile.
    """

    fast_init: bool = False
    """ Enable fast optimizer initialization when offloading to NVMe. """

    ratio: float = Field(1.0, ge=0.0, le=1.0)
    """ Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3."""

    @model_validator(mode="after")
    def set_pipeline(self):
        pipeline = self.pipeline_read or self.pipeline_write
        self.__dict__["pipeline"] = pipeline
        return self


class OffloadStateTypeEnum(str, Enum):
    """ Enum for internal buffer types """
    optim_states = "optim_states"
    hp_params = "hp_params"
    lp_params = "lp_params"
    lp_grads = "lp_grads"
    contiguous_grad_buffer = "contiguous_grad_buffer"
Repository: microsoft/DeepSpeed
File: tests/unit/ops/transformer/inference/inference_test_utils.py (Python)

Code before:
import torch
from deepspeed.accelerator import get_accelerator

TOLERANCES = None


def get_tolerances():
    global TOLERANCES
    if TOLERANCES is None:
        TOLERANCES = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}
        if get_accelerator().is_bf16_supported():
            # Note: BF16 tolerance is higher than FP16 because of the lower precision (7 (+1) bits vs
            # 10 (+1) bits)
            TOLERANCES[torch.bfloat16] = (4.8e-1, 3.2e-2)
    return TOLERANCES


DTYPES = None


def get_dtypes():
    global DTYPES
    if DTYPES is None:
        DTYPES = [torch.float32]
        try:
            if get_accelerator().is_fp16_supported():
                DTYPES.append(torch.float16)
            if get_accelerator().is_bf16_supported():
                DTYPES.append(torch.bfloat16)
        except (AssertionError, AttributeError):
            pass
    return DTYPES


def allclose(x, y):
    assert x.dtype == y.dtype
    rtol, atol = get_tolerances()[x.dtype]
    return torch.allclose(x, y, rtol=rtol, atol=atol)


def assert_almost_equal(x, y, decimal=2, err_msg=''):
    import numpy.testing as npt
    if isinstance(x, torch.Tensor):
        if x.dtype == torch.bfloat16:
            x = x.float()
        x = x.cpu().detach().numpy()
    if isinstance(y, torch.Tensor):
        if y.dtype == torch.bfloat16:
            y = y.float()
        y = y.cpu().detach().numpy()
    npt.assert_array_almost_equal(x, y, err_msg=err_msg, decimal=decimal)

Review comment (suggestion):
Would the following not be more robust?
```suggestion
            DTYPES += get_accelerator().supported_dtypes()
```

Code after:
import torch
from deepspeed.accelerator import get_accelerator

TOLERANCES = None


def get_tolerances():
    global TOLERANCES
    if TOLERANCES is None:
        TOLERANCES = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}
        if get_accelerator().is_bf16_supported():
            # Note: BF16 tolerance is higher than FP16 because of the lower precision (7 (+1) bits vs
            # 10 (+1) bits)
            TOLERANCES[torch.bfloat16] = (4.8e-1, 3.2e-2)
    return TOLERANCES


DTYPES = None


def get_dtypes():
    global DTYPES
    if DTYPES is None:
        DTYPES = get_accelerator().supported_dtypes()
    return DTYPES


def allclose(x, y):
    assert x.dtype == y.dtype
    rtol, atol = get_tolerances()[x.dtype]
    return torch.allclose(x, y, rtol=rtol, atol=atol)


def assert_almost_equal(x, y, decimal=2, err_msg=''):
    import numpy.testing as npt
    if isinstance(x, torch.Tensor):
        if x.dtype == torch.bfloat16:
            x = x.float()
        x = x.cpu().detach().numpy()
    if isinstance(y, torch.Tensor):
        if y.dtype == torch.bfloat16:
            y = y.float()
        y = y.cpu().detach().numpy()
    npt.assert_array_almost_equal(x, y, err_msg=err_msg, decimal=decimal)


def max_diff(a, b):
    a = a.to(torch.float32).flatten()
    b = b.to(torch.float32).flatten()
    diff = torch.abs(a - b)
    max_diff_indices = torch.argsort(diff)[-1]
Repository: microsoft/DeepSpeed
File: accelerator/mlu_accelerator.py (Python)

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

# DeepSpeed Team
import importlib
import inspect
import functools

from .abstract_accelerator import DeepSpeedAccelerator
import torch
# During setup stage torch may not be installed, pass on no torch will
# allow op builder related API to be executed.
try:
    import torch_mlu
except ImportError:
    pass


class MLU_Accelerator(DeepSpeedAccelerator):

    def __init__(self):
        self._name = 'mlu'
        self._communication_backend_name = 'cncl'
        self._compile_backend = "inductor"
        self.class_dict = None


Review comment (suggestion):
Don't you need to add `mlu` references 
1. [here](https://github.com/microsoft/DeepSpeed/blob/9b7fc5452471392b0f58844219fcfdd14a9cdc77/accelerator/real_accelerator.py#L23) 
2. [here](https://github.com/microsoft/DeepSpeed/blob/9b7fc5452471392b0f58844219fcfdd14a9cdc77/accelerator/real_accelerator.py#L51)

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

# DeepSpeed Team
import importlib
import inspect
import functools

from .abstract_accelerator import DeepSpeedAccelerator
import torch
# During setup stage torch may not be installed, pass on no torch will
# allow op builder related API to be executed.


class MLU_Accelerator(DeepSpeedAccelerator):

    def __init__(self):
        self._name = 'mlu'
        self._communication_backend_name = 'cncl'
        self._compile_backend = "inductor"
        self.class_dict = None

    def is_synchronized_device(self):
        return False

    def use_host_timers(self):
Repository: microsoft/DeepSpeed
File: deepspeed/utils/logging.py (Python)

Code before:

import functools
import logging
import sys
import os
import torch

log_levels = {
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warning": logging.WARNING,
    "error": logging.ERROR,
    "critical": logging.CRITICAL,
}


class LoggerFactory:

    def create_warning_filter(logger):
        warn = False

        def warn_once(record):
            nonlocal warn
            if torch.compiler.is_compiling() and not warn:
                warn = True
                logger.warning("It is recommended to set DISABLE_LOGS_WHILE_COMPILING=1,"
                               "to avoid graph breaks caused by logger.")
            return True

        return warn_once

    @staticmethod
    def logging_decorator(func):

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if torch._dynamo.is_compiling():
                return
            else:
                return func(*args, **kwargs)

        return wrapper

    @staticmethod
    def create_logger(name=None, level=logging.INFO):
        """create a logger

        Args:
            name (str): name of the logger
            level: level of logger


Review comment (suggestion):
Rephrase the warning to reflect that it applies to compile-mode execution and control by env var. For example, 
```suggestion
                logger.warning("To avoid graph breaks in compile-mode, it is recommended to disable logging by setting env var DISABLE_LOGS_WHILE_COMPILING=1,"
```

Code after:

import functools
import logging
import sys
import os
import torch
from deepspeed.runtime.compiler import is_compile_supported

log_levels = {
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warning": logging.WARNING,
    "error": logging.ERROR,
    "critical": logging.CRITICAL,
}


class LoggerFactory:

    def create_warning_filter(logger):
        warn = False

        def warn_once(record):
            nonlocal warn
            if is_compile_supported() and torch.compiler.is_compiling() and not warn:
                warn = True
                logger.warning("To avoid graph breaks caused by logger in compile-mode, it is recommended to"
                               " disable logging by setting env var DISABLE_LOGS_WHILE_COMPILING=1")
            return True

        return warn_once

    @staticmethod
    def logging_decorator(func):

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if torch.compiler.is_compiling():
                return
            else:
                return func(*args, **kwargs)

        return wrapper

    @staticmethod
    def create_logger(name=None, level=logging.INFO):
        """create a logger

        Args:
            name (str): name of the logger
            level: level of logger