diff --git a/.gitattributes b/.gitattributes index b18208408b22377b2b9b4d6979e2c1dcc85c4d4f..be02d6c046deddd6ed163bbe1db49955ac1ff81f 100644 --- a/.gitattributes +++ b/.gitattributes @@ -72,3 +72,6 @@ venv/lib/python3.10/site-packages/wrapt/_wrappers.cpython-310-x86_64-linux-gnu.s venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_processor.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/dataclasses.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_pipeline.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text +venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text +venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text +venv/lib/python3.10/site-packages/hf_transfer/hf_transfer.abi3.so filter=lfs diff=lfs merge=lfs -text diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca9c80e66833e667c4bdf4ec693a4052fd3e55b1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c005c697ef0add5080c18e222594cb3e894ade8e7a6e30d27b2365127f090b9b +size 126436 diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__init__.py b/venv/lib/python3.10/site-packages/accelerate/commands/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c9cbe26c257b515f657c05e1996d517e69613972 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2020 The HuggingFace Team. 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. diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45fa86fa8994780078de93b850a8dfee0613eff1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37abae224e0272336f45b3083b58fa739deb8cfa Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/env.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/env.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e3eb339d899d1abf541b7500c6eb52e1ad798db6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/env.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/estimate.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/estimate.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebfc97b2269681b3c6c494111bcbfac28aa705b6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/estimate.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/launch.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/launch.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb44c46afb04c5dc089f7faecb39498958f80576 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/launch.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/merge.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/merge.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87b67b3194bd1cf4242e9b9e1210a2308201ade5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/merge.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/test.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/test.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b70afeb31a99e87a6440d04e0da8d1b290c9df37 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/test.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22f32ac4eb5b3f5950f147aa698b5daff37735e4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/tpu.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/tpu.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0817c015174cea05b383cb659e5ccf88c54c7915 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/tpu.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..781957c07e932b17206d33e799902bf35ef6f7a9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py b/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py new file mode 100644 index 0000000000000000000000000000000000000000..b878c8debd874e1418b946775b11568c7487ad72 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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 accelerate.commands.config import get_config_parser +from accelerate.commands.env import env_command_parser +from accelerate.commands.estimate import estimate_command_parser +from accelerate.commands.launch import launch_command_parser +from accelerate.commands.merge import merge_command_parser +from accelerate.commands.test import test_command_parser +from accelerate.commands.to_fsdp2 import to_fsdp2_command_parser +from accelerate.commands.tpu import tpu_command_parser +from accelerate.commands.utils import CustomArgumentParser + + +def main(): + parser = CustomArgumentParser("Accelerate CLI tool", usage="accelerate []", allow_abbrev=False) + subparsers = parser.add_subparsers(help="accelerate command helpers") + + # Register commands + get_config_parser(subparsers=subparsers) + estimate_command_parser(subparsers=subparsers) + env_command_parser(subparsers=subparsers) + launch_command_parser(subparsers=subparsers) + merge_command_parser(subparsers=subparsers) + tpu_command_parser(subparsers=subparsers) + test_command_parser(subparsers=subparsers) + to_fsdp2_command_parser(subparsers=subparsers) + + # Let's go + args = parser.parse_args() + + if not hasattr(args, "func"): + parser.print_help() + exit(1) + + # Run + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__init__.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..649a15888cccd070b3d4ca9a600457c6ad59d4d3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/__init__.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import argparse + +from .config import config_command_parser +from .config_args import default_config_file, load_config_from_file # noqa: F401 +from .default import default_command_parser +from .update import update_command_parser + + +def get_config_parser(subparsers=None): + parent_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False) + # The main config parser + config_parser = config_command_parser(subparsers) + # The subparser to add commands to + subcommands = config_parser.add_subparsers(title="subcommands", dest="subcommand") + + # Then add other parsers with the parent parser + default_command_parser(subcommands, parents=[parent_parser]) + update_command_parser(subcommands, parents=[parent_parser]) + + return config_parser + + +def main(): + config_parser = get_config_parser() + args = config_parser.parse_args() + + if not hasattr(args, "func"): + config_parser.print_help() + exit(1) + + # Run + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..397787cf94c392eaa76d515f66acffda3ff0c573 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/cluster.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/cluster.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef96576540ccc60ad331d36ff449e0d42b3a72d6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/cluster.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bbc85f930efff617226439f8bbbedb316b9d81ca Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_args.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_args.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c812228dc8ea3f8906ed38aeefb38bf87631617 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_args.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e63e09bf1cb824846eaef542816047db9f3c6c7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/default.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/default.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f24657f24b3b20d6219e720b856163017b00fa4e Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/default.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b77e6c9816c00bd642ab092430188f370610cce Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/update.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/update.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32abad6203f10d9ca7356c26cab7cdb8fa84aa53 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/update.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/cluster.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/cluster.py new file mode 100644 index 0000000000000000000000000000000000000000..420fa5495ba7d3df064b7d97973db4f181d044ea --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/cluster.py @@ -0,0 +1,869 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import os + +from ...utils import ( + ComputeEnvironment, + DistributedType, + is_deepspeed_available, + is_fp8_available, + is_hpu_available, + is_mlu_available, + is_mps_available, + is_msamp_available, + is_musa_available, + is_npu_available, + is_sdaa_available, + is_transformer_engine_available, + is_transformers_available, + is_xpu_available, +) +from ...utils.constants import ( + DEEPSPEED_MULTINODE_LAUNCHERS, + FSDP2_STATE_DICT_TYPE, + FSDP_AUTO_WRAP_POLICY, + FSDP_BACKWARD_PREFETCH, + FSDP_SHARDING_STRATEGY, + FSDP_STATE_DICT_TYPE, + TORCH_DYNAMO_MODES, +) +from .config_args import ClusterConfig +from .config_utils import ( + DYNAMO_BACKENDS, + _ask_field, + _ask_options, + _convert_distributed_mode, + _convert_dynamo_backend, + _convert_fp8_backend, + _convert_mixed_precision, + _convert_yes_no_to_bool, +) + + +def get_cluster_input(): + distributed_type = _ask_options( + "Which type of machine are you using?", + [ + "No distributed training", + "multi-CPU", + "multi-XPU", + "multi-HPU", + "multi-GPU", + "multi-NPU", + "multi-MLU", + "multi-SDAA", + "multi-MUSA", + "TPU", + ], + _convert_distributed_mode, + ) + + machine_rank = 0 + num_machines = 1 + num_processes = 1 + gpu_ids = None + main_process_ip = None + main_process_port = None + rdzv_backend = "static" + same_network = True + debug = False + + if distributed_type in [ + DistributedType.MULTI_GPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_CPU, + DistributedType.MULTI_HPU, + ]: + num_machines = _ask_field( + "How many different machines will you use (use more than 1 for multi-node training)? [1]: ", + int, + default=1, + ) + if num_machines > 1: + machine_rank = _ask_options( + "What is the rank of this machine?", + list(range(num_machines)), + int, + ) + main_process_ip = _ask_field( + "What is the IP address of the machine that will host the main process? ", + ) + main_process_port = _ask_field( + "What is the port you will use to communicate with the main process? ", + int, + ) + same_network = _ask_field( + "Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on different network hosts [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + if not same_network: + rdzv_backend = _ask_field( + "What rendezvous backend will you use? ('static', 'c10d', ...): ", default="static" + ) + debug = _ask_field( + "Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + if distributed_type == DistributedType.NO: + use_cpu = _ask_field( + "Do you want to run your training on CPU only (even if a GPU / Apple Silicon / Ascend NPU device is available)? [yes/NO]:", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + elif distributed_type == DistributedType.MULTI_CPU: + use_cpu = True + else: + use_cpu = False + + ipex_config = {} + mpirun_config = {} + if use_cpu or is_xpu_available(): + ipex_config["ipex"] = _ask_field( + "Do you want to use Intel PyTorch Extension (IPEX) to speed up training on CPU/XPU? [yes/NO]:", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + if use_cpu: + if distributed_type == DistributedType.MULTI_CPU: + use_mpirun = _ask_field( + "Do you want accelerate to launch mpirun? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_mpirun: + mpirun_hostfile = _ask_field( + "Please enter the path to the hostfile to use with mpirun [~/hostfile]: ", + str, + default="~/hostfile", + ) + mpirun_config["mpirun_hostfile"] = os.path.expanduser(mpirun_hostfile.strip()) + mpirun_config["mpirun_ccl"] = _ask_field("Enter the number of oneCCL worker threads [1]: ", default=1) + + dynamo_config = {} + use_dynamo = _ask_field( + "Do you wish to optimize your script with torch dynamo?[yes/NO]:", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_dynamo: + prefix = "dynamo_" + dynamo_config[prefix + "backend"] = _ask_options( + "Which dynamo backend would you like to use?", + [x.lower() for x in DYNAMO_BACKENDS], + _convert_dynamo_backend, + default=2, + ) + use_custom_options = _ask_field( + "Do you want to customize the defaults sent to torch.compile? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + if use_custom_options: + dynamo_config[prefix + "mode"] = _ask_options( + "Which mode do you want to use?", + TORCH_DYNAMO_MODES, + lambda x: TORCH_DYNAMO_MODES[int(x)], + default=0, + ) + dynamo_config[prefix + "use_fullgraph"] = _ask_field( + "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + dynamo_config[prefix + "use_dynamic"] = _ask_field( + "Do you want to enable dynamic shape tracing? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + dynamo_config[prefix + "use_regional_compilation"] = _ask_field( + "Do you want to enable regional compilation? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + use_mps = not use_cpu and is_mps_available() + deepspeed_config = {} + if ( + distributed_type + in [ + DistributedType.MULTI_GPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + DistributedType.MULTI_NPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.NO, + ] + and not use_mps + ): + use_deepspeed = _ask_field( + "Do you want to use DeepSpeed? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_deepspeed: + distributed_type = DistributedType.DEEPSPEED + assert is_deepspeed_available(), ( + "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source" + ) + + if distributed_type == DistributedType.DEEPSPEED: + use_deepspeed_config = _ask_field( + "Do you want to specify a json file to a DeepSpeed config? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_deepspeed_config: + deepspeed_config["deepspeed_config_file"] = _ask_field( + "Please enter the path to the json DeepSpeed config file: ", + str, + default="none", + ) + else: + deepspeed_config["zero_stage"] = _ask_options( + "What should be your DeepSpeed's ZeRO optimization stage?", + [0, 1, 2, 3], + int, + default=2, + ) + + deepspeed_devices = ["none", "cpu", "nvme"] + if deepspeed_config["zero_stage"] >= 2: + deepspeed_config["offload_optimizer_device"] = _ask_options( + "Where to offload optimizer states?", deepspeed_devices, lambda x: deepspeed_devices[int(x)] + ) + deepspeed_config["offload_param_device"] = _ask_options( + "Where to offload parameters?", deepspeed_devices, lambda x: deepspeed_devices[int(x)] + ) + if deepspeed_config["offload_param_device"] == "nvme": + deepspeed_config["offload_param_nvme_path"] = _ask_field( + "Nvme Path to offload parameters?", + str, + default="/nvme", + ) + if deepspeed_config["offload_optimizer_device"] == "nvme": + deepspeed_config["offload_optimizer_nvme_path"] = _ask_field( + "Nvme Path to offload optimizer states?", + str, + default="/nvme", + ) + deepspeed_config["gradient_accumulation_steps"] = _ask_field( + "How many gradient accumulation steps you're passing in your script? [1]: ", + int, + default=1, + ) + use_gradient_clipping = _ask_field( + "Do you want to use gradient clipping? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_gradient_clipping: + deepspeed_config["gradient_clipping"] = _ask_field( + "What is the gradient clipping value? [1.0]: ", + float, + default=1.0, + ) + if deepspeed_config["zero_stage"] == 3: + deepspeed_config["zero3_save_16bit_model"] = _ask_field( + "Do you want to save 16-bit model weights when using ZeRO Stage-3? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + deepspeed_config["zero3_init_flag"] = _ask_field( + "Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if deepspeed_config["zero3_init_flag"]: + if not is_transformers_available(): + raise Exception( + "When `zero3_init_flag` is set, it requires Transformers to be installed. " + "Please run `pip3 install transformers`." + ) + use_moe = _ask_field( + "Do you want to enable Mixture-of-Experts training (MoE)? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_moe: + deepspeed_config["deepspeed_moe_layer_cls_names"] = _ask_field( + "Specify the comma-separated list of transformers MoE layer class names (case-sensitive), e.g : " + " `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ... : ", + str, + ) + + if num_machines > 1: + launcher_query = "Which Type of launcher do you want to use?" + deepspeed_config["deepspeed_multinode_launcher"] = _ask_options( + launcher_query, + DEEPSPEED_MULTINODE_LAUNCHERS, + lambda x: DEEPSPEED_MULTINODE_LAUNCHERS[int(x)], + ) + + if deepspeed_config["deepspeed_multinode_launcher"] != DEEPSPEED_MULTINODE_LAUNCHERS[1]: + deepspeed_config["deepspeed_hostfile"] = _ask_field( + "DeepSpeed configures multi-node compute resources with hostfile. " + "Each row is of the format `hostname slots=[num_gpus]`, e.g., `localhost slots=2`; " + "for more information please refer official [documentation]" + "(https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node). " + "Please specify the location of hostfile: ", + str, + ) + + is_exclusion_filter = _ask_field( + "Do you want to specify exclusion filter string? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if is_exclusion_filter: + deepspeed_config["deepspeed_exclusion_filter"] = _ask_field( + "DeepSpeed exclusion filter string: ", + str, + ) + + is_inclusion_filter = _ask_field( + "Do you want to specify inclusion filter string? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if is_inclusion_filter: + deepspeed_config["deepspeed_inclusion_filter"] = _ask_field( + "DeepSpeed inclusion filter string: ", + str, + ) + + fsdp_config = {} + + if distributed_type in [ + DistributedType.MULTI_GPU, + DistributedType.MULTI_NPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + ]: + use_fsdp = _ask_field( + "Do you want to use FullyShardedDataParallel? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_fsdp: + distributed_type = DistributedType.FSDP + if distributed_type == DistributedType.FSDP: + fsdp_config["fsdp_version"] = _ask_options( + "What should be your FSDP version? [2]: ", + [1, 2], + lambda x: int(x) + 1, + default=1, + ) + fsdp_version = fsdp_config["fsdp_version"] # extract to a variable to simplify usage later + + if fsdp_version == 1: + sharding_strategy_query = "What should be your sharding strategy?" + fsdp_config["fsdp_reshard_after_forward"] = _ask_options( + sharding_strategy_query, + FSDP_SHARDING_STRATEGY, + lambda x: FSDP_SHARDING_STRATEGY[int(x)], + ) + else: + fsdp_config["fsdp_reshard_after_forward"] = _ask_field( + "Do you want to enable resharding after forward? [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + + fsdp_config["fsdp_offload_params"] = _ask_field( + "Do you want to offload parameters and gradients to CPU? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + fsdp_wrap_query = "What should be your auto wrap policy?" + fsdp_config["fsdp_auto_wrap_policy"] = _ask_options( + fsdp_wrap_query, + FSDP_AUTO_WRAP_POLICY, + lambda x: FSDP_AUTO_WRAP_POLICY[int(x)], + ) + if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]: + use_no_split_modules = _ask_field( + "Do you want to use the model's `_no_split_modules` to wrap. Only applicable for 🤗 Transformers [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if not use_no_split_modules: + fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field( + "Specify the comma-separated list of transformer layer class names (case-sensitive) to wrap ,e.g, :" + "`BertLayer`, `GPTJBlock`, `T5Block`, `BertLayer,BertEmbeddings,BertSelfOutput` ...? : ", + str, + ) + elif fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[1]: + fsdp_config["fsdp_min_num_params"] = _ask_field( + "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ", + int, + default=100000000, + ) + # Removed in FSDP2, ask for user input for FSDP1 + if fsdp_version == 1: + fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?" + fsdp_config["fsdp_backward_prefetch"] = _ask_options( + fsdp_backward_prefetch_query, + FSDP_BACKWARD_PREFETCH, + lambda x: FSDP_BACKWARD_PREFETCH[int(x)], + ) + + fsdp_state_dict_type_query = "What should be your FSDP's state dict type?" + fsdp_config["fsdp_state_dict_type"] = _ask_options( + fsdp_state_dict_type_query, + FSDP_STATE_DICT_TYPE if fsdp_version == 1 else FSDP2_STATE_DICT_TYPE, + lambda x: FSDP_STATE_DICT_TYPE[int(x)] if fsdp_version == 1 else FSDP2_STATE_DICT_TYPE[int(x)], + default=0, + ) + # Not implemented in FSDP2, ask for user input for FSDP1 + if fsdp_version == 1: + fsdp_config["fsdp_forward_prefetch"] = _ask_field( + "Do you want to enable FSDP's forward prefetch policy? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + # Obsolete in FSDP2, ask for user input for FSDP1 + if fsdp_version == 1: + fsdp_config["fsdp_use_orig_params"] = _ask_field( + "Do you want to enable FSDP's `use_orig_params` feature? [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + fsdp_config["fsdp_cpu_ram_efficient_loading"] = _ask_field( + "Do you want to enable CPU RAM efficient model loading? Only applicable for 🤗 Transformers models. [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + # Obsolete in FSDP2, ask for user input for FSDP1 + if fsdp_version == 1: + if fsdp_config["fsdp_cpu_ram_efficient_loading"]: + fsdp_config["fsdp_sync_module_states"] = True + else: + fsdp_config["fsdp_sync_module_states"] = _ask_field( + "Do you want each individually wrapped FSDP unit to broadcast module parameters from rank 0 at the start? [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + fsdp_config["fsdp_activation_checkpointing"] = _ask_field( + "Do you want to enable FSDP activation checkpointing? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + megatron_lm_config = {} + if distributed_type in [DistributedType.MULTI_GPU]: + use_megatron_lm = _ask_field( + "Do you want to use Megatron-LM ? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_megatron_lm: + distributed_type = DistributedType.MEGATRON_LM + if distributed_type == DistributedType.MEGATRON_LM: + prefix = "megatron_lm_" + megatron_lm_config[prefix + "tp_degree"] = _ask_field( + "What is the Tensor Parallelism degree/size? [1]:", + int, + default=1, + error_message="Please enter an integer.", + ) + if megatron_lm_config[prefix + "tp_degree"] > 1: + megatron_lm_config[prefix + "sequence_parallelism"] = _ask_field( + "Do you want to enable Sequence Parallelism? [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + + megatron_lm_config[prefix + "pp_degree"] = _ask_field( + "What is the Pipeline Parallelism degree/size? [1]:", + int, + default=1, + error_message="Please enter an integer.", + ) + if megatron_lm_config[prefix + "pp_degree"] > 1: + megatron_lm_config[prefix + "num_micro_batches"] = _ask_field( + "What is the number of micro-batches? [1]:", + int, + default=1, + error_message="Please enter an integer.", + ) + + megatron_lm_config[prefix + "recompute_activations"] = _ask_field( + "Do you want to enable selective activation recomputation? [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + + megatron_lm_config[prefix + "use_distributed_optimizer"] = _ask_field( + "Do you want to use distributed optimizer " + "which shards optimizer state and gradients across data parallel ranks? [YES/no]: ", + _convert_yes_no_to_bool, + default=True, + error_message="Please enter yes or no.", + ) + + megatron_lm_config[prefix + "gradient_clipping"] = _ask_field( + "What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: ", + float, + default=1.0, + ) + # TPU specific defaults + tpu_commands = None + tpu_command_file = None + tpu_downcast_bf16 = "no" + tpu_env = [] + tpu_name = None + tpu_vm = None + tpu_zone = None + tpu_use_sudo = False + tpu_use_cluster = False + + if distributed_type in [ + DistributedType.MULTI_CPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + DistributedType.MULTI_GPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + DistributedType.XLA, + ]: + machine_type = str(distributed_type).split(".")[1].replace("MULTI_", "") + if machine_type == "TPU": + machine_type += " cores" + elif machine_type == "CPU": + machine_type = "processes" + else: + machine_type += "(s)" + num_processes = _ask_field( + f"How many {machine_type} should be used for distributed training? [1]:", + int, + default=1, + error_message="Please enter an integer.", + ) + elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]: + num_processes = _ask_field( + "How many GPU(s) should be used for distributed training? [1]:", + int, + default=1, + error_message="Please enter an integer.", + ) + else: + num_processes = 1 + + if (distributed_type == DistributedType.MULTI_GPU) and (num_machines == 1) and (num_processes == 1): + raise ValueError( + f"Specified distributed type {distributed_type} but only using 1 GPU on a single machine. Please select `No distributed training` for the type of machine you are using." + ) + + if ( + distributed_type + in [ + DistributedType.MULTI_GPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + DistributedType.NO, + ] + and not use_cpu + and not use_mps + ): + if is_npu_available(): + machine_type = "NPU(s)" + elif is_mlu_available(): + machine_type = "MLU(s)" + elif is_sdaa_available(): + machine_type = "SDAA(s)" + elif is_musa_available(): + machine_type = "MUSA(s)" + elif is_xpu_available(): + machine_type = "XPU(s)" + elif is_hpu_available(): + machine_type = "HPU(s)" + else: + machine_type = "GPU(s)" + gpu_ids = _ask_field( + f"What {machine_type} (by id) should be used for training on this machine as a comma-separated list? [all]:", + default="all", + ) + + # CPU affinity is only supported on NVIDIA hardware for now + enable_cpu_affinity = False + if distributed_type in (DistributedType.NO, DistributedType.MULTI_GPU) and not use_cpu and not use_mps: + enable_cpu_affinity = _ask_field( + "Would you like to enable numa efficiency? (Currently only supported on NVIDIA hardware). [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + fp8_config = None + if distributed_type == DistributedType.XLA: + mixed_precision = "no" + main_training_function = _ask_field( + "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ", + default="main", + ) + tpu_use_cluster = _ask_field( + "Are you using a TPU cluster? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if tpu_use_cluster: + tpu_name = _ask_field( + "What is the name of your TPU cluster? ", + default=None, + error_message="Please enter the name of your TPU cluster.", + ) + tpu_zone = _ask_field( + "What is the zone of your TPU cluster? ", + default=None, + error_message="Please enter the zone of your TPU cluster.", + ) + tpu_use_sudo = _ask_field( + "To run a python script in a TPU pod, should `sudo` be used? [yes/NO]: ", + default=False, + error_message="Please enter yes or no.", + ) + run_commands = _ask_field( + "Do you have code you wish to run on startup in each pod? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if run_commands: + use_command_file = _ask_field( + "Is this code located in a bash script? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_command_file: + tpu_command_file = _ask_field( + "What is the path to your bash script? ", + default=None, + error_message="Please enter the path to your bash script.", + ) + tpu_command_file = os.path.abspath(tpu_command_file) + else: + print("Please enter each command separately you wish to run on startup in each pod.") + tpu_commands = [] + another_command = True + while another_command: + tpu_commands.append( + _ask_field( + "Please enter a single command to be ran ", + default=None, + error_message="Please enter the commands you wish to run on startup in each pod as a single string.", + ) + ) + another_command = _ask_field( + "Do you wish to add another command? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + tpu_vm = _ask_field( + "If not using an instance group, what are the names of the Compute VM instances to be used, separated by a comma: ", + default="", + ).split(",") + tpu_env = _ask_field( + "What environment variables do you wish to set in each pod, separated by a comma: ", + default="", + ).split(",") + + else: + main_training_function = "main" + if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config: + mixed_precision = None + else: + mixed_precision = _ask_options( + "Do you wish to use mixed precision?", + ["no", "fp16", "bf16", "fp8"], + _convert_mixed_precision, + ) + if mixed_precision == "fp8": + if not is_fp8_available(): + raise ValueError("FP8 (either Transformer Engine or MSAMP) is not installed on this machine.") + fp8_config = {} + fp8_config["backend"] = _ask_options( + "Which FP8 backend do you want to use?", + ["te", "msamp"], + _convert_fp8_backend, + ) + if fp8_config["backend"] == "TE": + if not is_transformer_engine_available(): + raise ValueError("TransformersEngine was selected, but it is not installed on this machine.") + fp8_config["use_autocast_during_eval"] = _ask_field( + "Do you want to use FP8 autocast during eval mode? Generally better metrics are found when this is disabled [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + ) + fp8_config["margin"] = _ask_field( + "What margin should be used for gradient scaling? [0]: ", + int, + default=0, + ) + fp8_config["interval"] = _ask_field( + "What interval should be used for for how often the scaling factor is recomputed? [1]: ", + int, + default=1, + ) + fp8_config["fp8_format"] = _ask_options( + "Which weight format should be used?", + ["HYBRID", "E4M3", "E5M2"], + lambda i: ["HYBRID", "E4M3", "E5M2"][i], + default=0, + ) + fp8_config["amax_history_length"] = _ask_field( + "What length of history should be used for the amax scaling factor computation? [1024]: ", + int, + default=1024, + ) + fp8_config["amax_compute_algorithm"] = _ask_options( + "Which algorithm should be used for the amax scaling factor computation?", + ["max", "most_recent"], + lambda x: "max" if x == 0 else "most_recent", + default=0, + ) + fp8_config["override_linear_precision"] = _ask_field( + "Do you want to to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + ) + if fp8_config["override_linear_precision"]: + fprop = _ask_field( + "Should `fprop` be executed in higher precision? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + ) + dgrad = _ask_field( + "Should `dgrad` be executed in higher precision? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + ) + wgrad = _ask_field( + "Should `wgrad` be executed in higher precision? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + ) + fp8_config["override_linear_precision"] = (fprop, dgrad, wgrad) + else: + fp8_config["override_linear_precision"] = (False, False, False) + + elif fp8_config["backend"] == "MSAMP": + if not is_msamp_available(): + raise ValueError("MSAMP was selected, but it is not installed on this machine.") + fp8_config["optimization_level"] = _ask_options( + "Which optimization level should be used?", + ["O1", "O2"], + lambda x: "O1" if x == 0 else "O2", + default=1, + ) + + if use_dynamo and mixed_precision == "no" and not use_cpu: + print( + "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts." + ) + + if distributed_type == DistributedType.XLA and mixed_precision == "bf16": + tpu_downcast_bf16 = _ask_field( + "Should `torch.float` be cast as `bfloat16` and `torch.double` remain `float32` on TPUs?", default="no" + ) + + return ClusterConfig( + compute_environment=ComputeEnvironment.LOCAL_MACHINE, + distributed_type=distributed_type, + num_processes=num_processes, + gpu_ids=gpu_ids, + mixed_precision=mixed_precision, + downcast_bf16=tpu_downcast_bf16, + machine_rank=machine_rank, + num_machines=num_machines, + main_process_ip=main_process_ip, + main_process_port=main_process_port, + main_training_function=main_training_function, + fp8_config=fp8_config, + deepspeed_config=deepspeed_config, + fsdp_config=fsdp_config, + megatron_lm_config=megatron_lm_config, + ipex_config=ipex_config, + mpirun_config=mpirun_config, + use_cpu=use_cpu, + rdzv_backend=rdzv_backend, + same_network=same_network, + commands=tpu_commands, + command_file=tpu_command_file, + tpu_env=tpu_env, + tpu_name=tpu_name, + tpu_vm=tpu_vm, + tpu_zone=tpu_zone, + tpu_use_sudo=tpu_use_sudo, + tpu_use_cluster=tpu_use_cluster, + dynamo_config=dynamo_config, + debug=debug, + enable_cpu_affinity=enable_cpu_affinity, + ) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/config.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/config.py new file mode 100644 index 0000000000000000000000000000000000000000..72414f2abe62d76bd5133f4b0ed99bf34133f6f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/config.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import argparse +import os + +from accelerate.utils import ComputeEnvironment + +from .cluster import get_cluster_input +from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file # noqa: F401 +from .config_utils import _ask_field, _ask_options, _convert_compute_environment # noqa: F401 +from .sagemaker import get_sagemaker_input + + +description = "Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine" + + +def get_user_input(): + compute_environment = _ask_options( + "In which compute environment are you running?", + ["This machine", "AWS (Amazon SageMaker)"], + _convert_compute_environment, + ) + if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER: + config = get_sagemaker_input() + else: + config = get_cluster_input() + return config + + +def config_command_parser(subparsers=None): + if subparsers is not None: + parser = subparsers.add_parser("config", description=description) + else: + parser = argparse.ArgumentParser("Accelerate config command", description=description) + + parser.add_argument( + "--config_file", + default=None, + help=( + "The path to use to store the config file. Will default to a file named default_config.yaml in the cache " + "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " + "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " + "with 'huggingface'." + ), + ) + + if subparsers is not None: + parser.set_defaults(func=config_command) + return parser + + +def config_command(args): + config = get_user_input() + if args.config_file is not None: + config_file = args.config_file + else: + if not os.path.isdir(cache_dir): + os.makedirs(cache_dir) + config_file = default_yaml_config_file + + if config_file.endswith(".json"): + config.to_json_file(config_file) + else: + config.to_yaml_file(config_file) + print(f"accelerate configuration saved at {config_file}") + + +def main(): + parser = config_command_parser() + args = parser.parse_args() + config_command(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/config_args.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/config_args.py new file mode 100644 index 0000000000000000000000000000000000000000..3907329287b2f3ac033f4f74f84bc877b91af4f9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/config_args.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import json +import os +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Union + +import yaml + +from ...utils import ComputeEnvironment, DistributedType, SageMakerDistributedType +from ...utils.constants import SAGEMAKER_PYTHON_VERSION, SAGEMAKER_PYTORCH_VERSION, SAGEMAKER_TRANSFORMERS_VERSION + + +hf_cache_home = os.path.expanduser( + os.environ.get("HF_HOME", os.path.join(os.environ.get("XDG_CACHE_HOME", "~/.cache"), "huggingface")) +) +cache_dir = os.path.join(hf_cache_home, "accelerate") +default_json_config_file = os.path.join(cache_dir, "default_config.yaml") +default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml") + +# For backward compatibility: the default config is the json one if it's the only existing file. +if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file): + default_config_file = default_yaml_config_file +else: + default_config_file = default_json_config_file + + +def load_config_from_file(config_file): + if config_file is not None: + if not os.path.isfile(config_file): + raise FileNotFoundError( + f"The passed configuration file `{config_file}` does not exist. " + "Please pass an existing file to `accelerate launch`, or use the default one " + "created through `accelerate config` and run `accelerate launch` " + "without the `--config_file` argument." + ) + else: + config_file = default_config_file + with open(config_file, encoding="utf-8") as f: + if config_file.endswith(".json"): + if ( + json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE) + == ComputeEnvironment.LOCAL_MACHINE + ): + config_class = ClusterConfig + else: + config_class = SageMakerConfig + return config_class.from_json_file(json_file=config_file) + else: + if ( + yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE) + == ComputeEnvironment.LOCAL_MACHINE + ): + config_class = ClusterConfig + else: + config_class = SageMakerConfig + return config_class.from_yaml_file(yaml_file=config_file) + + +@dataclass +class BaseConfig: + compute_environment: ComputeEnvironment + distributed_type: Union[DistributedType, SageMakerDistributedType] + mixed_precision: str + use_cpu: bool + debug: bool + + def to_dict(self): + result = self.__dict__ + # For serialization, it's best to convert Enums to strings (or their underlying value type). + + def _convert_enums(value): + if isinstance(value, Enum): + return value.value + if isinstance(value, dict): + if not bool(value): + return None + for key1, value1 in value.items(): + value[key1] = _convert_enums(value1) + return value + + for key, value in result.items(): + result[key] = _convert_enums(value) + result = {k: v for k, v in result.items() if v is not None} + return result + + @staticmethod + def process_config(config_dict): + """ + Processes `config_dict` and sets default values for any missing keys + """ + if "compute_environment" not in config_dict: + config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE + if "distributed_type" not in config_dict: + raise ValueError("A `distributed_type` must be specified in the config file.") + if "num_processes" not in config_dict and config_dict["distributed_type"] == DistributedType.NO: + config_dict["num_processes"] = 1 + if "mixed_precision" not in config_dict: + config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else None + if "fp16" in config_dict: # Convert the config to the new format. + del config_dict["fp16"] + if "dynamo_backend" in config_dict: # Convert the config to the new format. + dynamo_backend = config_dict.pop("dynamo_backend") + config_dict["dynamo_config"] = {} if dynamo_backend == "NO" else {"dynamo_backend": dynamo_backend} + if "use_cpu" not in config_dict: + config_dict["use_cpu"] = False + if "debug" not in config_dict: + config_dict["debug"] = False + if "enable_cpu_affinity" not in config_dict: + config_dict["enable_cpu_affinity"] = False + return config_dict + + @classmethod + def from_json_file(cls, json_file=None): + json_file = default_json_config_file if json_file is None else json_file + with open(json_file, encoding="utf-8") as f: + config_dict = json.load(f) + config_dict = cls.process_config(config_dict) + extra_keys = sorted(set(config_dict.keys()) - set(cls.__dataclass_fields__.keys())) + if len(extra_keys) > 0: + raise ValueError( + f"The config file at {json_file} had unknown keys ({extra_keys}), please try upgrading your `accelerate`" + " version or fix (and potentially remove) these keys from your config file." + ) + + return cls(**config_dict) + + def to_json_file(self, json_file): + with open(json_file, "w", encoding="utf-8") as f: + content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n" + f.write(content) + + @classmethod + def from_yaml_file(cls, yaml_file=None): + yaml_file = default_yaml_config_file if yaml_file is None else yaml_file + with open(yaml_file, encoding="utf-8") as f: + config_dict = yaml.safe_load(f) + config_dict = cls.process_config(config_dict) + extra_keys = sorted(set(config_dict.keys()) - set(cls.__dataclass_fields__.keys())) + if len(extra_keys) > 0: + raise ValueError( + f"The config file at {yaml_file} had unknown keys ({extra_keys}), please try upgrading your `accelerate`" + " version or fix (and potentially remove) these keys from your config file." + ) + return cls(**config_dict) + + def to_yaml_file(self, yaml_file): + with open(yaml_file, "w", encoding="utf-8") as f: + yaml.safe_dump(self.to_dict(), f) + + def __post_init__(self): + if isinstance(self.compute_environment, str): + self.compute_environment = ComputeEnvironment(self.compute_environment) + if isinstance(self.distributed_type, str): + if self.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER: + self.distributed_type = SageMakerDistributedType(self.distributed_type) + else: + self.distributed_type = DistributedType(self.distributed_type) + if getattr(self, "dynamo_config", None) is None: + self.dynamo_config = {} + + +@dataclass +class ClusterConfig(BaseConfig): + num_processes: int = -1 # For instance if we use SLURM and the user manually passes it in + machine_rank: int = 0 + num_machines: int = 1 + gpu_ids: Optional[str] = None + main_process_ip: Optional[str] = None + main_process_port: Optional[int] = None + rdzv_backend: Optional[str] = "static" + same_network: Optional[bool] = False + main_training_function: str = "main" + enable_cpu_affinity: bool = False + + # args for FP8 training + fp8_config: dict = None + # args for deepspeed_plugin + deepspeed_config: dict = None + # args for fsdp + fsdp_config: dict = None + # args for megatron_lm + megatron_lm_config: dict = None + # args for ipex + ipex_config: dict = None + # args for mpirun + mpirun_config: dict = None + # args for TPU + downcast_bf16: bool = False + + # args for TPU pods + tpu_name: str = None + tpu_zone: str = None + tpu_use_cluster: bool = False + tpu_use_sudo: bool = False + command_file: str = None + commands: list[str] = None + tpu_vm: list[str] = None + tpu_env: list[str] = None + + # args for dynamo + dynamo_config: dict = None + + def __post_init__(self): + if self.deepspeed_config is None: + self.deepspeed_config = {} + if self.fsdp_config is None: + self.fsdp_config = {} + if self.megatron_lm_config is None: + self.megatron_lm_config = {} + if self.ipex_config is None: + self.ipex_config = {} + if self.mpirun_config is None: + self.mpirun_config = {} + if self.fp8_config is None: + self.fp8_config = {} + return super().__post_init__() + + +@dataclass +class SageMakerConfig(BaseConfig): + ec2_instance_type: str + iam_role_name: str + image_uri: Optional[str] = None + profile: Optional[str] = None + region: str = "us-east-1" + num_machines: int = 1 + gpu_ids: str = "all" + base_job_name: str = f"accelerate-sagemaker-{num_machines}" + pytorch_version: str = SAGEMAKER_PYTORCH_VERSION + transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION + py_version: str = SAGEMAKER_PYTHON_VERSION + sagemaker_inputs_file: str = None + sagemaker_metrics_file: str = None + additional_args: dict = None + dynamo_config: dict = None + enable_cpu_affinity: bool = False diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/config_utils.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/config_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..24ee971383c8fdda1491e2b5795446790755ac70 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/config_utils.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import argparse + +from ...utils.dataclasses import ( + ComputeEnvironment, + DistributedType, + DynamoBackend, + FP8BackendType, + PrecisionType, + SageMakerDistributedType, +) +from ..menu import BulletMenu + + +DYNAMO_BACKENDS = [ + "EAGER", + "AOT_EAGER", + "INDUCTOR", + "AOT_TS_NVFUSER", + "NVPRIMS_NVFUSER", + "CUDAGRAPHS", + "OFI", + "FX2TRT", + "ONNXRT", + "TENSORRT", + "AOT_TORCHXLA_TRACE_ONCE", + "TORHCHXLA_TRACE_ONCE", + "IPEX", + "TVM", +] + + +def _ask_field(input_text, convert_value=None, default=None, error_message=None): + ask_again = True + while ask_again: + result = input(input_text) + try: + if default is not None and len(result) == 0: + return default + return convert_value(result) if convert_value is not None else result + except Exception: + if error_message is not None: + print(error_message) + + +def _ask_options(input_text, options=[], convert_value=None, default=0): + menu = BulletMenu(input_text, options) + result = menu.run(default_choice=default) + return convert_value(result) if convert_value is not None else result + + +def _convert_compute_environment(value): + value = int(value) + return ComputeEnvironment(["LOCAL_MACHINE", "AMAZON_SAGEMAKER"][value]) + + +def _convert_distributed_mode(value): + value = int(value) + return DistributedType( + [ + "NO", + "MULTI_CPU", + "MULTI_XPU", + "MULTI_HPU", + "MULTI_GPU", + "MULTI_NPU", + "MULTI_MLU", + "MULTI_SDAA", + "MULTI_MUSA", + "XLA", + ][value] + ) + + +def _convert_dynamo_backend(value): + value = int(value) + return DynamoBackend(DYNAMO_BACKENDS[value]).value + + +def _convert_mixed_precision(value): + value = int(value) + return PrecisionType(["no", "fp16", "bf16", "fp8"][value]) + + +def _convert_sagemaker_distributed_mode(value): + value = int(value) + return SageMakerDistributedType(["NO", "DATA_PARALLEL", "MODEL_PARALLEL"][value]) + + +def _convert_fp8_backend(value): + value = int(value) + return FP8BackendType(["TE", "MSAMP"][value]) + + +def _convert_yes_no_to_bool(value): + return {"yes": True, "no": False}[value.lower()] + + +class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter): + """ + A custom formatter that will remove the usage line from the help message for subcommands. + """ + + def _format_usage(self, usage, actions, groups, prefix): + usage = super()._format_usage(usage, actions, groups, prefix) + usage = usage.replace(" [] ", "") + return usage diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/default.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/default.py new file mode 100644 index 0000000000000000000000000000000000000000..f5f267f8c47cccc80ce9ef53f970c2266472c117 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/default.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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 pathlib import Path + +import torch + +from ...utils import ( + is_hpu_available, + is_mlu_available, + is_musa_available, + is_npu_available, + is_sdaa_available, + is_xpu_available, +) +from .config_args import ClusterConfig, default_json_config_file +from .config_utils import SubcommandHelpFormatter + + +description = "Create a default config file for Accelerate with only a few flags set." + + +def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file): + """ + Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also + set CPU if it is a CPU-only machine. + + Args: + mixed_precision (`str`, *optional*, defaults to "no"): + Mixed Precision to use. Should be one of "no", "fp16", or "bf16" + save_location (`str`, *optional*, defaults to `default_json_config_file`): + Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default + location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overridden by setting + the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`. + """ + path = Path(save_location) + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + print( + f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`." + ) + return False + mixed_precision = mixed_precision.lower() + if mixed_precision not in ["no", "fp16", "bf16", "fp8"]: + raise ValueError( + f"`mixed_precision` should be one of 'no', 'fp16', 'bf16', or 'fp8'. Received {mixed_precision}" + ) + config = { + "compute_environment": "LOCAL_MACHINE", + "mixed_precision": mixed_precision, + } + if is_mlu_available(): + num_mlus = torch.mlu.device_count() + config["num_processes"] = num_mlus + config["use_cpu"] = False + if num_mlus > 1: + config["distributed_type"] = "MULTI_MLU" + else: + config["distributed_type"] = "NO" + if is_sdaa_available(): + num_sdaas = torch.sdaa.device_count() + config["num_processes"] = num_sdaas + config["use_cpu"] = False + if num_sdaas > 1: + config["distributed_type"] = "MULTI_SDAA" + else: + config["distributed_type"] = "NO" + elif is_musa_available(): + num_musas = torch.musa.device_count() + config["num_processes"] = num_musas + config["use_cpu"] = False + if num_musas > 1: + config["distributed_type"] = "MULTI_MUSA" + else: + config["distributed_type"] = "NO" + elif is_hpu_available(): + num_hpus = torch.hpu.device_count() + config["num_processes"] = num_hpus + config["use_cpu"] = False + if num_hpus > 1: + config["distributed_type"] = "MULTI_HPU" + else: + config["distributed_type"] = "NO" + elif torch.cuda.is_available(): + num_gpus = torch.cuda.device_count() + config["num_processes"] = num_gpus + config["use_cpu"] = False + if num_gpus > 1: + config["distributed_type"] = "MULTI_GPU" + else: + config["distributed_type"] = "NO" + elif is_xpu_available(): + num_xpus = torch.xpu.device_count() + config["num_processes"] = num_xpus + config["use_cpu"] = False + if num_xpus > 1: + config["distributed_type"] = "MULTI_XPU" + else: + config["distributed_type"] = "NO" + elif is_npu_available(): + num_npus = torch.npu.device_count() + config["num_processes"] = num_npus + config["use_cpu"] = False + if num_npus > 1: + config["distributed_type"] = "MULTI_NPU" + else: + config["distributed_type"] = "NO" + else: + num_xpus = 0 + config["use_cpu"] = True + config["num_processes"] = 1 + config["distributed_type"] = "NO" + config["debug"] = False + config["enable_cpu_affinity"] = False + config = ClusterConfig(**config) + config.to_json_file(path) + return path + + +def default_command_parser(parser, parents): + parser = parser.add_parser("default", parents=parents, help=description, formatter_class=SubcommandHelpFormatter) + parser.add_argument( + "--config_file", + default=default_json_config_file, + help=( + "The path to use to store the config file. Will default to a file named default_config.yaml in the cache " + "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " + "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " + "with 'huggingface'." + ), + dest="save_location", + ) + + parser.add_argument( + "--mixed_precision", + choices=["no", "fp16", "bf16"], + type=str, + help="Whether or not to use mixed precision training. " + "Choose between FP16 and BF16 (bfloat16) training. " + "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.", + default="no", + ) + parser.set_defaults(func=default_config_command) + return parser + + +def default_config_command(args): + config_file = write_basic_config(args.mixed_precision, args.save_location) + if config_file: + print(f"accelerate configuration saved at {config_file}") diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/sagemaker.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/sagemaker.py new file mode 100644 index 0000000000000000000000000000000000000000..5092ef31fc4715f901be6c1e7bfe80c0b140d767 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/sagemaker.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. +import json +import os + +from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES +from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType +from ...utils.imports import is_boto3_available +from .config_args import SageMakerConfig +from .config_utils import ( + DYNAMO_BACKENDS, + _ask_field, + _ask_options, + _convert_dynamo_backend, + _convert_mixed_precision, + _convert_sagemaker_distributed_mode, + _convert_yes_no_to_bool, +) + + +if is_boto3_available(): + import boto3 # noqa: F401 + + +def _create_iam_role_for_sagemaker(role_name): + iam_client = boto3.client("iam") + + sagemaker_trust_policy = { + "Version": "2012-10-17", + "Statement": [ + {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"} + ], + } + try: + # create the role, associated with the chosen trust policy + iam_client.create_role( + RoleName=role_name, AssumeRolePolicyDocument=json.dumps(sagemaker_trust_policy, indent=2) + ) + policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "sagemaker:*", + "ecr:GetDownloadUrlForLayer", + "ecr:BatchGetImage", + "ecr:BatchCheckLayerAvailability", + "ecr:GetAuthorizationToken", + "cloudwatch:PutMetricData", + "cloudwatch:GetMetricData", + "cloudwatch:GetMetricStatistics", + "cloudwatch:ListMetrics", + "logs:CreateLogGroup", + "logs:CreateLogStream", + "logs:DescribeLogStreams", + "logs:PutLogEvents", + "logs:GetLogEvents", + "s3:CreateBucket", + "s3:ListBucket", + "s3:GetBucketLocation", + "s3:GetObject", + "s3:PutObject", + ], + "Resource": "*", + } + ], + } + # attach policy to role + iam_client.put_role_policy( + RoleName=role_name, + PolicyName=f"{role_name}_policy_permission", + PolicyDocument=json.dumps(policy_document, indent=2), + ) + except iam_client.exceptions.EntityAlreadyExistsException: + print(f"role {role_name} already exists. Using existing one") + + +def _get_iam_role_arn(role_name): + iam_client = boto3.client("iam") + return iam_client.get_role(RoleName=role_name)["Role"]["Arn"] + + +def get_sagemaker_input(): + credentials_configuration = _ask_options( + "How do you want to authorize?", + ["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "], + int, + ) + aws_profile = None + if credentials_configuration == 0: + aws_profile = _ask_field("Enter your AWS Profile name: [default] ", default="default") + os.environ["AWS_PROFILE"] = aws_profile + else: + print( + "Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with," + "`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`" + ) + aws_access_key_id = _ask_field("AWS Access Key ID: ") + os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id + + aws_secret_access_key = _ask_field("AWS Secret Access Key: ") + os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key + + aws_region = _ask_field("Enter your AWS Region: [us-east-1]", default="us-east-1") + os.environ["AWS_DEFAULT_REGION"] = aws_region + + role_management = _ask_options( + "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?", + ["Provide IAM Role name", "Create new IAM role using credentials"], + int, + ) + if role_management == 0: + iam_role_name = _ask_field("Enter your IAM role name: ") + else: + iam_role_name = "accelerate_sagemaker_execution_role" + print(f'Accelerate will create an iam role "{iam_role_name}" using the provided credentials') + _create_iam_role_for_sagemaker(iam_role_name) + + is_custom_docker_image = _ask_field( + "Do you want to use custom Docker image? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + docker_image = None + if is_custom_docker_image: + docker_image = _ask_field("Enter your Docker image: ", lambda x: str(x).lower()) + + is_sagemaker_inputs_enabled = _ask_field( + "Do you want to provide SageMaker input channels with data locations? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + sagemaker_inputs_file = None + if is_sagemaker_inputs_enabled: + sagemaker_inputs_file = _ask_field( + "Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): ", + lambda x: str(x).lower(), + ) + + is_sagemaker_metrics_enabled = _ask_field( + "Do you want to enable SageMaker metrics? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + sagemaker_metrics_file = None + if is_sagemaker_metrics_enabled: + sagemaker_metrics_file = _ask_field( + "Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): ", + lambda x: str(x).lower(), + ) + + distributed_type = _ask_options( + "What is the distributed mode?", + ["No distributed training", "Data parallelism"], + _convert_sagemaker_distributed_mode, + ) + dynamo_config = {} + use_dynamo = _ask_field( + "Do you wish to optimize your script with torch dynamo?[yes/NO]:", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + if use_dynamo: + prefix = "dynamo_" + dynamo_config[prefix + "backend"] = _ask_options( + "Which dynamo backend would you like to use?", + [x.lower() for x in DYNAMO_BACKENDS], + _convert_dynamo_backend, + default=2, + ) + use_custom_options = _ask_field( + "Do you want to customize the defaults sent to torch.compile? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + if use_custom_options: + dynamo_config[prefix + "mode"] = _ask_options( + "Which mode do you want to use?", + TORCH_DYNAMO_MODES, + lambda x: TORCH_DYNAMO_MODES[int(x)], + default="default", + ) + dynamo_config[prefix + "use_fullgraph"] = _ask_field( + "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + dynamo_config[prefix + "use_dynamic"] = _ask_field( + "Do you want to enable dynamic shape tracing? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + dynamo_config[prefix + "use_regional_compilation"] = _ask_field( + "Do you want to enable regional compilation? [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + ec2_instance_query = "Which EC2 instance type you want to use for your training?" + if distributed_type != SageMakerDistributedType.NO: + ec2_instance_type = _ask_options( + ec2_instance_query, SAGEMAKER_PARALLEL_EC2_INSTANCES, lambda x: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(x)] + ) + else: + ec2_instance_query += "? [ml.p3.2xlarge]:" + ec2_instance_type = _ask_field(ec2_instance_query, lambda x: str(x).lower(), default="ml.p3.2xlarge") + + debug = False + if distributed_type != SageMakerDistributedType.NO: + debug = _ask_field( + "Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: ", + _convert_yes_no_to_bool, + default=False, + error_message="Please enter yes or no.", + ) + + num_machines = 1 + if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL): + num_machines = _ask_field( + "How many machines do you want use? [1]: ", + int, + default=1, + ) + + mixed_precision = _ask_options( + "Do you wish to use FP16 or BF16 (mixed precision)?", + ["no", "fp16", "bf16", "fp8"], + _convert_mixed_precision, + ) + + if use_dynamo and mixed_precision == "no": + print( + "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts." + ) + + return SageMakerConfig( + image_uri=docker_image, + compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER, + distributed_type=distributed_type, + use_cpu=False, + dynamo_config=dynamo_config, + ec2_instance_type=ec2_instance_type, + profile=aws_profile, + region=aws_region, + iam_role_name=iam_role_name, + mixed_precision=mixed_precision, + num_machines=num_machines, + sagemaker_inputs_file=sagemaker_inputs_file, + sagemaker_metrics_file=sagemaker_metrics_file, + debug=debug, + ) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/update.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/update.py new file mode 100644 index 0000000000000000000000000000000000000000..5f025594b04ada3e3a78687befc5c1bc1d236adf --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/config/update.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python + +# Copyright 2022 The HuggingFace Team. 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 pathlib import Path + +from .config_args import default_config_file, load_config_from_file +from .config_utils import SubcommandHelpFormatter + + +description = "Update an existing config file with the latest defaults while maintaining the old configuration." + + +def update_config(args): + """ + Update an existing config file with the latest defaults while maintaining the old configuration. + """ + config_file = args.config_file + if config_file is None and Path(default_config_file).exists(): + config_file = default_config_file + elif not Path(config_file).exists(): + raise ValueError(f"The passed config file located at {config_file} doesn't exist.") + config = load_config_from_file(config_file) + + if config_file.endswith(".json"): + config.to_json_file(config_file) + else: + config.to_yaml_file(config_file) + return config_file + + +def update_command_parser(parser, parents): + parser = parser.add_parser("update", parents=parents, help=description, formatter_class=SubcommandHelpFormatter) + parser.add_argument( + "--config_file", + default=None, + help=( + "The path to the config file to update. Will default to a file named default_config.yaml in the cache " + "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " + "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " + "with 'huggingface'." + ), + ) + + parser.set_defaults(func=update_config_command) + return parser + + +def update_config_command(args): + config_file = update_config(args) + print(f"Sucessfully updated the configuration file at {config_file}.") diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/env.py b/venv/lib/python3.10/site-packages/accelerate/commands/env.py new file mode 100644 index 0000000000000000000000000000000000000000..3dd2170aea8ba48a08e31c8031c0447825ec4797 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/env.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python + +# Copyright 2022 The HuggingFace Team. 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. + +import argparse +import os +import platform +import subprocess + +import numpy as np +import psutil +import torch + +from accelerate import __version__ as version +from accelerate.commands.config import default_config_file, load_config_from_file + +from ..utils import is_mlu_available, is_musa_available, is_npu_available, is_sdaa_available, is_xpu_available + + +def env_command_parser(subparsers=None): + if subparsers is not None: + parser = subparsers.add_parser("env") + else: + parser = argparse.ArgumentParser("Accelerate env command") + + parser.add_argument( + "--config_file", default=None, help="The config file to use for the default values in the launching script." + ) + + if subparsers is not None: + parser.set_defaults(func=env_command) + return parser + + +def env_command(args): + pt_version = torch.__version__ + pt_cuda_available = torch.cuda.is_available() + pt_xpu_available = is_xpu_available() + pt_mlu_available = is_mlu_available() + pt_sdaa_available = is_sdaa_available() + pt_musa_available = is_musa_available() + pt_npu_available = is_npu_available() + + accelerator = "N/A" + if pt_cuda_available: + accelerator = "CUDA" + elif pt_xpu_available: + accelerator = "XPU" + elif pt_mlu_available: + accelerator = "MLU" + elif pt_sdaa_available: + accelerator = "SDAA" + elif pt_musa_available: + accelerator = "MUSA" + elif pt_npu_available: + accelerator = "NPU" + + accelerate_config = "Not found" + # Get the default from the config file. + if args.config_file is not None or os.path.isfile(default_config_file): + accelerate_config = load_config_from_file(args.config_file).to_dict() + + # if we can run which, get it + command = None + bash_location = "Not found" + if os.name == "nt": + command = ["where", "accelerate"] + elif os.name == "posix": + command = ["which", "accelerate"] + if command is not None: + bash_location = subprocess.check_output(command, text=True, stderr=subprocess.STDOUT).strip() + info = { + "`Accelerate` version": version, + "Platform": platform.platform(), + "`accelerate` bash location": bash_location, + "Python version": platform.python_version(), + "Numpy version": np.__version__, + "PyTorch version": f"{pt_version}", + "PyTorch accelerator": accelerator, + "System RAM": f"{psutil.virtual_memory().total / 1024**3:.2f} GB", + } + if pt_cuda_available: + info["GPU type"] = torch.cuda.get_device_name() + elif pt_xpu_available: + info["XPU type"] = torch.xpu.get_device_name() + elif pt_mlu_available: + info["MLU type"] = torch.mlu.get_device_name() + elif pt_sdaa_available: + info["SDAA type"] = torch.sdaa.get_device_name() + elif pt_musa_available: + info["MUSA type"] = torch.musa.get_device_name() + elif pt_npu_available: + info["CANN version"] = torch.version.cann + + print("\nCopy-and-paste the text below in your GitHub issue\n") + print("\n".join([f"- {prop}: {val}" for prop, val in info.items()])) + + print("- `Accelerate` default config:" if args.config_file is None else "- `Accelerate` config passed:") + accelerate_config_str = ( + "\n".join([f"\t- {prop}: {val}" for prop, val in accelerate_config.items()]) + if isinstance(accelerate_config, dict) + else f"\t{accelerate_config}" + ) + print(accelerate_config_str) + + info["`Accelerate` configs"] = accelerate_config + + return info + + +def main() -> int: + parser = env_command_parser() + args = parser.parse_args() + env_command(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/estimate.py b/venv/lib/python3.10/site-packages/accelerate/commands/estimate.py new file mode 100644 index 0000000000000000000000000000000000000000..77571777c6c8e78b7729e4b29e1a98cea6e95de7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/estimate.py @@ -0,0 +1,312 @@ +#!/usr/bin/env python + +# Copyright 2023 The HuggingFace Team. 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. +import torch +from huggingface_hub import model_info +from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError + +from accelerate import init_empty_weights +from accelerate.commands.utils import CustomArgumentParser +from accelerate.utils import ( + calculate_maximum_sizes, + convert_bytes, + is_timm_available, + is_transformers_available, +) + + +if is_transformers_available(): + import transformers + from transformers import AutoConfig, AutoModel + +if is_timm_available(): + import timm + + +def verify_on_hub(repo: str, token: str = None): + "Verifies that the model is on the hub and returns the model info." + try: + return model_info(repo, token=token) + except (OSError, GatedRepoError): + return "gated" + except RepositoryNotFoundError: + return "repo" + + +def check_has_model(error): + """ + Checks what library spawned `error` when a model is not found + """ + if is_timm_available() and isinstance(error, RuntimeError) and "Unknown model" in error.args[0]: + return "timm" + elif ( + is_transformers_available() + and isinstance(error, OSError) + and "does not appear to have a file named" in error.args[0] + ): + return "transformers" + else: + return "unknown" + + +def create_empty_model(model_name: str, library_name: str, trust_remote_code: bool = False, access_token: str = None): + """ + Creates an empty model in full precision from its parent library on the `Hub` to calculate the overall memory + consumption. + + Args: + model_name (`str`): + The model name on the Hub + library_name (`str`): + The library the model has an integration with, such as `transformers`. Will be used if `model_name` has no + metadata on the Hub to determine the library. + trust_remote_code (`bool`, `optional`, defaults to `False`): + Whether or not to allow for custom models defined on the Hub in their own modeling files. This option + should only be set to `True` for repositories you trust and in which you have read the code, as it will + execute code present on the Hub on your local machine. + access_token (`str`, `optional`, defaults to `None`): + The access token to use to access private or gated models on the Hub. (for use on the Gradio app) + + Returns: + `torch.nn.Module`: The torch model that has been initialized on the `meta` device. + + """ + model_info = verify_on_hub(model_name, access_token) + # Simplified errors + if model_info == "gated": + raise GatedRepoError( + f"Repo for model `{model_name}` is gated. You must be authenticated to access it. Please run `huggingface-cli login`." + ) + elif model_info == "repo": + raise RepositoryNotFoundError( + f"Repo for model `{model_name}` does not exist on the Hub. If you are trying to access a private repo," + " make sure you are authenticated via `huggingface-cli login` and have access." + ) + if library_name is None: + library_name = getattr(model_info, "library_name", False) + if not library_name: + raise ValueError( + f"Model `{model_name}` does not have any library metadata on the Hub, please manually pass in a `--library_name` to use (such as `transformers`)" + ) + if library_name == "transformers": + if not is_transformers_available(): + raise ImportError( + f"To check `{model_name}`, `transformers` must be installed. Please install it via `pip install transformers`" + ) + print(f"Loading pretrained config for `{model_name}` from `transformers`...") + if model_info.config is None: + raise RuntimeError(f"Tried to load `{model_name}` with `transformers` but it does not have any metadata.") + + auto_map = model_info.config.get("auto_map", False) + config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code, token=access_token) + with init_empty_weights(): + # remote code could specify a specific `AutoModel` class in the `auto_map` + constructor = AutoModel + if isinstance(auto_map, dict): + value = None + for key in auto_map.keys(): + if key.startswith("AutoModelFor"): + value = key + break + if value is not None: + constructor = getattr(transformers, value) + # we need to pass the dtype, otherwise it is going to use the torch_dtype that is saved in the config + model = constructor.from_config(config, torch_dtype=torch.float32, trust_remote_code=trust_remote_code) + elif library_name == "timm": + if not is_timm_available(): + raise ImportError( + f"To check `{model_name}`, `timm` must be installed. Please install it via `pip install timm`" + ) + print(f"Loading pretrained config for `{model_name}` from `timm`...") + with init_empty_weights(): + model = timm.create_model(model_name, pretrained=False) + else: + raise ValueError( + f"Library `{library_name}` is not supported yet, please open an issue on GitHub for us to add support." + ) + return model + + +def create_ascii_table(headers: list, rows: list, title: str): + "Creates a pretty table from a list of rows, minimal version of `tabulate`." + sep_char, in_between = "│", "─" + column_widths = [] + for i in range(len(headers)): + column_values = [row[i] for row in rows] + [headers[i]] + max_column_width = max(len(value) for value in column_values) + column_widths.append(max_column_width) + + formats = [f"%{column_widths[i]}s" for i in range(len(rows[0]))] + + pattern = f"{sep_char}{sep_char.join(formats)}{sep_char}" + diff = 0 + + def make_row(left_char, middle_char, right_char): + return f"{left_char}{middle_char.join([in_between * n for n in column_widths])}{in_between * diff}{right_char}" + + separator = make_row("├", "┼", "┤") + if len(title) > sum(column_widths): + diff = abs(len(title) - len(separator)) + column_widths[-1] += diff + + # Update with diff + separator = make_row("├", "┼", "┤") + initial_rows = [ + make_row("┌", in_between, "┐"), + f"{sep_char}{title.center(len(separator) - 2)}{sep_char}", + make_row("├", "┬", "┤"), + ] + table = "\n".join(initial_rows) + "\n" + column_widths[-1] += diff + centered_line = [text.center(column_widths[i]) for i, text in enumerate(headers)] + table += f"{pattern % tuple(centered_line)}\n{separator}\n" + for i, line in enumerate(rows): + centered_line = [t.center(column_widths[i]) for i, t in enumerate(line)] + table += f"{pattern % tuple(centered_line)}\n" + table += f"└{'┴'.join([in_between * n for n in column_widths])}┘" + + return table + + +def estimate_command_parser(subparsers=None): + if subparsers is not None: + parser = subparsers.add_parser("estimate-memory") + else: + parser = CustomArgumentParser(description="Model size estimator for fitting a model onto CUDA memory.") + + parser.add_argument("model_name", type=str, help="The model name on the Hugging Face Hub.") + parser.add_argument( + "--library_name", + type=str, + help="The library the model has an integration with, such as `transformers`, needed only if this information is not stored on the Hub.", + choices=["timm", "transformers"], + ) + parser.add_argument( + "--dtypes", + type=str, + nargs="+", + default=["float32", "float16", "int8", "int4"], + help="The dtypes to use for the model, must be one (or many) of `float32`, `float16`, `int8`, and `int4`", + choices=["float32", "float16", "int8", "int4"], + ) + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="""Whether or not to allow for custom models defined on the Hub in their own modeling files. This flag + should only be used for repositories you trust and in which you have read the code, as it will execute + code present on the Hub on your local machine.""", + default=False, + ) + + if subparsers is not None: + parser.set_defaults(func=estimate_command) + return parser + + +def estimate_training_usage(bytes: int, mixed_precision: str, msamp_config: str = None) -> dict: + """ + Given an amount of `bytes` and `mixed_precision`, calculates how much training memory is needed for a batch size of + 1. + + Args: + bytes (`int`): + The size of the model being trained. + mixed_precision (`str`): + The mixed precision that would be ran. + msamp_config (`str`): + The msamp config to estimate the training memory for if `mixed_precision` is set to `"fp8"`. + """ + memory_sizes = {"model": -1, "optimizer": -1, "gradients": -1, "step": -1} + fp32_size = bytes + fp16_size = bytes // 2 + + if mixed_precision == "float32": + memory_sizes["model"] = fp32_size + memory_sizes["gradients"] = fp32_size + memory_sizes["optimizer"] = fp32_size * 2 + memory_sizes["step"] = fp32_size * 4 + elif mixed_precision in ("float16", "bfloat16") or (mixed_precision == "fp8" and msamp_config is None): + # With native `TransformersEngine`, there is no memory savings with FP8 + # With mixed precision training, the model has weights stored + # in FP16 and FP32 + memory_sizes["model"] = fp32_size + # 1.5 from weight gradient + computation (GEMM) + memory_sizes["gradients"] = fp32_size + fp16_size + # 2x from optimizer states + memory_sizes["optimizer"] = fp32_size * 2 # Optimizer states + memory_sizes["step"] = memory_sizes["optimizer"] + return memory_sizes + + +def gather_data(args): + "Creates an empty model and gathers the data for the sizes" + try: + model = create_empty_model( + args.model_name, library_name=args.library_name, trust_remote_code=args.trust_remote_code + ) + except (RuntimeError, OSError) as e: + library = check_has_model(e) + if library != "unknown": + raise RuntimeError( + f"Tried to load `{args.model_name}` with `{library}` but a possible model to load was not found inside the repo." + ) + raise e + + total_size, largest_layer = calculate_maximum_sizes(model) + + data = [] + + for dtype in args.dtypes: + dtype_total_size = total_size + dtype_largest_layer = largest_layer[0] + dtype_training_size = estimate_training_usage(dtype_total_size, dtype) + if dtype == "float16": + dtype_total_size /= 2 + dtype_largest_layer /= 2 + elif dtype == "int8": + dtype_total_size /= 4 + dtype_largest_layer /= 4 + elif dtype == "int4": + dtype_total_size /= 8 + dtype_largest_layer /= 8 + data.append([dtype, dtype_largest_layer, dtype_total_size, dtype_training_size]) + return data + + +def estimate_command(args): + data = gather_data(args) + for row in data: + for i, item in enumerate(row): + if isinstance(item, (int, float)): + row[i] = convert_bytes(item) + elif isinstance(item, dict): + training_usage = max(item.values()) + row[i] = convert_bytes(training_usage) if training_usage != -1 else "N/A" + + headers = ["dtype", "Largest Layer", "Total Size", "Training using Adam"] + + title = f"Memory Usage for loading `{args.model_name}`" + table = create_ascii_table(headers, data, title) + print(table) + + +def main(): + parser = estimate_command_parser() + args = parser.parse_args() + estimate_command(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/launch.py b/venv/lib/python3.10/site-packages/accelerate/commands/launch.py new file mode 100644 index 0000000000000000000000000000000000000000..7f4304143a3dd67c45ee9ffb9b3e195f8b874edb --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/launch.py @@ -0,0 +1,1209 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import argparse +import importlib +import logging +import os +import subprocess +import sys +from pathlib import Path + +import psutil +import torch + +from accelerate.commands.config import default_config_file, load_config_from_file +from accelerate.commands.config.config_args import SageMakerConfig +from accelerate.commands.config.config_utils import DYNAMO_BACKENDS +from accelerate.commands.utils import CustomArgumentParser +from accelerate.state import get_int_from_env +from accelerate.utils import ( + ComputeEnvironment, + DistributedType, + PrepareForLaunch, + _filter_args, + check_cuda_p2p_ib_support, + convert_dict_to_env_variables, + is_bf16_available, + is_deepspeed_available, + is_hpu_available, + is_mlu_available, + is_musa_available, + is_npu_available, + is_rich_available, + is_sagemaker_available, + is_sdaa_available, + is_torch_xla_available, + is_xpu_available, + patch_environment, + prepare_deepspeed_cmd_env, + prepare_multi_gpu_env, + prepare_sagemager_args_inputs, + prepare_simple_launcher_cmd_env, + prepare_tpu, + str_to_bool, +) +from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS, TORCH_DYNAMO_MODES + + +if is_rich_available(): + from rich import get_console + from rich.logging import RichHandler + + FORMAT = "%(message)s" + logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()]) + + +logger = logging.getLogger(__name__) + + +options_to_group = { + "multi_gpu": "Distributed GPUs", + "tpu": "TPU", + "use_deepspeed": "DeepSpeed Arguments", + "use_fsdp": "FSDP Arguments", + "use_megatron_lm": "Megatron-LM Arguments", + "fp8_backend": "FP8 Arguments", +} + + +def clean_option(option): + "Finds all cases of - after the first two characters and changes them to _" + if "fp8_backend" in option: + option = "--fp8_backend" + if option.startswith("--"): + return option[2:].replace("-", "_") + + +class CustomHelpFormatter(argparse.HelpFormatter): + """ + This is a custom help formatter that will hide all arguments that are not used in the command line when the help is + called. This is useful for the case where the user is using a specific platform and only wants to see the arguments + for that platform. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.titles = [ + "Hardware Selection Arguments", + "Resource Selection Arguments", + "Training Paradigm Arguments", + "positional arguments", + "optional arguments", + ] + + def add_argument(self, action: argparse.Action): + if "accelerate" in sys.argv[0] and "launch" in sys.argv[1:]: + args = sys.argv[2:] + else: + args = sys.argv[1:] + + if len(args) > 1: + args = list(map(clean_option, args)) + used_platforms = [arg for arg in args if arg in options_to_group.keys()] + used_titles = [options_to_group[o] for o in used_platforms] + if action.container.title not in self.titles + used_titles: + action.help = argparse.SUPPRESS + elif action.container.title == "Hardware Selection Arguments": + if set(action.option_strings).isdisjoint(set(args)): + action.help = argparse.SUPPRESS + else: + action.help = action.help + " (currently selected)" + elif action.container.title == "Training Paradigm Arguments": + if set(action.option_strings).isdisjoint(set(args)): + action.help = argparse.SUPPRESS + else: + action.help = action.help + " (currently selected)" + + action.option_strings = [s for s in action.option_strings if "-" not in s[2:]] + super().add_argument(action) + + def end_section(self): + if len(self._current_section.items) < 2: + self._current_section.items = [] + self._current_section.heading = "" + super().end_section() + + +def launch_command_parser(subparsers=None): + description = "Launch a python script in a distributed scenario. Arguments can be passed in with either hyphens (`--num-processes=2`) or underscores (`--num_processes=2`)" + if subparsers is not None: + parser = subparsers.add_parser( + "launch", description=description, add_help=False, allow_abbrev=False, formatter_class=CustomHelpFormatter + ) + else: + parser = CustomArgumentParser( + "Accelerate launch command", + description=description, + add_help=False, + allow_abbrev=False, + formatter_class=CustomHelpFormatter, + ) + + parser.add_argument("-h", "--help", action="help", help="Show this help message and exit.") + + parser.add_argument( + "--config_file", + default=None, + help="The config file to use for the default values in the launching script.", + ) + parser.add_argument( + "--quiet", + "-q", + action="store_true", + help="Silence subprocess errors from the launch stack trace and only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations)", + ) + # Hardware selection arguments + hardware_args = parser.add_argument_group( + "Hardware Selection Arguments", "Arguments for selecting the hardware to be used." + ) + hardware_args.add_argument( + "--cpu", default=False, action="store_true", help="Whether or not to force the training on the CPU." + ) + hardware_args.add_argument( + "--multi_gpu", + default=False, + action="store_true", + help="Whether or not this should launch a distributed GPU training.", + ) + hardware_args.add_argument( + "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training." + ) + hardware_args.add_argument( + "--ipex", + default=False, + action="store_true", + help="Whether or not this should launch a Intel PyTorch Extension (IPEX) training.", + ) + + # Resource selection arguments + resource_args = parser.add_argument_group( + "Resource Selection Arguments", "Arguments for fine-tuning how available hardware should be used." + ) + resource_args.add_argument( + "--mixed_precision", + type=str, + choices=["no", "fp16", "bf16", "fp8"], + help="Whether or not to use mixed precision training. " + "Choose between FP16 and BF16 (bfloat16) training. " + "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.", + ) + resource_args.add_argument( + "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel." + ) + resource_args.add_argument( + "--num_machines", type=int, default=None, help="The total number of machines used in this training." + ) + resource_args.add_argument( + "--num_cpu_threads_per_process", + type=int, + default=None, + help="The number of CPU threads per process. Can be tuned for optimal performance.", + ) + resource_args.add_argument( + "--enable_cpu_affinity", + default=False, + action="store_true", + help="Whether or not CPU affinity and balancing should be enabled. Currently only supported on NVIDIA hardware.", + ) + # Dynamo arguments + resource_args.add_argument( + "--dynamo_backend", + type=str, + choices=["no"] + [b.lower() for b in DYNAMO_BACKENDS], + help="Choose a backend to optimize your training with dynamo, see more at " + "https://github.com/pytorch/torchdynamo.", + ) + resource_args.add_argument( + "--dynamo_mode", + type=str, + default="default", + choices=TORCH_DYNAMO_MODES, + help="Choose a mode to optimize your training with dynamo.", + ) + resource_args.add_argument( + "--dynamo_use_fullgraph", + default=False, + action="store_true", + help="Whether to use full graph mode for dynamo or it is ok to break model into several subgraphs", + ) + resource_args.add_argument( + "--dynamo_use_dynamic", + default=False, + action="store_true", + help="Whether to enable dynamic shape tracing.", + ) + resource_args.add_argument( + "--dynamo_use_regional_compilation", + default=False, + action="store_true", + help="Whether to enable regional compilation.", + ) + + # Training Paradigm arguments + paradigm_args = parser.add_argument_group( + "Training Paradigm Arguments", "Arguments for selecting which training paradigm to be used." + ) + paradigm_args.add_argument( + "--use_deepspeed", + default=False, + action="store_true", + help="Whether to use deepspeed.", + ) + paradigm_args.add_argument( + "--use_fsdp", + default=False, + action="store_true", + help="Whether to use fsdp.", + ) + paradigm_args.add_argument( + "--use_megatron_lm", + default=False, + action="store_true", + help="Whether to use Megatron-LM.", + ) + + paradigm_args.add_argument( + "--use_xpu", + default=None, + action="store_true", + help="Whether to use IPEX plugin to speed up training on XPU specifically. This argument is deprecated and ignored, will be removed in Accelerate v1.20.", + ) + + # distributed GPU training arguments + distributed_args = parser.add_argument_group("Distributed GPUs", "Arguments related to distributed GPU training.") + distributed_args.add_argument( + "--gpu_ids", + default=None, + help="What GPUs (by id) should be used for training on this machine as a comma-separated list", + ) + distributed_args.add_argument( + "--same_network", + default=False, + action="store_true", + help="Whether all machines used for multinode training exist on the same local network.", + ) + distributed_args.add_argument( + "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched." + ) + distributed_args.add_argument( + "--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0." + ) + distributed_args.add_argument( + "--main_process_port", + type=int, + default=None, + help="The port to use to communicate with the machine of rank 0.", + ) + distributed_args.add_argument( + "-t", + "--tee", + default="0", + type=str, + help="Tee std streams into a log file and also to console.", + ) + distributed_args.add_argument( + "--log_dir", + type=str, + default=None, + help=( + "Base directory to use for log files when using torchrun/torch.distributed.run as launcher. " + "Use with --tee to redirect std streams info log files." + ), + ) + distributed_args.add_argument( + "--role", + type=str, + default="default", + help="User-defined role for the workers.", + ) + # Rendezvous related arguments + distributed_args.add_argument( + "--rdzv_backend", + type=str, + default="static", + help="The rendezvous method to use, such as 'static' (the default) or 'c10d'", + ) + distributed_args.add_argument( + "--rdzv_conf", + type=str, + default="", + help="Additional rendezvous configuration (=,=,...).", + ) + distributed_args.add_argument( + "--max_restarts", + type=int, + default=0, + help="Maximum number of worker group restarts before failing.", + ) + distributed_args.add_argument( + "--monitor_interval", + type=float, + default=0.1, + help="Interval, in seconds, to monitor the state of workers.", + ) + parser.add_argument( + "-m", + "--module", + action="store_true", + help="Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.", + ) + parser.add_argument( + "--no_python", + action="store_true", + help="Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.", + ) + + # TPU arguments + tpu_args = parser.add_argument_group("TPU", "Arguments related to TPU.") + tpu_args.add_argument( + "--tpu_cluster", + action="store_true", + dest="tpu_use_cluster", + help="Whether to use a GCP TPU pod for training.", + ) + tpu_args.add_argument( + "--no_tpu_cluster", + action="store_false", + dest="tpu_use_cluster", + help="Should not be passed explicitly, this is for internal use only.", + ) + tpu_args.add_argument( + "--tpu_use_sudo", + action="store_true", + help="Whether to use `sudo` when running the TPU training script in each pod.", + ) + tpu_args.add_argument( + "--vm", + type=str, + action="append", + help=( + "List of single Compute VM instance names. " + "If not provided we assume usage of instance groups. For TPU pods." + ), + ) + tpu_args.add_argument( + "--env", + type=str, + action="append", + help="List of environment variables to set on the Compute VM instances. For TPU pods.", + ) + tpu_args.add_argument( + "--main_training_function", + type=str, + default=None, + help="The name of the main function to be executed in your script (only for TPU training).", + ) + tpu_args.add_argument( + "--downcast_bf16", + action="store_true", + help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.", + ) + + # DeepSpeed arguments + deepspeed_args = parser.add_argument_group("DeepSpeed Arguments", "Arguments related to DeepSpeed.") + deepspeed_args.add_argument( + "--deepspeed_config_file", + default=None, + type=str, + help="DeepSpeed config file.", + ) + deepspeed_args.add_argument( + "--zero_stage", + default=None, + type=int, + help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to `2`.", + ) + deepspeed_args.add_argument( + "--offload_optimizer_device", + default=None, + type=str, + help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to 'none'.", + ) + deepspeed_args.add_argument( + "--offload_param_device", + default=None, + type=str, + help="Decides where (none|cpu|nvme) to offload parameters (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to 'none'.", + ) + deepspeed_args.add_argument( + "--offload_optimizer_nvme_path", + default=None, + type=str, + help="Decides Nvme Path to offload optimizer states (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to 'none'.", + ) + deepspeed_args.add_argument( + "--offload_param_nvme_path", + default=None, + type=str, + help="Decides Nvme Path to offload parameters (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to 'none'.", + ) + deepspeed_args.add_argument( + "--gradient_accumulation_steps", + default=None, + type=int, + help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to `1`.", + ) + deepspeed_args.add_argument( + "--gradient_clipping", + default=None, + type=float, + help="gradient clipping value used in your training script (useful only when `use_deepspeed` flag is passed). " + "If unspecified, will default to `1.0`.", + ) + deepspeed_args.add_argument( + "--zero3_init_flag", + default=None, + type=str, + help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. " + "Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will default to `true`.", + ) + deepspeed_args.add_argument( + "--zero3_save_16bit_model", + default=None, + type=str, + help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. " + "Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will default to `false`.", + ) + deepspeed_args.add_argument( + "--deepspeed_hostfile", + default=None, + type=str, + help="DeepSpeed hostfile for configuring multi-node compute resources.", + ) + deepspeed_args.add_argument( + "--deepspeed_exclusion_filter", + default=None, + type=str, + help="DeepSpeed exclusion filter string when using mutli-node setup.", + ) + deepspeed_args.add_argument( + "--deepspeed_inclusion_filter", + default=None, + type=str, + help="DeepSpeed inclusion filter string when using mutli-node setup.", + ) + deepspeed_args.add_argument( + "--deepspeed_multinode_launcher", + default=None, + type=str, + help="DeepSpeed multi-node launcher to use, e.g. `pdsh`, `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5). If unspecified, will default to `pdsh`.", + ) + deepspeed_args.add_argument( + "--deepspeed_moe_layer_cls_names", + default=None, + type=str, + help="comma-separated list of transformer MoE layer class names (case-sensitive) to wrap ,e.g, `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ..." + " (useful only when `use_deepspeed` flag is passed).", + ) + + # fsdp arguments + fsdp_args = parser.add_argument_group("FSDP Arguments", "Arguments related to Fully Shared Data Parallelism.") + fsdp_args.add_argument( + "--fsdp_version", + type=str, + default="1", + choices=["1", "2"], + help="FSDP version to use. (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_offload_params", + default="false", + type=str, + help="Decides Whether (true|false) to offload parameters and gradients to CPU. (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_min_num_params", + type=int, + default=1e8, + help="FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `use_fsdp` flag is passed).", + ) + # We enable this for backwards compatibility, throw a warning if this is set in `FullyShardedDataParallelPlugin` + fsdp_args.add_argument( + "--fsdp_sharding_strategy", + type=str, + default="FULL_SHARD", + help="FSDP's sharding strategy. (useful only when `use_fsdp` flag is passed and `fsdp_version=1`).", + ) + fsdp_args.add_argument( + "--fsdp_reshard_after_forward", + type=str, + default="true", + help="FSDP's Reshard After Forward Strategy. (useful only when `use_fsdp` flag is passed). Supports either boolean (FSDP2) or `FULL_SHARD | SHARD_GRAD_OP | NO_RESHARD` (FSDP1).", + ) + fsdp_args.add_argument( + "--fsdp_auto_wrap_policy", + type=str, + default=None, + help="FSDP's auto wrap policy. (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_transformer_layer_cls_to_wrap", + default=None, + type=str, + help="Transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` .... " + "(useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_backward_prefetch", + default=None, + type=str, + help="FSDP's backward prefetch policy. (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_state_dict_type", + default=None, + type=str, + help="FSDP's state dict type. (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_forward_prefetch", + default="false", + type=str, + help="If True, then FSDP explicitly prefetches the next upcoming " + "all-gather while executing in the forward pass (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_use_orig_params", + default="true", + type=str, + help="If True, allows non-uniform `requires_grad` during init, which means support for interspersed frozen and trainable paramteres." + " (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_cpu_ram_efficient_loading", + default="true", + type=str, + help="If True, only the first process loads the pretrained model checkoint while all other processes have empty weights. " + "Only applicable for 🤗 Transformers. When using this, `--fsdp_sync_module_states` needs to True. " + "(useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_sync_module_states", + default="true", + type=str, + help="If True, each individually wrapped FSDP unit will broadcast module parameters from rank 0." + " (useful only when `use_fsdp` flag is passed).", + ) + fsdp_args.add_argument( + "--fsdp_activation_checkpointing", + default="false", + type=str, + help="Decides Whether (true|false) intermediate activations are freed during the forward pass, and a checkpoint is left as a placeholder. (useful only when `use_fsdp` flag is passed).", + ) + + # megatron_lm args + megatron_lm_args = parser.add_argument_group("Megatron-LM Arguments", "Arguments related to Megatron-LM.") + megatron_lm_args.add_argument( + "--megatron_lm_tp_degree", + type=int, + default=1, + help="Megatron-LM's Tensor Parallelism (TP) degree. (useful only when `use_megatron_lm` flag is passed).", + ) + megatron_lm_args.add_argument( + "--megatron_lm_pp_degree", + type=int, + default=1, + help="Megatron-LM's Pipeline Parallelism (PP) degree. (useful only when `use_megatron_lm` flag is passed).", + ) + megatron_lm_args.add_argument( + "--megatron_lm_num_micro_batches", + type=int, + default=None, + help="Megatron-LM's number of micro batches when PP degree > 1. (useful only when `use_megatron_lm` flag is passed).", + ) + megatron_lm_args.add_argument( + "--megatron_lm_sequence_parallelism", + default=None, + type=str, + help="Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1. " + "(useful only when `use_megatron_lm` flag is passed).", + ) + megatron_lm_args.add_argument( + "--megatron_lm_recompute_activations", + default=None, + type=str, + help="Decides Whether (true|false) to enable Selective Activation Recomputation. " + "(useful only when `use_megatron_lm` flag is passed).", + ) + megatron_lm_args.add_argument( + "--megatron_lm_use_distributed_optimizer", + default=None, + type=str, + help="Decides Whether (true|false) to use distributed optimizer " + "which shards optimizer state and gradients across Data Pralellel (DP) ranks. " + "(useful only when `use_megatron_lm` flag is passed).", + ) + megatron_lm_args.add_argument( + "--megatron_lm_gradient_clipping", + default=1.0, + type=float, + help="Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable). " + "(useful only when `use_megatron_lm` flag is passed).", + ) + + # FP8 arguments + fp8_args = parser.add_argument_group( + "FP8 Arguments", "Arguments related to FP8 training (requires `--mixed_precision=fp8`)" + ) + fp8_args.add_argument( + "--fp8_backend", + type=str, + choices=["te", "msamp"], + help="Choose a backend to train with FP8 (te: TransformerEngine, msamp: MS-AMP)", + ) + fp8_args.add_argument( + "--fp8_use_autocast_during_eval", + default=False, + action="store_true", + help="Whether to use FP8 autocast during eval mode (useful only when `--fp8_backend=te` is passed). Generally better metrics are found when this is not passed.", + ) + fp8_args.add_argument( + "--fp8_margin", + type=int, + default=0, + help="The margin to use for the gradient scaling (useful only when `--fp8_backend=te` is passed).", + ) + fp8_args.add_argument( + "--fp8_interval", + type=int, + default=1, + help="The interval to use for how often the scaling factor is recomputed (useful only when `--fp8_backend=te` is passed).", + ) + fp8_args.add_argument( + "--fp8_format", + type=str, + default="HYBRID", + choices=["HYBRID", "E4M3", "E5M2"], + help="The format to use for the FP8 recipe (useful only when `--fp8_backend=te` is passed).", + ) + fp8_args.add_argument( + "--fp8_amax_history_len", + type=int, + default=1024, + help="The length of the history to use for the scaling factor computation (useful only when `--fp8_backend=te` is passed).", + ) + fp8_args.add_argument( + "--fp8_amax_compute_algo", + type=str, + default="most_recent", + choices=["max", "most_recent"], + help="The algorithm to use for the scaling factor computation. (useful only when `--fp8_backend=te` is passed).", + ) + fp8_args.add_argument( + "--fp8_override_linear_precision", + type=lambda x: tuple(map(str_to_bool, x.split(","))), + default=(False, False, False), + help="Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision. Should be passed in a comma-separated string of booleans (useful only when `--fp8_backend=te` is passed).", + ) + fp8_args.add_argument( + "--fp8_opt_level", + type=str, + default="O2", + choices=["O1", "O2"], + help="What level of 8-bit collective communication should be used with MS-AMP (useful only when `--fp8_backend=msamp` is passed).", + ) + + # AWS arguments + aws_args = parser.add_argument_group("AWS Arguments", "Arguments related to AWS.") + aws_args.add_argument( + "--aws_access_key_id", + type=str, + default=None, + help="The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job", + ) + aws_args.add_argument( + "--aws_secret_access_key", + type=str, + default=None, + help="The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job.", + ) + parser.add_argument( + "--debug", + action="store_true", + help="Whether to print out the torch.distributed stack trace when something fails.", + ) + parser.add_argument( + "training_script", + type=str, + help=( + "The full path to the script to be launched in parallel, followed by all the arguments for the training " + "script." + ), + ) + + # MPI arguments + mpirun_args = parser.add_argument_group("MPI Arguments", "Arguments related to mpirun for Multi-CPU") + mpirun_args.add_argument( + "--mpirun_hostfile", + type=str, + default=None, + help="Location for a hostfile for using Accelerate to launch a multi-CPU training job with mpirun. This will " + "get passed to the MPI --hostfile or -f parameter, depending on which MPI program is installed.", + ) + mpirun_args.add_argument( + "--mpirun_ccl", + type=int, + default=1, + help="The number of oneCCL worker threads when using Accelerate to launch multi-CPU training with mpirun.", + ) + + # Other arguments of the training scripts + parser.add_argument("training_script_args", nargs=argparse.REMAINDER, help="Arguments of the training script.") + + if subparsers is not None: + parser.set_defaults(func=launch_command) + return parser + + +def simple_launcher(args): + cmd, current_env = prepare_simple_launcher_cmd_env(args) + + process = subprocess.Popen(cmd, env=current_env) + process.wait() + if process.returncode != 0: + if not args.quiet: + raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd) + else: + sys.exit(1) + + +def multi_gpu_launcher(args): + import torch.distributed.run as distrib_run + + current_env = prepare_multi_gpu_env(args) + if not check_cuda_p2p_ib_support(): + message = "Using RTX 4000 series which doesn't support faster communication speedups. Ensuring P2P and IB communications are disabled." + warn = False + if "NCCL_P2P_DISABLE" not in current_env: + current_env["NCCL_P2P_DISABLE"] = "1" + warn = True + if "NCCL_IB_DISABLE" not in current_env: + current_env["NCCL_IB_DISABLE"] = "1" + warn = True + if warn: + logger.warning(message) + + debug = getattr(args, "debug", False) + args = _filter_args( + args, + distrib_run.get_args_parser(), + ["--training_script", args.training_script, "--training_script_args", args.training_script_args], + ) + + with patch_environment(**current_env): + try: + distrib_run.run(args) + except Exception: + if is_rich_available() and debug: + console = get_console() + console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]") + console.print_exception(suppress=[__file__], show_locals=False) + else: + raise + + +def deepspeed_launcher(args): + import torch.distributed.run as distrib_run + + if not is_deepspeed_available(): + raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.") + else: + from deepspeed.launcher.runner import DEEPSPEED_ENVIRONMENT_NAME + + cmd, current_env = prepare_deepspeed_cmd_env(args) + if not check_cuda_p2p_ib_support(): + message = "Using RTX 4000 series which doesn't support faster communication speedups. Ensuring P2P and IB communications are disabled." + warn = False + if "NCCL_P2P_DISABLE" not in current_env: + current_env["NCCL_P2P_DISABLE"] = "1" + warn = True + if "NCCL_IB_DISABLE" not in current_env: + current_env["NCCL_IB_DISABLE"] = "1" + warn = True + if warn: + logger.warning(message) + + if args.num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]: + with open(DEEPSPEED_ENVIRONMENT_NAME, "a") as f: + valid_env_items = convert_dict_to_env_variables(current_env) + if len(valid_env_items) > 1: + f.writelines(valid_env_items) + + process = subprocess.Popen(cmd, env=current_env) + process.wait() + if process.returncode != 0: + if not args.quiet: + raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd) + else: + sys.exit(1) + else: + debug = getattr(args, "debug", False) + args = _filter_args( + args, + distrib_run.get_args_parser(), + ["--training_script", args.training_script, "--training_script_args", args.training_script_args], + ) + with patch_environment(**current_env): + try: + distrib_run.run(args) + except Exception: + if is_rich_available() and debug: + console = get_console() + console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]") + console.print_exception(suppress=[__file__], show_locals=False) + else: + raise + + +def tpu_launcher(args): + import torch_xla.distributed.xla_multiprocessing as xmp + + if args.no_python: + raise ValueError("--no_python cannot be used with TPU launcher") + + args, current_env = prepare_tpu(args, {}) + + if args.module: + mod_name = args.training_script + else: + # Import training_script as a module + script_path = Path(args.training_script) + sys.path.append(str(script_path.parent.resolve())) + mod_name = script_path.stem + + mod = importlib.import_module(mod_name) + if not hasattr(mod, args.main_training_function): + raise ValueError( + f"Your training script should have a function named {args.main_training_function}, or you should pass a " + "different value to `--main_training_function`." + ) + + # Patch sys.argv + sys.argv = [mod.__file__] + args.training_script_args + + main_function = getattr(mod, args.main_training_function) + with patch_environment(**current_env): + xmp.spawn(PrepareForLaunch(main_function), args=()) + + +def tpu_pod_launcher(args): + from torch_xla.distributed import xla_dist + + current_env = {} + args, current_env = prepare_tpu(args, current_env, True) + debug = getattr(args, "debug", False) + + training_script = args.training_script + training_script_args = args.training_script_args + new_args = _filter_args( + args, xla_dist.get_args_parser(), ["--tpu", args.tpu_name, "--positional", "", "--restart-tpuvm-pod-server"] + ) + + if args.tpu_use_sudo: + new_cmd = ["sudo"] + else: + new_cmd = [] + + new_cmd += [ + "accelerate-launch", + "--tpu", + "--no_tpu_cluster", + "--num_machines", + "1", + "--mixed_precision", + "no", + "--dynamo_backend", + "no", + "--num_processes", + str(args.num_processes), + "--main_training_function", + str(args.main_training_function), + training_script, + ] + training_script_args + + new_args.positional = new_cmd + bad_flags = "" + for arg in vars(new_args): + if arg.startswith("docker_"): + value = getattr(new_args, arg) + if value != "" and value is not None: + bad_flags += f'{arg}="{value}"\n' + if bad_flags != "": + raise ValueError( + f"Docker containers are not supported for TPU pod launcher currently, please remove the following flags:\n{bad_flags}" + ) + new_args.env = [f"{k}={v}" for k, v in current_env.items()] + new_args.env.append("ACCELERATE_IN_TPU_POD=1") + try: + xla_dist.resolve_and_execute(new_args) + except Exception: + if is_rich_available() and debug: + console = get_console() + console.print("\n[bold red]Using --debug, `torch_xla.xla_dist` Stack Trace:[/bold red]") + console.print_exception(suppress=[__file__], show_locals=False) + else: + raise + + +def sagemaker_launcher(sagemaker_config: SageMakerConfig, args): + if not is_sagemaker_available(): + raise ImportError( + "Please install sagemaker to be able to launch training on Amazon SageMaker with `pip install accelerate[sagemaker]`" + ) + if args.module or args.no_python: + raise ValueError( + "SageMaker requires a python training script file and cannot be used with --module or --no_python" + ) + + from sagemaker.huggingface import HuggingFace + + args, sagemaker_inputs = prepare_sagemager_args_inputs(sagemaker_config, args) + + huggingface_estimator = HuggingFace(**args) + + huggingface_estimator.fit(inputs=sagemaker_inputs) + print(f"You can find your model data at: {huggingface_estimator.model_data}") + + +def _validate_launch_command(args): + # Sanity checks + if sum([args.multi_gpu, args.cpu, args.tpu, args.use_deepspeed, args.use_fsdp]) > 1: + raise ValueError( + "You can only use one of `--cpu`, `--multi_gpu`, `--tpu`, `--use_deepspeed`, `--use_fsdp` at a time." + ) + if args.multi_gpu and (args.num_processes is not None) and (args.num_processes < 2): + raise ValueError("You need to use at least 2 processes to use `--multi_gpu`.") + + defaults = None + warned = [] + mp_from_config_flag = False + # Get the default from the config file. + if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu: + defaults = load_config_from_file(args.config_file) + if ( + not args.multi_gpu + and not args.tpu + and not args.tpu_use_cluster + and not args.use_deepspeed + and not args.use_fsdp + and not args.use_megatron_lm + ): + args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED + args.multi_gpu = ( + True + if defaults.distributed_type + in ( + DistributedType.MULTI_GPU, + DistributedType.MULTI_NPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + ) + else False + ) + args.tpu = defaults.distributed_type == DistributedType.XLA + args.use_fsdp = defaults.distributed_type == DistributedType.FSDP + args.use_megatron_lm = defaults.distributed_type == DistributedType.MEGATRON_LM + args.tpu_use_cluster = defaults.tpu_use_cluster if args.tpu else False + if args.gpu_ids is None: + if defaults.gpu_ids is not None: + args.gpu_ids = defaults.gpu_ids + else: + args.gpu_ids = "all" + + if args.multi_gpu and args.num_machines is None: + args.num_machines = defaults.num_machines + + if len(args.gpu_ids.split(",")) < 2 and (args.gpu_ids != "all") and args.multi_gpu and args.num_machines <= 1: + raise ValueError( + "Less than two GPU ids were configured and tried to run on on multiple GPUs. " + "Please ensure at least two are specified for `--gpu_ids`, or use `--gpu_ids='all'`." + ) + if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE: + # Update args with the defaults + for name, attr in defaults.__dict__.items(): + if isinstance(attr, dict): + # Copy defaults.somedict.somearg to args.somearg and + # defaults.fsdp_config.x to args.fsdp_x + for key, value in attr.items(): + if name == "fsdp_config" and not key.startswith("fsdp"): + key = "fsdp_" + key + elif name == "fp8_config" and not key.startswith("fp8"): + key = "fp8_" + key + if hasattr(args, "nondefault") and key not in args.nondefault: + setattr(args, key, value) + elif ( + name not in ["compute_environment", "mixed_precision", "distributed_type"] + and getattr(args, name, None) is None + ): + # Those args are handled separately + setattr(args, name, attr) + if not args.debug: + args.debug = defaults.debug + + if not args.mixed_precision: + if defaults.mixed_precision is None: + args.mixed_precision = "no" + else: + args.mixed_precision = defaults.mixed_precision + mp_from_config_flag = True + else: + native_amp = is_bf16_available(True) + if ( + args.mixed_precision == "bf16" + and not native_amp + and not (args.tpu and is_torch_xla_available(check_is_tpu=True)) + ): + raise ValueError("bf16 mixed precision requires PyTorch >= 1.10 and a supported device.") + + # Silently set the default here + if args.dynamo_backend is None: + args.dynamo_backend = "no" + if args.num_processes == -1: + raise ValueError("You need to manually pass in `--num_processes` using this config yaml.") + else: + if args.num_processes is None: + if is_xpu_available(): + args.num_processes = torch.xpu.device_count() + elif is_mlu_available(): + args.num_processes = torch.mlu.device_count() + elif is_sdaa_available(): + args.num_processes = torch.sdaa.device_count() + elif is_musa_available(): + args.num_processes = torch.musa.device_count() + elif is_npu_available(): + args.num_processes = torch.npu.device_count() + elif is_hpu_available(): + args.num_processes = torch.hpu.device_count() + else: + args.num_processes = torch.cuda.device_count() + warned.append(f"\t`--num_processes` was set to a value of `{args.num_processes}`") + if args.debug is None: + args.debug = False + if ( + not args.multi_gpu + and args.num_processes > 1 + and ( + (is_xpu_available() and torch.xpu.device_count() > 1) + or (is_npu_available() and torch.npu.device_count() > 1) + or (is_hpu_available() and torch.hpu.device_count() > 1) + or (is_mlu_available() and torch.mlu.device_count() > 1) + or (is_sdaa_available() and torch.sdaa.device_count() > 1) + or (is_musa_available() and torch.musa.device_count() > 1) + or (torch.cuda.is_available() and torch.cuda.device_count() > 1) + ) + ): + warned.append( + "\t\tMore than one GPU was found, enabling multi-GPU training.\n" + "\t\tIf this was unintended please pass in `--num_processes=1`." + ) + args.multi_gpu = True + if args.num_machines is None: + warned.append("\t`--num_machines` was set to a value of `1`") + args.num_machines = 1 + if args.mixed_precision is None: + warned.append("\t`--mixed_precision` was set to a value of `'no'`") + args.mixed_precision = "no" + if not hasattr(args, "use_cpu"): + args.use_cpu = args.cpu + if args.dynamo_backend is None: + warned.append("\t`--dynamo_backend` was set to a value of `'no'`") + args.dynamo_backend = "no" + if args.debug: + logger.debug("Running script in debug mode, expect distributed operations to be slightly slower.") + + is_aws_env_disabled = defaults is None or ( + defaults is not None and defaults.compute_environment != ComputeEnvironment.AMAZON_SAGEMAKER + ) + if is_aws_env_disabled and args.num_cpu_threads_per_process is None: + args.num_cpu_threads_per_process = get_int_from_env(["OMP_NUM_THREADS"], 1) + if args.use_cpu and args.num_processes >= 1 and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0: + local_size = get_int_from_env( + ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], + max(int(args.num_processes / args.num_machines), 1), + ) + threads_per_process = int(psutil.cpu_count(logical=False) / local_size) + if threads_per_process > 1: + args.num_cpu_threads_per_process = threads_per_process + warned.append( + f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance when training on CPUs" + ) + + if args.ipex is not None: + logger.warning( + "ipex flag is deprecated, will be removed in Accelerate v1.10. " + "From 2.7.0, PyTorch has all needed optimizations for Intel CPU and XPU." + ) + + if args.use_xpu is not None: + logger.warning( + "use_xpu is deprecated and ignored, will be removed in Accelerate v1.20. " + "XPU is a PyTorch native citizen now, we don't need extra argument to enable it any more." + ) + + if any(warned): + message = "The following values were not passed to `accelerate launch` and had defaults used instead:\n" + message += "\n".join(warned) + message += ( + "\nTo avoid this warning pass in values for each of the problematic parameters or run `accelerate config`." + ) + logger.warning(message) + return args, defaults, mp_from_config_flag + + +def launch_command(args): + args, defaults, mp_from_config_flag = _validate_launch_command(args) + # Use the proper launcher + if args.use_deepspeed and not args.cpu: + args.deepspeed_fields_from_accelerate_config = list(defaults.deepspeed_config.keys()) if defaults else [] + if mp_from_config_flag: + args.deepspeed_fields_from_accelerate_config.append("mixed_precision") + args.deepspeed_fields_from_accelerate_config = ",".join(args.deepspeed_fields_from_accelerate_config) + deepspeed_launcher(args) + elif args.use_fsdp and not args.cpu: + multi_gpu_launcher(args) + elif args.use_megatron_lm and not args.cpu: + multi_gpu_launcher(args) + elif args.multi_gpu and not args.cpu: + multi_gpu_launcher(args) + elif args.tpu and not args.cpu: + if args.tpu_use_cluster: + tpu_pod_launcher(args) + else: + tpu_launcher(args) + elif defaults is not None and defaults.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER: + sagemaker_launcher(defaults, args) + else: + simple_launcher(args) + + +def main(): + parser = launch_command_parser() + args = parser.parse_args() + launch_command(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__init__.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c2c851cc0b192ab8207d3fa68d7409868c84354c --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2022 The HuggingFace Team. 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 .selection_menu import BulletMenu diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c565b5bf6119aca3772d5bfbf7413927cbeb7594 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11616afac2865b4ba21ef276ea052b3085d50725 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..375f54ea3b9b20839c4506226c7ed1ac366f3220 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/input.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/input.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a90a63e8fe6a7dd0852c9070aa7d3041aca0c473 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/input.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ccbd09290b2217c4694527eb1373bff993a99d9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3abda192fb278fbc69b88e573a2a5a4d78a7117b Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/cursor.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/cursor.py new file mode 100644 index 0000000000000000000000000000000000000000..c1f0bb7b68025ae4fe0c2c76c095eb36b4e64f2c --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/menu/cursor.py @@ -0,0 +1,65 @@ +# Copyright 2022 The HuggingFace Team and Brian Chao. 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. + +""" +A utility for showing and hiding the terminal cursor on Windows and Linux, based on https://github.com/bchao1/bullet +""" + +import os +import sys +from contextlib import contextmanager + + +# Windows only +if os.name == "nt": + import ctypes + import msvcrt # noqa + + class CursorInfo(ctypes.Structure): + # _fields is a specific attr expected by ctypes + _fields_ = [("size", ctypes.c_int), ("visible", ctypes.c_byte)] + + +def hide_cursor(): + if os.name == "nt": + ci = CursorInfo() + handle = ctypes.windll.kernel32.GetStdHandle(-11) + ctypes.windll.kernel32.GetConsoleCursorInfo(handle, ctypes.byref(ci)) + ci.visible = False + ctypes.windll.kernel32.SetConsoleCursorInfo(handle, ctypes.byref(ci)) + elif os.name == "posix": + sys.stdout.write("\033[?25l") + sys.stdout.flush() + + +def show_cursor(): + if os.name == "nt": + ci = CursorInfo() + handle = ctypes.windll.kernel32.GetStdHandle(-11) + ctypes.windll.kernel32.GetConsoleCursorInfo(handle, ctypes.byref(ci)) + ci.visible = True + ctypes.windll.kernel32.SetConsoleCursorInfo(handle, ctypes.byref(ci)) + elif os.name == "posix": + sys.stdout.write("\033[?25h") + sys.stdout.flush() + + +@contextmanager +def hide(): + "Context manager to hide the terminal cursor" + try: + hide_cursor() + yield + finally: + show_cursor() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/helpers.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..de46f37ddcf4591167e3e01791391e4b1729034f --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/menu/helpers.py @@ -0,0 +1,59 @@ +# Copyright 2022 The HuggingFace Team and Brian Chao. 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. + +""" +A variety of helper functions and constants when dealing with terminal menu choices, based on +https://github.com/bchao1/bullet +""" + +import enum +import shutil +import sys + + +TERMINAL_WIDTH, _ = shutil.get_terminal_size() + +CURSOR_TO_CHAR = {"UP": "A", "DOWN": "B", "RIGHT": "C", "LEFT": "D"} + + +class Direction(enum.Enum): + UP = 0 + DOWN = 1 + + +def forceWrite(content, end=""): + sys.stdout.write(str(content) + end) + sys.stdout.flush() + + +def writeColor(content, color, end=""): + forceWrite(f"\u001b[{color}m{content}\u001b[0m", end) + + +def reset_cursor(): + forceWrite("\r") + + +def move_cursor(num_lines: int, direction: str): + forceWrite(f"\033[{num_lines}{CURSOR_TO_CHAR[direction.upper()]}") + + +def clear_line(): + forceWrite(" " * TERMINAL_WIDTH) + reset_cursor() + + +def linebreak(): + reset_cursor() + forceWrite("-" * TERMINAL_WIDTH) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/input.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/input.py new file mode 100644 index 0000000000000000000000000000000000000000..f1270eaece9d4243e7282dcb31166feeeb9bdfc1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/menu/input.py @@ -0,0 +1,84 @@ +# Copyright 2022 The HuggingFace Team and Brian Chao. 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. + +""" +This file contains utilities for handling input from the user and registering specific keys to specific functions, +based on https://github.com/bchao1/bullet +""" + +from .keymap import KEYMAP, get_character + + +def mark(key: str): + """ + Mark the function with the key code so it can be handled in the register + """ + + def decorator(func): + handle = getattr(func, "handle_key", []) + handle += [key] + func.handle_key = handle + return func + + return decorator + + +def mark_multiple(*keys: list[str]): + """ + Mark the function with the key codes so it can be handled in the register + """ + + def decorator(func): + handle = getattr(func, "handle_key", []) + handle += keys + func.handle_key = handle + return func + + return decorator + + +class KeyHandler(type): + """ + Metaclass that adds the key handlers to the class + """ + + def __new__(cls, name, bases, attrs): + new_cls = super().__new__(cls, name, bases, attrs) + if not hasattr(new_cls, "key_handler"): + new_cls.key_handler = {} + new_cls.handle_input = KeyHandler.handle_input + + for value in attrs.values(): + handled_keys = getattr(value, "handle_key", []) + for key in handled_keys: + new_cls.key_handler[key] = value + return new_cls + + @staticmethod + def handle_input(cls): + "Finds and returns the selected character if it exists in the handler" + char = get_character() + if char != KEYMAP["undefined"]: + char = ord(char) + handler = cls.key_handler.get(char) + if handler: + cls.current_selection = char + return handler(cls) + else: + return None + + +def register(cls): + """Adds KeyHandler metaclass to the class""" + return KeyHandler(cls.__name__, cls.__bases__, cls.__dict__.copy()) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/keymap.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/keymap.py new file mode 100644 index 0000000000000000000000000000000000000000..787db12860fe21c6786dda69c34fcccab114f2f8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/menu/keymap.py @@ -0,0 +1,133 @@ +# Copyright 2022 The HuggingFace Team and Brian Chao. 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. + +""" +Utilities relating to parsing raw characters from the keyboard, based on https://github.com/bchao1/bullet +""" + +import os +import string +import sys + + +ARROW_KEY_FLAG = 1 << 8 + +KEYMAP = { + "tab": ord("\t"), + "newline": ord("\r"), + "esc": 27, + "up": 65 + ARROW_KEY_FLAG, + "down": 66 + ARROW_KEY_FLAG, + "right": 67 + ARROW_KEY_FLAG, + "left": 68 + ARROW_KEY_FLAG, + "mod_int": 91, + "undefined": sys.maxsize, + "interrupt": 3, + "insert": 50, + "delete": 51, + "pg_up": 53, + "pg_down": 54, +} + +KEYMAP["arrow_begin"] = KEYMAP["up"] +KEYMAP["arrow_end"] = KEYMAP["left"] + +if sys.platform == "win32": + WIN_CH_BUFFER = [] + WIN_KEYMAP = { + b"\xe0H": KEYMAP["up"] - ARROW_KEY_FLAG, + b"\x00H": KEYMAP["up"] - ARROW_KEY_FLAG, + b"\xe0P": KEYMAP["down"] - ARROW_KEY_FLAG, + b"\x00P": KEYMAP["down"] - ARROW_KEY_FLAG, + b"\xe0M": KEYMAP["right"] - ARROW_KEY_FLAG, + b"\x00M": KEYMAP["right"] - ARROW_KEY_FLAG, + b"\xe0K": KEYMAP["left"] - ARROW_KEY_FLAG, + b"\x00K": KEYMAP["left"] - ARROW_KEY_FLAG, + } + +for i in range(10): + KEYMAP[str(i)] = ord(str(i)) + + +def get_raw_chars(): + "Gets raw characters from inputs" + if os.name == "nt": + import msvcrt + + encoding = "mbcs" + # Flush the keyboard buffer + while msvcrt.kbhit(): + msvcrt.getch() + if len(WIN_CH_BUFFER) == 0: + # Read the keystroke + ch = msvcrt.getch() + + # If it is a prefix char, get second part + if ch in (b"\x00", b"\xe0"): + ch2 = ch + msvcrt.getch() + # Translate actual Win chars to bullet char types + try: + chx = chr(WIN_KEYMAP[ch2]) + WIN_CH_BUFFER.append(chr(KEYMAP["mod_int"])) + WIN_CH_BUFFER.append(chx) + if ord(chx) in ( + KEYMAP["insert"] - 1 << 9, + KEYMAP["delete"] - 1 << 9, + KEYMAP["pg_up"] - 1 << 9, + KEYMAP["pg_down"] - 1 << 9, + ): + WIN_CH_BUFFER.append(chr(126)) + ch = chr(KEYMAP["esc"]) + except KeyError: + ch = ch2[1] + else: + ch = ch.decode(encoding) + else: + ch = WIN_CH_BUFFER.pop(0) + elif os.name == "posix": + import termios + import tty + + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) + try: + tty.setraw(fd) + ch = sys.stdin.read(1) + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + return ch + + +def get_character(): + "Gets a character from the keyboard and returns the key code" + char = get_raw_chars() + if ord(char) in [KEYMAP["interrupt"], KEYMAP["newline"]]: + return char + + elif ord(char) == KEYMAP["esc"]: + combo = get_raw_chars() + if ord(combo) == KEYMAP["mod_int"]: + key = get_raw_chars() + if ord(key) >= KEYMAP["arrow_begin"] - ARROW_KEY_FLAG and ord(key) <= KEYMAP["arrow_end"] - ARROW_KEY_FLAG: + return chr(ord(key) + ARROW_KEY_FLAG) + else: + return KEYMAP["undefined"] + else: + return get_raw_chars() + + else: + if char in string.printable: + return char + else: + return KEYMAP["undefined"] diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/selection_menu.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/selection_menu.py new file mode 100644 index 0000000000000000000000000000000000000000..ee9a771a54ef666ee46b67ae6c75fb957d49efdd --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/menu/selection_menu.py @@ -0,0 +1,144 @@ +# Copyright 2022 The HuggingFace Team and Brian Chao. 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. + +""" +Main driver for the selection menu, based on https://github.com/bchao1/bullet +""" + +import builtins +import sys + +from ...utils.imports import _is_package_available +from . import cursor, input +from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor +from .keymap import KEYMAP + + +in_colab = False +try: + in_colab = _is_package_available("google.colab") +except ModuleNotFoundError: + pass + + +@input.register +class BulletMenu: + """ + A CLI menu to select a choice from a list of choices using the keyboard. + """ + + def __init__(self, prompt: str = None, choices: list = []): + self.position = 0 + self.choices = choices + self.prompt = prompt + if sys.platform == "win32": + self.arrow_char = "*" + else: + self.arrow_char = "➔ " + + def write_choice(self, index, end: str = ""): + if sys.platform != "win32": + writeColor(self.choices[index], 32, end) + else: + forceWrite(self.choices[index], end) + + def print_choice(self, index: int): + "Prints the choice at the given index" + if index == self.position: + forceWrite(f" {self.arrow_char} ") + self.write_choice(index) + else: + forceWrite(f" {self.choices[index]}") + reset_cursor() + + def move_direction(self, direction: Direction, num_spaces: int = 1): + "Should not be directly called, used to move a direction of either up or down" + old_position = self.position + if direction == Direction.DOWN: + if self.position + 1 >= len(self.choices): + return + self.position += num_spaces + else: + if self.position - 1 < 0: + return + self.position -= num_spaces + clear_line() + self.print_choice(old_position) + move_cursor(num_spaces, direction.name) + self.print_choice(self.position) + + @input.mark(KEYMAP["up"]) + def move_up(self): + self.move_direction(Direction.UP) + + @input.mark(KEYMAP["down"]) + def move_down(self): + self.move_direction(Direction.DOWN) + + @input.mark(KEYMAP["newline"]) + def select(self): + move_cursor(len(self.choices) - self.position, "DOWN") + return self.position + + @input.mark(KEYMAP["interrupt"]) + def interrupt(self): + move_cursor(len(self.choices) - self.position, "DOWN") + raise KeyboardInterrupt + + @input.mark_multiple(*[KEYMAP[str(number)] for number in range(10)]) + def select_row(self): + index = int(chr(self.current_selection)) + movement = index - self.position + if index == self.position: + return + if index < len(self.choices): + if self.position > index: + self.move_direction(Direction.UP, -movement) + elif self.position < index: + self.move_direction(Direction.DOWN, movement) + else: + return + else: + return + + def run(self, default_choice: int = 0): + "Start the menu and return the selected choice" + if self.prompt: + linebreak() + forceWrite(self.prompt, "\n") + if in_colab: + forceWrite("Please input a choice index (starting from 0), and press enter", "\n") + else: + forceWrite("Please select a choice using the arrow or number keys, and selecting with enter", "\n") + self.position = default_choice + for i in range(len(self.choices)): + self.print_choice(i) + forceWrite("\n") + move_cursor(len(self.choices) - self.position, "UP") + with cursor.hide(): + while True: + if in_colab: + try: + choice = int(builtins.input()) + except ValueError: + choice = default_choice + else: + choice = self.handle_input() + if choice is not None: + reset_cursor() + for _ in range(len(self.choices) + 1): + move_cursor(1, "UP") + clear_line() + self.write_choice(choice, "\n") + return choice diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/merge.py b/venv/lib/python3.10/site-packages/accelerate/commands/merge.py new file mode 100644 index 0000000000000000000000000000000000000000..475b53b5bbb71b959057126f8667d7f61eb9d0e1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/merge.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python + +# Copyright 2024 The HuggingFace Team. 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 accelerate.commands.utils import CustomArgumentParser +from accelerate.utils import merge_fsdp_weights + + +description = """Utility to merge the weights from multiple FSDP checkpoints into a single combined checkpoint. Should be used if +`SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}`. + +This is a CPU-bound process and requires enough RAM to load the entire model state dict.""" + + +def merge_command(args): + merge_fsdp_weights( + args.checkpoint_directory, args.output_path, not args.unsafe_serialization, args.remove_checkpoint_dir + ) + + +def merge_command_parser(subparsers=None): + if subparsers is not None: + parser = subparsers.add_parser("merge-weights", description=description) + else: + parser = CustomArgumentParser(description=description) + + parser.add_argument("checkpoint_directory", type=str, help="A directory containing sharded weights saved by FSDP.") + parser.add_argument( + "output_path", + type=str, + help="The path to save the merged weights. Defaults to the current directory. ", + ) + parser.add_argument( + "--unsafe_serialization", + action="store_true", + default=False, + help="Whether to save the merged weights as `.bin` rather than `.safetensors` (not recommended).", + ) + parser.add_argument( + "--remove_checkpoint_dir", + action="store_true", + help="Whether to remove the checkpoint directory after merging.", + default=False, + ) + + if subparsers is not None: + parser.set_defaults(func=merge_command) + return parser + + +def main(): + parser = merge_command_parser() + args = parser.parse_args() + merge_command(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/test.py b/venv/lib/python3.10/site-packages/accelerate/commands/test.py new file mode 100644 index 0000000000000000000000000000000000000000..a0d2f7bcf14727aa13e3438f4cd6e6f140f5bb2f --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/test.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Team. 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. + +import argparse + +from accelerate.test_utils import execute_subprocess_async, path_in_accelerate_package + + +def test_command_parser(subparsers=None): + if subparsers is not None: + parser = subparsers.add_parser("test") + else: + parser = argparse.ArgumentParser("Accelerate test command") + + parser.add_argument( + "--config_file", + default=None, + help=( + "The path to use to store the config file. Will default to a file named default_config.yaml in the cache " + "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " + "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " + "with 'huggingface'." + ), + ) + + if subparsers is not None: + parser.set_defaults(func=test_command) + return parser + + +def test_command(args): + script_name = path_in_accelerate_package("test_utils", "scripts", "test_script.py") + + if args.config_file is None: + test_args = [script_name] + else: + test_args = f"--config_file={args.config_file} {script_name}".split() + + cmd = ["accelerate-launch"] + test_args + result = execute_subprocess_async(cmd) + if result.returncode == 0: + print("Test is a success! You are ready for your distributed training!") + + +def main(): + parser = test_command_parser() + args = parser.parse_args() + test_command(args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/to_fsdp2.py b/venv/lib/python3.10/site-packages/accelerate/commands/to_fsdp2.py new file mode 100644 index 0000000000000000000000000000000000000000..443407cd983dcb31711b75c2a6337f9a7af24584 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/to_fsdp2.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Team. 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. + +import enum +import logging +from pathlib import Path + +import yaml + +from accelerate.commands.utils import CustomArgumentParser + + +class ConversionStatus(enum.Enum): + NOT_YET_IMPLEMENTED = 0 + REMOVED = -1 + + +ARGUMENT_KEY_MAPPING = { + # New keys in FSDP2 + "fsdp_version": "fsdp_version", + "fsdp_reshard_after_forward": "fsdp_reshard_after_forward", + # https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md + # https://huggingface.co/docs/accelerate/en/usage_guides/fsdp + "fsdp_auto_wrap_policy": "fsdp_auto_wrap_policy", + "fsdp_backward_prefetch": ConversionStatus.REMOVED, + "fsdp_forward_prefetch": ConversionStatus.NOT_YET_IMPLEMENTED, + "fsdp_cpu_ram_efficient_loading": "fsdp_cpu_ram_efficient_loading", + "fsdp_offload_params": "fsdp_offload_params", + "fsdp_sharding_strategy": "fsdp_reshard_after_forward", + "fsdp_state_dict_type": "fsdp_state_dict_type", + "fsdp_sync_module_states": ConversionStatus.REMOVED, + "fsdp_transformer_layer_cls_to_wrap": "fsdp_transformer_layer_cls_to_wrap", + "fsdp_min_num_params": "fsdp_min_num_params", + "fsdp_use_orig_params": ConversionStatus.REMOVED, + "fsdp_activation_checkpointing": "fsdp_activation_checkpointing", +} + +ARGUMENT_VALUE_MAPPING = { + "fsdp_sharding_strategy": { + "FULL_SHARD": True, + "SHARD_GRAD_OP": False, + "HYBRID_SHARD": True, + "HYBRID_SHARD_ZERO2": False, + "NO_SHARD": False, + }, + "fsdp_reshard_after_forward": { # Needed to convert newly created configs using FSDP1 to FSDP2 + "FULL_SHARD": True, + "SHARD_GRAD_OP": False, + "HYBRID_SHARD": True, + "HYBRID_SHARD_ZERO2": False, + "NO_SHARD": False, + }, +} + +logger = logging.getLogger(__name__) + + +def _validate_to_fsdp2_args(args): + if not Path(args.config_file).exists(): + raise FileNotFoundError(f"Config file {args.config_file} not found") + + if not args.overwrite and args.output_file is None: + raise ValueError("If --overwrite is not set, --output_file must be provided") + + if not args.overwrite and Path(args.output_file).exists(): + raise FileExistsError(f"Output file {args.output_file} already exists and --overwrite is not set") + + +def convert_config_to_fsdp2(config: dict) -> dict: + fsdp_config = config.get("fsdp_config", {}) + + if not fsdp_config: + logger.info("No FSDP config found in the config file, skipping conversion...") + return config + + new_fsdp_config = {} + + if fsdp_config.get("fsdp_version", 1) == 2: + logger.warning("Config already specfies FSDP2, skipping conversion...") + logger.warning( + "If the config doesn't use new argument names, change `fsdp_version` to `1` and rerun the command." + ) + return config + + for key, value in fsdp_config.items(): + conversion_status = ARGUMENT_KEY_MAPPING.get(key, None) + if isinstance(conversion_status, ConversionStatus) or conversion_status is None: + conversion_status = key + new_fsdp_config[conversion_status] = value + continue + + if conversion_status == ConversionStatus.REMOVED: + logger.warning(f"Argument {key} has been removed in FSDP2, skipping this key...") + continue + + if conversion_status == ConversionStatus.NOT_YET_IMPLEMENTED: + logger.warning(f"Argument {key} is not yet implemented in FSDP2, skipping this key...") + continue + + if conversion_status is None: + logger.warning(f"Argument {key} is not being converted, skipping this key...") + new_fsdp_config[key] = value + else: + if key in ARGUMENT_VALUE_MAPPING: + value = ARGUMENT_VALUE_MAPPING[key].get(value, value) + new_fsdp_config[ARGUMENT_KEY_MAPPING[key]] = value + + new_fsdp_config["fsdp_version"] = 2 + config["fsdp_config"] = new_fsdp_config + return config + + +def to_fsdp2_command_parser(subparsers=None): + description = "Convert an Accelerate config from FSDP1 to FSDP2" + + if subparsers is not None: + parser = subparsers.add_parser("to-fsdp2", description=description) + else: + parser = CustomArgumentParser(description=description) + + parser.add_argument("--config_file", type=str, help="The config file to convert to FSDP2", required=True) + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite the config file if it exists", + default=False, + ) + parser.add_argument( + "--output_file", + type=str, + help="The path to the output file to write the converted config to. If not provided, the input file will be overwritten (if --overwrite is set)", + default=None, + ) + if subparsers is not None: + parser.set_defaults(func=to_fsdp2_command) + + return parser + + +def load_config(config_file: str) -> dict: + with open(config_file) as f: + config = yaml.safe_load(f) + if not config: + raise ValueError("Config file is empty") + + return config + + +def to_fsdp2_command(args): + _validate_to_fsdp2_args(args) + config = load_config(args.config_file) + + if args.overwrite and args.output_file is None: + args.output_file = args.config_file + + new_config = convert_config_to_fsdp2(config) + + with open(args.output_file, "w") as f: + yaml.dump(new_config, f) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/tpu.py b/venv/lib/python3.10/site-packages/accelerate/commands/tpu.py new file mode 100644 index 0000000000000000000000000000000000000000..fc0f07bf8697bfdb6484d3bf817f2e18b1313b00 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/tpu.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python + +# Copyright 2022 The HuggingFace Team. 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. + +import argparse +import os +import subprocess + +from packaging.version import Version, parse + +from accelerate.commands.config.config_args import default_config_file, load_config_from_file + + +_description = "Run commands across TPU VMs for initial setup before running `accelerate launch`." + + +def tpu_command_parser(subparsers=None): + if subparsers is not None: + parser = subparsers.add_parser("tpu-config", description=_description) + else: + parser = argparse.ArgumentParser("Accelerate tpu-config command", description=_description) + # Core arguments + config_args = parser.add_argument_group( + "Config Arguments", "Arguments that can be configured through `accelerate config`." + ) + config_args.add_argument( + "--config_file", + type=str, + default=None, + help="Path to the config file to use for accelerate.", + ) + config_args.add_argument( + "--tpu_name", + default=None, + help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.", + ) + config_args.add_argument( + "--tpu_zone", + default=None, + help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.", + ) + pod_args = parser.add_argument_group("TPU Arguments", "Arguments for options ran inside the TPU.") + pod_args.add_argument( + "--use_alpha", + action="store_true", + help="Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`.", + ) + pod_args.add_argument( + "--command_file", + default=None, + help="The path to the file containing the commands to run on the pod on startup.", + ) + pod_args.add_argument( + "--command", + action="append", + nargs="+", + help="A command to run on the pod. Can be passed multiple times.", + ) + pod_args.add_argument( + "--install_accelerate", + action="store_true", + help="Whether to install accelerate on the pod. Defaults to False.", + ) + pod_args.add_argument( + "--accelerate_version", + default="latest", + help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.", + ) + pod_args.add_argument( + "--debug", action="store_true", help="If set, will print the command that would be run instead of running it." + ) + + if subparsers is not None: + parser.set_defaults(func=tpu_command_launcher) + return parser + + +def tpu_command_launcher(args): + defaults = None + + # Get the default from the config file if it exists. + if args.config_file is not None or os.path.isfile(default_config_file): + defaults = load_config_from_file(args.config_file) + if not args.command_file and defaults.command_file is not None and not args.command: + args.command_file = defaults.command_file + if not args.command and defaults.commands is not None: + args.command = defaults.commands + if not args.tpu_name: + args.tpu_name = defaults.tpu_name + if not args.tpu_zone: + args.tpu_zone = defaults.tpu_zone + if args.accelerate_version == "dev": + args.accelerate_version = "git+https://github.com/huggingface/accelerate.git" + elif args.accelerate_version == "latest": + args.accelerate_version = "accelerate -U" + elif isinstance(parse(args.accelerate_version), Version): + args.accelerate_version = f"accelerate=={args.accelerate_version}" + + if not args.command_file and not args.command: + raise ValueError("You must specify either a command file or a command to run on the pod.") + + if args.command_file: + with open(args.command_file) as f: + args.command = [f.read().splitlines()] + + # To turn list of lists into list of strings + if isinstance(args.command[0], list): + args.command = [line for cmd in args.command for line in cmd] + # Default to the shared folder and install accelerate + new_cmd = ["cd /usr/share"] + if args.install_accelerate: + new_cmd += [f"pip install {args.accelerate_version}"] + new_cmd += args.command + args.command = "; ".join(new_cmd) + + # Then send it to gcloud + # Eventually try to use google-api-core to do this instead of subprocess + cmd = ["gcloud"] + if args.use_alpha: + cmd += ["alpha"] + cmd += [ + "compute", + "tpus", + "tpu-vm", + "ssh", + args.tpu_name, + "--zone", + args.tpu_zone, + "--command", + args.command, + "--worker", + "all", + ] + if args.debug: + print(f"Running {' '.join(cmd)}") + return + subprocess.run(cmd) + print("Successfully setup pod.") + + +def main(): + parser = tpu_command_parser() + args = parser.parse_args() + + tpu_command_launcher(args) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/utils.py b/venv/lib/python3.10/site-packages/accelerate/commands/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..326f37d7f93de2417e4171e5ffe91193fb97225c --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/commands/utils.py @@ -0,0 +1,123 @@ +# Copyright 2024 The HuggingFace Team. 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. + +import argparse + + +class _StoreAction(argparse.Action): + """ + Custom action that allows for `-` or `_` to be passed in for an argument. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + new_option_strings = [] + for option_string in self.option_strings: + new_option_strings.append(option_string) + if "_" in option_string[2:]: + # Add `-` version to the option string + new_option_strings.append(option_string.replace("_", "-")) + self.option_strings = new_option_strings + + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, values) + if not hasattr(namespace, "nondefault"): + namespace.nondefault = set() + namespace.nondefault.add(self.dest) + + +class _StoreConstAction(_StoreAction): + """ + Same as `argparse._StoreConstAction` but uses the custom `_StoreAction`. + """ + + def __init__(self, option_strings, dest, const, default=None, required=False, help=None): + super().__init__( + option_strings=option_strings, + dest=dest, + nargs=0, + const=const, + default=default, + required=required, + help=help, + ) + + def __call__(self, parser, namespace, values, option_string=None): + super().__call__(parser, namespace, self.const, option_string) + + +class _StoreTrueAction(_StoreConstAction): + """ + Same as `argparse._StoreTrueAction` but uses the custom `_StoreConstAction`. + """ + + def __init__( + self, + option_strings, + dest, + default=None, + required=False, + help=None, + ): + super().__init__( + option_strings=option_strings, dest=dest, const=True, default=default, required=required, help=help + ) + + +class CustomArgumentGroup(argparse._ArgumentGroup): + """ + Custom argument group that allows for the use of `-` or `_` in arguments passed and overrides the help for each + when applicable. + """ + + def _add_action(self, action): + args = vars(action) + if isinstance(action, argparse._StoreTrueAction): + action = _StoreTrueAction( + args["option_strings"], args["dest"], args["default"], args["required"], args["help"] + ) + elif isinstance(action, argparse._StoreConstAction): + action = _StoreConstAction( + args["option_strings"], + args["dest"], + args["const"], + args["default"], + args["required"], + args["help"], + ) + elif isinstance(action, argparse._StoreAction): + action = _StoreAction(**args) + action = super()._add_action(action) + return action + + +class CustomArgumentParser(argparse.ArgumentParser): + """ + Custom argument parser that allows for the use of `-` or `_` in arguments passed and overrides the help for each + when applicable. + """ + + def add_argument(self, *args, **kwargs): + if "action" in kwargs: + # Translate action -> class + if kwargs["action"] == "store_true": + kwargs["action"] = _StoreTrueAction + else: + kwargs["action"] = _StoreAction + super().add_argument(*args, **kwargs) + + def add_argument_group(self, *args, **kwargs): + group = CustomArgumentGroup(self, *args, **kwargs) + self._action_groups.append(group) + return group diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..faeec3e0ed0ecfd33f1a41a92a7046a0b2f72048 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7cce7bfc91448b93175c668e04d6eff0451e946 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aab3548b518620e89472f2e5091825f376202b2d Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1bdaaeec24bc57c5296873e44cc2488b8ebe62b3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c9cbe26c257b515f657c05e1996d517e69613972 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2020 The HuggingFace Team. 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. diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c3bc486ef902410251a6c0fecb247a42c9801d2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..083a957f2d83c8049f0292eee8f9fec04a37aaf4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7520b0bf1e4aae11ed29f8b5becc760eb7483ef Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52729838b0bad19cc293ca2604a260a90937b65e Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7cc060478d6f3917a92e4a89f1d126e651e608f2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa6c2eb4453a0b039a90e71e48d110f7d814760d Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a530a1f570d585e532f9320ce9a1e12d54c870ec Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17c32f7f7e22597140a684a86aba7f373ec0854a Binary files /dev/null and b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..6a1553898ec3d55e64822c204ddf7e705069ce8a --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py @@ -0,0 +1,269 @@ +# Copyright 2022 The HuggingFace Inc. team. 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. +import argparse +import json +import os + +import evaluate +import torch +from datasets import load_dataset +from torch.optim import AdamW +from torch.utils.data import DataLoader +from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed + +from accelerate import Accelerator, DistributedType +from accelerate.utils.deepspeed import DummyOptim, DummyScheduler + + +MAX_GPU_BATCH_SIZE = 16 +EVAL_BATCH_SIZE = 32 + + +def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"): + """ + Creates a set of `DataLoader`s for the `glue` dataset. + + Args: + accelerator (`Accelerator`): + An `Accelerator` object + batch_size (`int`, *optional*): + The batch size for the train and validation DataLoaders. + model_name (`str`, *optional*): + """ + tokenizer = AutoTokenizer.from_pretrained(model_name) + datasets = load_dataset("glue", "mrpc") + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + + # Apply the method we just defined to all the examples in all the splits of the dataset + tokenized_datasets = datasets.map( + tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False + ) + + # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the + # transformers library + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + + def collate_fn(examples): + # On TPU it's best to pad everything to the same length or training will be very slow. + if accelerator.distributed_type == DistributedType.XLA: + return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") + return tokenizer.pad(examples, padding="longest", return_tensors="pt") + + # Instantiate dataloaders. + train_dataloader = DataLoader( + tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size + ) + eval_dataloader = DataLoader( + tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE + ) + + return train_dataloader, eval_dataloader + + +def evaluation_loop(accelerator, model, eval_dataloader, metric): + model.eval() + samples_seen = 0 + for step, batch in enumerate(eval_dataloader): + # We could avoid this line since we set the accelerator with `device_placement=True`. + batch.to(accelerator.device) + with torch.no_grad(): + outputs = model(**batch) + predictions = outputs.logits.argmax(dim=-1) + # It is slightly faster to call this once, than multiple times + predictions, references = accelerator.gather( + (predictions, batch["labels"]) + ) # If we are in a multiprocess environment, the last batch has duplicates + if accelerator.use_distributed: + if step == len(eval_dataloader) - 1: + predictions = predictions[: len(eval_dataloader.dataset) - samples_seen] + references = references[: len(eval_dataloader.dataset) - samples_seen] + else: + samples_seen += references.shape[0] + metric.add_batch( + predictions=predictions, + references=references, + ) + + eval_metric = metric.compute() + return eval_metric["accuracy"] + + +def training_function(config, args): + # Initialize accelerator + accelerator = Accelerator() + + # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs + lr = config["lr"] + num_epochs = int(config["num_epochs"]) + seed = int(config["seed"]) + batch_size = int(config["batch_size"]) + model_name = args.model_name_or_path + + set_seed(seed) + train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name) + + # Instantiate the model (we build the model here so that the seed also control new weights initialization) + model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True) + + # Instantiate optimizer + optimizer_cls = ( + AdamW + if accelerator.state.deepspeed_plugin is None + or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config + else DummyOptim + ) + optimizer = optimizer_cls(params=model.parameters(), lr=lr) + + if accelerator.state.deepspeed_plugin is not None: + gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[ + "gradient_accumulation_steps" + ] + else: + gradient_accumulation_steps = 1 + max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps + + # Instantiate scheduler + if ( + accelerator.state.deepspeed_plugin is None + or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config + ): + lr_scheduler = get_linear_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=0, + num_training_steps=max_training_steps, + ) + else: + lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0) + + # Prepare everything + # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the + # prepare method. + model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, eval_dataloader, lr_scheduler + ) + + # We need to keep track of how many total steps we have iterated over + overall_step = 0 + # We also need to keep track of the stating epoch so files are named properly + starting_epoch = 0 + metric = evaluate.load("glue", "mrpc") + ending_epoch = num_epochs + + if args.partial_train_epoch is not None: + ending_epoch = args.partial_train_epoch + + if args.resume_from_checkpoint: + accelerator.load_state(args.resume_from_checkpoint) + epoch_string = args.resume_from_checkpoint.split("epoch_")[1] + state_epoch_num = "" + for char in epoch_string: + if char.isdigit(): + state_epoch_num += char + else: + break + starting_epoch = int(state_epoch_num) + 1 + accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric) + accelerator.print("resumed checkpoint performance:", accuracy) + accelerator.print("resumed checkpoint's scheduler's lr:", lr_scheduler.get_lr()[0]) + accelerator.print("resumed optimizers's lr:", optimizer.param_groups[0]["lr"]) + with open(os.path.join(args.output_dir, f"state_{starting_epoch - 1}.json")) as f: + resumed_state = json.load(f) + assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed" + assert resumed_state["lr"] == lr_scheduler.get_lr()[0], ( + "Scheduler learning rate mismatch, loading from checkpoint failed" + ) + assert resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"], ( + "Optimizer learning rate mismatch, loading from checkpoint failed" + ) + assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed" + return + + # Now we train the model + state = {} + for epoch in range(starting_epoch, ending_epoch): + model.train() + for step, batch in enumerate(train_dataloader): + outputs = model(**batch) + loss = outputs.loss + loss = loss / gradient_accumulation_steps + accelerator.backward(loss) + if step % gradient_accumulation_steps == 0: + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + overall_step += 1 + output_dir = f"epoch_{epoch}" + output_dir = os.path.join(args.output_dir, output_dir) + accelerator.save_state(output_dir) + accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric) + state["accuracy"] = accuracy + state["lr"] = lr_scheduler.get_lr()[0] + state["optimizer_lr"] = optimizer.param_groups[0]["lr"] + state["epoch"] = epoch + state["step"] = overall_step + accelerator.print(f"epoch {epoch}:", state) + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + with open(os.path.join(args.output_dir, f"state_{epoch}.json"), "w") as f: + json.dump(state, f) + accelerator.end_training() + + +def main(): + parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") + parser.add_argument( + "--model_name_or_path", + type=str, + default="bert-base-cased", + help="Path to pretrained model or model identifier from huggingface.co/models.", + required=False, + ) + parser.add_argument( + "--output_dir", + type=str, + default=".", + help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.", + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help="If the training should continue from a checkpoint folder.", + ) + parser.add_argument( + "--partial_train_epoch", + type=int, + default=None, + help="If passed, the training will stop after this number of epochs.", + ) + parser.add_argument( + "--num_epochs", + type=int, + default=2, + help="Number of train epochs.", + ) + args = parser.parse_args() + config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} + + training_function(config, args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py new file mode 100644 index 0000000000000000000000000000000000000000..3729ecf4c72190bf865d620b6941206ab904818c --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py @@ -0,0 +1,332 @@ +# Copyright 2024 The HuggingFace Inc. team. 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. + +""" +Test script for verifying multiple models can be utilized with Accelerate + DeepSpeed: + +Scenario 1: One model is training, another model is being used for inference/logits to impact training in some form. +Scenario 2: Two models are training simultaneously, which means two optimizers, etc. +""" + +import argparse +from pathlib import Path + +import evaluate +import torch +from datasets import load_dataset +from torch.optim import AdamW +from torch.utils.data import DataLoader +from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup + +from accelerate import Accelerator, DeepSpeedPlugin, DistributedType +from accelerate.state import AcceleratorState +from accelerate.utils.deepspeed import get_active_deepspeed_plugin + + +MAX_GPU_BATCH_SIZE = 16 +EVAL_BATCH_SIZE = 32 + + +class NoiseModel(torch.nn.Module): + def __init__(self, noise_factor=0.1): + super().__init__() + self.noise_factor = torch.nn.Parameter(torch.tensor(noise_factor, dtype=torch.float32)) + + def forward(self, loss): + return loss * self.noise_factor + + +def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"): + """ + Creates a set of `DataLoader`s for the `glue` dataset. + + Args: + accelerator (`Accelerator`): + An `Accelerator` object + batch_size (`int`, *optional*): + The batch size for the train and validation DataLoaders. + model_name (`str`, *optional*): + """ + tokenizer = AutoTokenizer.from_pretrained(model_name) + datasets = load_dataset("glue", "mrpc") + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + + # Apply the method we just defined to all the examples in all the splits of the dataset + tokenized_datasets = datasets.map( + tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False + ) + + # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the + # transformers library + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + + def collate_fn(examples): + # On TPU it's best to pad everything to the same length or training will be very slow. + if accelerator.distributed_type == DistributedType.XLA: + return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") + return tokenizer.pad(examples, padding="longest", return_tensors="pt") + + # Instantiate dataloaders. + train_dataloader = DataLoader( + tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size + ) + eval_dataloader = DataLoader( + tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE + ) + + return train_dataloader, eval_dataloader + + +test_file_path = __file__ +path = Path(test_file_path).resolve() +test_file_dir_str = str(path.parent.parent.parent.parent.parent.parent) + +# Create our DS plugins +# We use custom schedulers and optimizers, hence `model_only` +ds_config_file = dict( + zero2=f"{test_file_dir_str}/tests/deepspeed/ds_config_zero2_model_only.json", + zero3=f"{test_file_dir_str}/tests/deepspeed/ds_config_zero3_model_only.json", +) + + +def single_model_training(config, args): + # Training a single model, we have a `noise` model that is untrainable used to inject some noise into the training process + num_epochs = config["num_epochs"] + zero2_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero2"]) + zero3_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero3"]) + + deepspeed_plugins = {"training": zero2_plugin, "inference": zero3_plugin} + + # Initialize accelerator + accelerator = Accelerator( + deepspeed_plugins=deepspeed_plugins, + mixed_precision="bf16", + ) + + # Initialize model under zero2 plugin + assert get_active_deepspeed_plugin(accelerator.state) is zero2_plugin + train_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path) + train_dataloader, eval_dataloader = get_dataloaders( + accelerator, batch_size=config["batch_size"], model_name=args.model_name_or_path + ) + max_training_steps = len(train_dataloader) * config["num_epochs"] + optimizer = AdamW(train_model.parameters(), lr=config["lr"]) + lr_scheduler = get_linear_schedule_with_warmup( + optimizer, num_warmup_steps=0, num_training_steps=max_training_steps + ) + + train_dataloader, eval_dataloader, train_model, optimizer, lr_scheduler = accelerator.prepare( + train_dataloader, eval_dataloader, train_model, optimizer, lr_scheduler + ) + + # Now prepare the model under zero3 plugin + accelerator.state.select_deepspeed_plugin("inference") + assert get_active_deepspeed_plugin(accelerator.state) is zero3_plugin + inference_model = NoiseModel() + inference_model = accelerator.prepare(inference_model) + inference_model.eval() + + # Run training loop + accelerator.state.select_deepspeed_plugin("training") + # We also need to keep track of the stating epoch so files are named properly + starting_epoch = 0 + + # Now we train the model + best_performance = 0 + metric = evaluate.load("glue", "mrpc") + performance_metric = {} + for epoch in range(starting_epoch, num_epochs): + train_model.train() + inference_model.train() + for step, batch in enumerate(train_dataloader): + with accelerator.accumulate(train_model): + outputs_1 = train_model(**batch) + with torch.no_grad(): + outputs_2 = inference_model(outputs_1.loss) + # Combine the losses + loss = outputs_1.loss + outputs_2 + accelerator.backward(loss) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + train_model.eval() + for step, batch in enumerate(eval_dataloader): + with torch.no_grad(): + outputs = train_model(**batch) + predictions = outputs.logits.argmax(dim=-1) + # It is slightly faster to call this once, than multiple times + predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"])) + metric.add_batch( + predictions=predictions, + references=references, + ) + + eval_metric = metric.compute() + # Use accelerator.print to print only on the main process. + accelerator.print(f"epoch {epoch}:", eval_metric) + performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"] + + if best_performance < eval_metric["accuracy"]: + best_performance = eval_metric["accuracy"] + assert best_performance > performance_metric["epoch-0"] + + +def multiple_model_training(config, args): + # This will essentially be like a k-fold model, but one model is Zero-2 and another model is Zero-3 + num_epochs = config["num_epochs"] + zero2_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero2"]) + zero3_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero3"]) + + deepspeed_plugins = {"zero2": zero2_plugin, "zero3": zero3_plugin} + + # Initialize accelerator + zero2_accelerator = Accelerator( + deepspeed_plugins=deepspeed_plugins, + mixed_precision="bf16", + ) + + # Since an `AcceleratorState` has already been made, we can just reuse it here + zero3_accelerator = Accelerator() + + # Initialize model under zero2 plugin + assert get_active_deepspeed_plugin(zero2_accelerator.state) is zero2_plugin + zero2_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path) + train_dataloader, eval_dataloader = get_dataloaders( + zero2_accelerator, batch_size=config["batch_size"], model_name=args.model_name_or_path + ) + max_training_steps = len(train_dataloader) * config["num_epochs"] + zero2_optimizer = AdamW(zero2_model.parameters(), lr=config["lr"]) + zero2_lr_scheduler = get_linear_schedule_with_warmup( + zero2_optimizer, num_warmup_steps=0, num_training_steps=max_training_steps + ) + + train_dataloader, eval_dataloader, zero2_model, zero2_optimizer, zero2_lr_scheduler = zero2_accelerator.prepare( + train_dataloader, eval_dataloader, zero2_model, zero2_optimizer, zero2_lr_scheduler + ) + assert zero2_accelerator.deepspeed_engine_wrapped.engine is zero2_model + + # now do Zero3 + zero3_accelerator.state.select_deepspeed_plugin("zero3") + zero3_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = zero2_plugin.deepspeed_config[ + "train_micro_batch_size_per_gpu" + ] + assert get_active_deepspeed_plugin(zero3_accelerator.state) is zero3_plugin + zero3_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path) + zero3_optimizer = AdamW(zero3_model.parameters(), lr=config["lr"]) + zero3_lr_scheduler = get_linear_schedule_with_warmup( + zero3_optimizer, num_warmup_steps=0, num_training_steps=max_training_steps + ) + zero3_model, zero3_optimizer, zero3_lr_scheduler = zero3_accelerator.prepare( + zero3_model, zero3_optimizer, zero3_lr_scheduler + ) + assert zero3_accelerator.deepspeed_engine_wrapped.engine is zero3_model + + # Run training loop + starting_epoch = 0 + + # Now we train the model + best_performance_a = 0 + best_performance_b = 0 + metric_a = evaluate.load("glue", "mrpc") + metric_b = evaluate.load("glue", "mrpc") + performance_metric_a = {} + performance_metric_b = {} + for epoch in range(starting_epoch, num_epochs): + zero2_model.train() + zero3_model.train() + for step, batch in enumerate(train_dataloader): + with zero2_accelerator.accumulate(zero2_model, zero3_model): + outputs_1 = zero2_model(**batch) + zero2_accelerator.backward(outputs_1.loss) + zero2_optimizer.step() + zero2_lr_scheduler.step() + zero2_optimizer.zero_grad() + outputs_2 = zero3_model(**batch) + zero3_accelerator.backward(outputs_2.loss) + zero3_optimizer.step() + zero3_lr_scheduler.step() + zero3_optimizer.zero_grad() + + zero2_model.eval() + zero3_model.eval() + for step, batch in enumerate(eval_dataloader): + with torch.no_grad(): + logits_a = zero2_model(**batch).logits + logits_b = zero3_model(**batch).logits + # Combine the logits from both models + predictions_a = logits_a.argmax(dim=-1) + predictions_b = logits_b.argmax(dim=-1) + # It is slightly faster to call this once, than multiple times + predictions_a, predictions_b, references = zero2_accelerator.gather_for_metrics( + (predictions_a, predictions_b, batch["labels"]) + ) + metric_a.add_batch( + predictions=predictions_a, + references=references, + ) + metric_b.add_batch( + predictions=predictions_b, + references=references, + ) + + eval_metric_a = metric_a.compute() + eval_metric_b = metric_b.compute() + # Use accelerator.print to print only on the main process. + zero2_accelerator.print(f"epoch {epoch}:", eval_metric_a, eval_metric_b) + performance_metric_a[f"epoch-{epoch}"] = eval_metric_a["accuracy"] + performance_metric_b[f"epoch-{epoch}"] = eval_metric_b["accuracy"] + + if best_performance_a < eval_metric_a["accuracy"]: + best_performance_a = eval_metric_a["accuracy"] + if best_performance_b < eval_metric_b["accuracy"]: + best_performance_b = eval_metric_b["accuracy"] + assert best_performance_a > performance_metric_a["epoch-0"] + assert best_performance_b > performance_metric_b["epoch-0"] + + +def main(): + parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") + parser.add_argument( + "--model_name_or_path", + type=str, + default="bert-base-cased", + help="Path to pretrained model or model identifier from huggingface.co/models.", + required=False, + ) + parser.add_argument( + "--performance_lower_bound", + type=float, + default=None, + help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.", + ) + parser.add_argument( + "--num_epochs", + type=int, + default=2, + help="Number of train epochs.", + ) + args = parser.parse_args() + config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} + single_model_training(config, args) + AcceleratorState._reset_state(True) + multiple_model_training(config, args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..d1bfe351509148ebc48067584e9d61b93e7210a6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py @@ -0,0 +1,307 @@ +# Copyright 2022 The HuggingFace Team. 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. + +import logging +import math +import os +from copy import deepcopy + +import datasets +import evaluate +import torch +import transformers +from datasets import load_dataset +from torch.utils.data import DataLoader, IterableDataset +from transformers import AutoModelForSequenceClassification, AutoTokenizer + +from accelerate import Accelerator, DataLoaderConfiguration, DistributedType +from accelerate.data_loader import DataLoaderDispatcher +from accelerate.test_utils import RegressionDataset, RegressionModel, torch_device +from accelerate.utils import is_torch_xla_available, set_seed + + +os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" + + +class ListHandler(logging.Handler): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.logs = [] + + def emit(self, record): + self.logs.append(record) + + +def get_basic_setup(accelerator, num_samples=82, batch_size=16): + "Returns everything needed to perform basic training" + set_seed(42) + model = RegressionModel() + ddp_model = deepcopy(model) + dset = RegressionDataset(length=num_samples) + dataloader = DataLoader(dset, batch_size=batch_size) + model.to(accelerator.device) + ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader) + return model, ddp_model, dataloader + + +def get_dataloader(accelerator: Accelerator, use_longest=False): + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/mrpc-bert-base-cased") + dataset = load_dataset("glue", "mrpc", split="validation") + + def tokenize_function(examples): + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + + with accelerator.main_process_first(): + tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], + ) + + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + + def collate_fn(examples): + if use_longest: + return tokenizer.pad(examples, padding="longest", return_tensors="pt") + return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") + + return DataLoader(tokenized_datasets, shuffle=False, collate_fn=collate_fn, batch_size=16) + + +def get_mrpc_setup(dispatch_batches, split_batches): + dataloader_config = DataLoaderConfiguration(dispatch_batches=dispatch_batches, split_batches=split_batches) + accelerator = Accelerator(dataloader_config=dataloader_config) + dataloader = get_dataloader(accelerator, not dispatch_batches) + model = AutoModelForSequenceClassification.from_pretrained( + "hf-internal-testing/mrpc-bert-base-cased", return_dict=True + ) + ddp_model, ddp_dataloader = accelerator.prepare(model, dataloader) + return { + "ddp": [ddp_model, ddp_dataloader, torch_device], + "no": [model, dataloader, accelerator.device], + }, accelerator + + +def generate_predictions(model, dataloader, accelerator): + logits_and_targets = [] + for batch in dataloader: + input, target = batch.values() + with torch.no_grad(): + logit = model(input) + logit, target = accelerator.gather_for_metrics((logit, target)) + logits_and_targets.append((logit, target)) + logits, targs = [], [] + for logit, targ in logits_and_targets: + logits.append(logit) + targs.append(targ) + logits, targs = torch.cat(logits), torch.cat(targs) + return logits, targs + + +def test_torch_metrics( + accelerator: Accelerator, num_samples=82, dispatch_batches=False, split_batches=False, batch_size=16 +): + _, ddp_model, dataloader = get_basic_setup(accelerator, num_samples, batch_size) + logits, _ = generate_predictions(ddp_model, dataloader, accelerator) + assert len(logits) == num_samples, ( + f"Unexpected number of inputs:\n Expected: {num_samples}\n Actual: {len(logits)}" + ) + + +def test_mrpc(dispatch_batches: bool = False, split_batches: bool = False): + metric = evaluate.load("glue", "mrpc") + setup, accelerator = get_mrpc_setup(dispatch_batches, split_batches) + # First do baseline + model, dataloader, device = setup["no"] + model.to(device) + model.eval() + for batch in dataloader: + batch.to(device) + with torch.inference_mode(): + outputs = model(**batch) + preds = outputs.logits.argmax(dim=-1) + metric.add_batch(predictions=preds, references=batch["labels"]) + baseline = metric.compute() + + # Then do distributed + model, dataloader, device = setup["ddp"] + model.eval() + for batch in dataloader: + with torch.inference_mode(): + outputs = model(**batch) + preds = outputs.logits.argmax(dim=-1) + references = batch["labels"] + preds, references = accelerator.gather_for_metrics((preds, references)) + metric.add_batch(predictions=preds, references=references) + distributed = metric.compute() + + for key in "accuracy f1".split(): + assert math.isclose(baseline[key], distributed[key]), ( + f"Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n" + ) + + +def test_gather_for_metrics_with_non_tensor_objects_iterable_dataset(): + class DummyIterableDataset(IterableDataset): + def __init__(self, data): + self.data = data + + def __len__(self): + return len(self.data) + + def __iter__(self): + yield from self.data + + iterable_dataset = DummyIterableDataset([n for n in range(30)]) + dataloader = DataLoader(iterable_dataset, batch_size=4) + accelerator = Accelerator() + prepared_dataloader = accelerator.prepare(dataloader) + + if accelerator.is_main_process: + logger = logging.root.manager.loggerDict["accelerate.accelerator"] + list_handler = ListHandler() + logger.addHandler(list_handler) + + batches_for_metrics = [] + for batch in prepared_dataloader: + batches_for_metrics.append(accelerator.gather_for_metrics(batch)) + + assert torch.cat(batches_for_metrics).size(0) == 30 + + if accelerator.is_main_process: + assert len(list_handler.logs) == 0 + logger.removeHandler(list_handler) + + +def test_gather_for_metrics_with_iterable_dataset(): + class DummyIterableDataset(IterableDataset): + def __init__(self, data): + self.data = data + + def __len__(self): + return len(self.data) + + def __iter__(self): + yield from self.data + + iterable_dataset = DummyIterableDataset(torch.as_tensor(range(30))) + dataloader = DataLoader(iterable_dataset, batch_size=4) + + accelerator = Accelerator() + prepared_dataloader = accelerator.prepare(dataloader) + + assert isinstance(prepared_dataloader, DataLoaderDispatcher) + + if accelerator.is_main_process: + logger = logging.root.manager.loggerDict["accelerate.accelerator"] + list_handler = ListHandler() + logger.addHandler(list_handler) + + batches_for_metrics = [] + for batch in prepared_dataloader: + batches_for_metrics.append(accelerator.gather_for_metrics(batch)) + + assert torch.cat(batches_for_metrics).size(0) == 30 + + if accelerator.is_main_process: + assert len(list_handler.logs) == 0 + + logger.removeHandler(list_handler) + + +def test_gather_for_metrics_drop_last(): + accelerator = Accelerator() + per_device_batch_size = 5 + num_items = (10 * accelerator.num_processes) + 1 + dataloader = DataLoader(range(num_items), batch_size=per_device_batch_size, drop_last=True) + dataloader = accelerator.prepare(dataloader) + + iterator = iter(dataloader) + next(iterator) # Skip first batch tensor([0, 1, 2, 3, 4], device='cuda:0') + batch = next(iterator) + gathered_items = accelerator.gather_for_metrics(batch) + + # Should return a full set of complete batches from each GPU + num_expected_items = per_device_batch_size * accelerator.num_processes + assert gathered_items.size(0) == (num_expected_items), ( + f"Expected number of items: {num_expected_items}, Actual: {gathered_items.size(0)}" + ) + + +def main(): + dataloader_config = DataLoaderConfiguration(split_batches=False, dispatch_batches=False) + accelerator = Accelerator(dataloader_config=dataloader_config) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + # TorchXLA does not support batch dispatching. 'put_on_device' is always False for + # TorchXLA, which can cause a value error in 'prepare_data_loader' function. + dispatch_batches_options = [False] if accelerator.state.distributed_type == DistributedType.XLA else [True, False] + + # Temporarily close this test for TorchXLA due to the 'Cannot set version_counter for + # inference tensor' error in inference mode. Reopen it after TorchXLA fixes this bug. + # These are a bit slower so they should only be ran on the GPU or TPU + if accelerator.device.type != "cpu" and not is_torch_xla_available(): + if accelerator.is_local_main_process: + print("**Testing gather_for_metrics**") + for split_batches in [True, False]: + for dispatch_batches in dispatch_batches_options: + if accelerator.is_local_main_process: + print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`") + test_mrpc(dispatch_batches, split_batches) + accelerator.state._reset_state() + print("test_gather_for_metrics_with_iterable_dataset") + test_gather_for_metrics_with_iterable_dataset() + print("test gather_for_metrics_with_non_tensor_objects_iterable_dataset") + test_gather_for_metrics_with_non_tensor_objects_iterable_dataset() + + # MpDeviceLoader in TorchXLA is an asynchronous loader that preloads several batches into cache. + # This can cause the 'end_of_dataloader' of DataLoaderStateMixin to be set earlier than intended. + # Skip this test when TorchXLA is enabled. + if accelerator.state.distributed_type != DistributedType.XLA: + if accelerator.is_local_main_process: + print("**Test torch metrics**") + for split_batches in [True, False]: + for dispatch_batches in dispatch_batches_options: + dataloader_config = DataLoaderConfiguration( + split_batches=split_batches, dispatch_batches=dispatch_batches + ) + accelerator = Accelerator(dataloader_config=dataloader_config) + if accelerator.is_local_main_process: + print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99") + test_torch_metrics(accelerator, 99) + accelerator.state._reset_state() + if accelerator.is_local_main_process: + print("**Test last batch is not dropped when perfectly divisible**") + accelerator = Accelerator() + test_torch_metrics(accelerator, 512) + accelerator.state._reset_state() + if accelerator.is_local_main_process: + print("**Test that `drop_last` is taken into account**") + test_gather_for_metrics_drop_last() + accelerator.end_training() + accelerator.state._reset_state() + + +def _mp_fn(index): + # For xla_spawn (TPUs) + main() + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py new file mode 100644 index 0000000000000000000000000000000000000000..723e5497656020516bf072cf4112f61f59c2e5cb --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py @@ -0,0 +1,314 @@ +# Copyright 2022 The HuggingFace Inc. team. 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. +import argparse +import gc +import json +import os + +import torch +from datasets import load_dataset +from torch.optim import AdamW +from torch.utils.data import DataLoader +from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed + +from accelerate import Accelerator, DistributedType +from accelerate.utils import ( + is_hpu_available, + is_mlu_available, + is_musa_available, + is_npu_available, + is_sdaa_available, + is_xpu_available, +) +from accelerate.utils.deepspeed import DummyOptim, DummyScheduler + + +MAX_GPU_BATCH_SIZE = 16 +EVAL_BATCH_SIZE = 32 + + +# Converting Bytes to Megabytes +def b2mb(x): + return int(x / 2**20) + + +# This context manager is used to track the peak memory usage of the process +class TorchTracemalloc: + def __enter__(self): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.reset_max_memory_allocated() # reset the peak gauge to zero + self.begin = torch.cuda.memory_allocated() + elif is_mlu_available(): + torch.mlu.empty_cache() + torch.mlu.reset_max_memory_allocated() # reset the peak gauge to zero + self.begin = torch.mlu.memory_allocated() + elif is_sdaa_available(): + torch.sdaa.empty_cache() + torch.sdaa.reset_max_memory_allocated() # reset the peak gauge to zero + self.begin = torch.sdaa.memory_allocated() + elif is_musa_available(): + torch.musa.empty_cache() + torch.musa.reset_max_memory_allocated() # reset the peak gauge to zero + self.begin = torch.musa.memory_allocated() + elif is_npu_available(): + torch.npu.empty_cache() + torch.npu.reset_max_memory_allocated() # reset the peak gauge to zero + self.begin = torch.npu.memory_allocated() + elif is_xpu_available(): + torch.xpu.empty_cache() + torch.xpu.reset_max_memory_allocated() # reset the peak gauge to zero + self.begin = torch.xpu.memory_allocated() + elif is_hpu_available(): + # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process + torch.hpu.reset_peak_memory_stats() # reset the peak gauge to zero + self.begin = torch.hpu.memory_allocated() + return self + + def __exit__(self, *exc): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self.end = torch.cuda.memory_allocated() + self.peak = torch.cuda.max_memory_allocated() + elif is_mlu_available(): + torch.mlu.empty_cache() + self.end = torch.mlu.memory_allocated() + self.begin = torch.mlu.max_memory_allocated() + elif is_sdaa_available(): + torch.sdaa.empty_cache() + self.end = torch.sdaa.memory_allocated() + self.begin = torch.sdaa.max_memory_allocated() + elif is_musa_available(): + torch.musa.empty_cache() + self.end = torch.musa.memory_allocated() + self.begin = torch.musa.max_memory_allocated() + elif is_npu_available(): + torch.npu.empty_cache() + self.end = torch.npu.memory_allocated() + self.peak = torch.npu.max_memory_allocated() + elif is_xpu_available(): + torch.xpu.empty_cache() + self.end = torch.xpu.memory_allocated() + self.peak = torch.xpu.max_memory_allocated() + elif is_hpu_available(): + # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process + self.end = torch.hpu.memory_allocated() + self.peak = torch.hpu.max_memory_allocated() + self.used = b2mb(self.end - self.begin) + self.peaked = b2mb(self.peak - self.begin) + # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}") + + +def get_dataloaders( + accelerator: Accelerator, + batch_size: int = 16, + model_name: str = "bert-base-cased", + n_train: int = 320, + n_val: int = 160, +): + """ + Creates a set of `DataLoader`s for the `glue` dataset. + + Args: + accelerator (`Accelerator`): + An `Accelerator` object + batch_size (`int`, *optional*): + The batch size for the train and validation DataLoaders. + model_name (`str`, *optional*): + The name of the model to use. + n_train (`int`, *optional*): + The number of training examples to use. + n_val (`int`, *optional*): + The number of validation examples to use. + """ + tokenizer = AutoTokenizer.from_pretrained(model_name) + datasets = load_dataset( + "glue", "mrpc", split={"train": f"train[:{n_train}]", "validation": f"validation[:{n_val}]"} + ) + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + + # Apply the method we just defined to all the examples in all the splits of the dataset + tokenized_datasets = datasets.map( + tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False + ) + + # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the + # transformers library + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + + def collate_fn(examples): + # On TPU it's best to pad everything to the same length or training will be very slow. + if accelerator.distributed_type == DistributedType.XLA: + return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") + return tokenizer.pad(examples, padding="longest", return_tensors="pt") + + # Instantiate dataloaders. + train_dataloader = DataLoader( + tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size + ) + eval_dataloader = DataLoader( + tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE + ) + + return train_dataloader, eval_dataloader + + +def training_function(config, args): + # Initialize accelerator + accelerator = Accelerator() + + # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs + lr = config["lr"] + num_epochs = int(config["num_epochs"]) + seed = int(config["seed"]) + batch_size = int(config["batch_size"]) + model_name = args.model_name_or_path + + set_seed(seed) + train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name, args.n_train, args.n_val) + + # Instantiate the model (we build the model here so that the seed also control new weights initialization) + model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True) + + # Instantiate optimizer + optimizer_cls = ( + AdamW + if accelerator.state.deepspeed_plugin is None + or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config + else DummyOptim + ) + optimizer = optimizer_cls(params=model.parameters(), lr=lr) + + if accelerator.state.deepspeed_plugin is not None: + gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[ + "gradient_accumulation_steps" + ] + else: + gradient_accumulation_steps = 1 + max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps + + # Instantiate scheduler + if ( + accelerator.state.deepspeed_plugin is None + or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config + ): + lr_scheduler = get_linear_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=0, + num_training_steps=max_training_steps, + ) + else: + lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0) + + # Prepare everything + # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the + # prepare method. + model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, eval_dataloader, lr_scheduler + ) + + # We need to keep track of how many total steps we have iterated over + overall_step = 0 + # We also need to keep track of the stating epoch so files are named properly + starting_epoch = 0 + + # Now we train the model + train_total_peak_memory = {} + for epoch in range(starting_epoch, num_epochs): + with TorchTracemalloc() as tracemalloc: + model.train() + for step, batch in enumerate(train_dataloader): + outputs = model(**batch) + loss = outputs.loss + loss = loss / gradient_accumulation_steps + accelerator.backward(loss) + if step % gradient_accumulation_steps == 0: + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + overall_step += 1 + + # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage + accelerator.print(f"Memory before entering the train : {b2mb(tracemalloc.begin)}") + accelerator.print(f"Memory consumed at the end of the train (end-begin): {tracemalloc.used}") + accelerator.print(f"Peak Memory consumed during the train (max-begin): {tracemalloc.peaked}") + accelerator.print( + f"Total Peak Memory consumed during the train (max): {tracemalloc.peaked + b2mb(tracemalloc.begin)}" + ) + train_total_peak_memory[f"epoch-{epoch}"] = tracemalloc.peaked + b2mb(tracemalloc.begin) + if args.peak_memory_upper_bound is not None: + assert train_total_peak_memory[f"epoch-{epoch}"] <= args.peak_memory_upper_bound, ( + "Peak memory usage exceeded the upper bound" + ) + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + with open(os.path.join(args.output_dir, "peak_memory_utilization.json"), "w") as f: + json.dump(train_total_peak_memory, f) + accelerator.end_training() + + +def main(): + parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") + parser.add_argument( + "--model_name_or_path", + type=str, + default="bert-base-cased", + help="Path to pretrained model or model identifier from huggingface.co/models.", + required=False, + ) + parser.add_argument( + "--output_dir", + type=str, + default=".", + help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.", + ) + parser.add_argument( + "--peak_memory_upper_bound", + type=float, + default=None, + help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.", + ) + parser.add_argument( + "--n_train", + type=int, + default=320, + help="Number of training examples to use.", + ) + parser.add_argument( + "--n_val", + type=int, + default=160, + help="Number of validation examples to use.", + ) + parser.add_argument( + "--num_epochs", + type=int, + default=1, + help="Number of train epochs.", + ) + args = parser.parse_args() + config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} + training_function(config, args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py new file mode 100644 index 0000000000000000000000000000000000000000..d4396ec4d6f7ce23db3678930c2de22af39681ca --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py @@ -0,0 +1,298 @@ +# Copyright 2022 The HuggingFace Inc. team. 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. +import argparse +import json +import os +from contextlib import nullcontext +from pathlib import Path + +import evaluate +import torch +from datasets import load_dataset +from torch.optim import AdamW +from torch.utils.data import DataLoader +from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup + +from accelerate import Accelerator, DistributedType +from accelerate.utils import SAFE_WEIGHTS_NAME, TorchTensorParallelPlugin, set_seed +from accelerate.utils.deepspeed import DummyOptim, DummyScheduler + + +MAX_GPU_BATCH_SIZE = 16 +EVAL_BATCH_SIZE = 32 + + +def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"): + """ + Creates a set of `DataLoader`s for the `glue` dataset. + + Args: + accelerator (`Accelerator`): + An `Accelerator` object + batch_size (`int`, *optional*): + The batch size for the train and validation DataLoaders. + model_name (`str`, *optional*): + """ + tokenizer = AutoTokenizer.from_pretrained(model_name) + + datasets = load_dataset("glue", "mrpc") + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + + # Apply the method we just defined to all the examples in all the splits of the dataset + tokenized_datasets = datasets.map( + tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False + ) + + # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the + # transformers library + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + + def collate_fn(examples): + # On TPU it's best to pad everything to the same length or training will be very slow. + if accelerator.distributed_type == DistributedType.XLA: + return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") + return tokenizer.pad(examples, padding="longest", return_tensors="pt") + + # Instantiate dataloaders. + train_dataloader = DataLoader( + tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size + ) + eval_dataloader = DataLoader( + tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE + ) + + return train_dataloader, eval_dataloader + + +def training_function(config, args): + accelerator_kwargs = {} + # need this for DeepSpeed tests as `args.tp_size` would be None and `torch.distributed.init_device_mesh` would fail + if args.tp_size is not None: + accelerator_kwargs["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=args.tp_size) + + # Initialize accelerator + accelerator = Accelerator(**accelerator_kwargs) + + # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs + lr = config["lr"] + num_epochs = int(config["num_epochs"]) + seed = int(config["seed"]) + batch_size = int(config["batch_size"]) + model_name = args.model_name_or_path + + set_seed(seed) + train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name) + + # Add TP related kwargs if provided + model_kwargs = {} + if args.tp_plan is not None: + model_kwargs["tp_plan"] = args.tp_plan + if args.tp_size is not None: + model_kwargs["tp_size"] = args.tp_size + + # Instantiate the model (we build the model here so that the seed also control new weights initialization) + model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True, **model_kwargs) + + if args.add_pad_token: + if model.config.pad_token_id is None: + model.config.pad_token_id = 0 + + # Instantiate optimizer + optimizer_cls = ( + AdamW + if accelerator.state.deepspeed_plugin is None + or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config + else DummyOptim + ) + optimizer = optimizer_cls(params=model.parameters(), lr=lr) + + max_training_steps = len(train_dataloader) * num_epochs + + # Instantiate scheduler + linear_decay_scheduler = False + if ( + accelerator.state.deepspeed_plugin is None + or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config + ): + lr_scheduler = get_linear_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=0, + num_training_steps=max_training_steps, + ) + linear_decay_scheduler = True + else: + lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0) + + # Prepare everything + # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the + # prepare method. + model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, eval_dataloader, lr_scheduler + ) + + # We also need to keep track of the stating epoch so files are named properly + starting_epoch = 0 + + # Now we train the model + metric = evaluate.load("glue", "mrpc") + best_performance = 0 + performance_metric = {} + expected_lr_after_first_optim_step = lr * ( + 1 - 1 / (max_training_steps / accelerator.num_processes / accelerator.gradient_accumulation_steps) + ) + lr_scheduler_check_completed = False + for epoch in range(starting_epoch, num_epochs): + model.train() + for step, batch in enumerate(train_dataloader): + with accelerator.accumulate(model): + outputs = model(**batch) + loss = outputs.loss + accelerator.backward(loss) + context = nullcontext + if args.tp_plan is not None: + from torch.distributed._tensor.experimental import implicit_replication + + context = implicit_replication + with context(): + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # assert the learning rate after first optimizer step + if ( + accelerator.sync_gradients + and not lr_scheduler_check_completed + and linear_decay_scheduler + and accelerator.state.mixed_precision == "no" + ): + assert lr_scheduler.get_last_lr()[0] == expected_lr_after_first_optim_step, ( + f"Wrong lr found at second step, expected {expected_lr_after_first_optim_step}, got {lr_scheduler.get_last_lr()[0]}" + ) + lr_scheduler_check_completed = True + + model.eval() + samples_seen = 0 + for step, batch in enumerate(eval_dataloader): + # We could avoid this line since we set the accelerator with `device_placement=True`. + batch.to(accelerator.device) + with torch.no_grad(): + outputs = model(**batch) + predictions = outputs.logits.argmax(dim=-1) + # It is slightly faster to call this once, than multiple times + predictions, references = accelerator.gather( + (predictions, batch["labels"]) + ) # If we are in a multiprocess environment, the last batch has duplicates + if accelerator.use_distributed: + if step == len(eval_dataloader) - 1: + predictions = predictions[: len(eval_dataloader.dataset) - samples_seen] + references = references[: len(eval_dataloader.dataset) - samples_seen] + else: + samples_seen += references.shape[0] + metric.add_batch( + predictions=predictions, + references=references, + ) + + eval_metric = metric.compute() + # Use accelerator.print to print only on the main process. + accelerator.print(f"epoch {epoch}:", eval_metric) + performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"] + + if best_performance < eval_metric["accuracy"]: + best_performance = eval_metric["accuracy"] + + # check that the LR is 0 + if linear_decay_scheduler and accelerator.state.mixed_precision == "no": + assert lr_scheduler.get_last_lr()[0] == 0, ( + f"Wrong lr found at last step, expected 0, got {lr_scheduler.get_last_lr()[0]}" + ) + + if args.performance_lower_bound is not None: + assert args.performance_lower_bound <= best_performance, ( + f"Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}" + ) + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + with open(os.path.join(args.output_dir, "all_results.json"), "w") as f: + json.dump(performance_metric, f) + + # TODO: skip saving of the model test for TP until the feature lands + if args.tp_plan is None: + # Finally try saving the model + accelerator.save_model(model, args.output_dir) + accelerator.wait_for_everyone() + if args.tp_plan is None: + assert Path(args.output_dir, SAFE_WEIGHTS_NAME).exists(), ( + "Model was not saved when calling `Accelerator.save_model`" + ) + accelerator.end_training() + + +def main(): + parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") + parser.add_argument( + "--model_name_or_path", + type=str, + default="bert-base-cased", + help="Path to pretrained model or model identifier from huggingface.co/models.", + required=False, + ) + parser.add_argument( + "--output_dir", + type=str, + default=".", + help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.", + ) + parser.add_argument( + "--performance_lower_bound", + type=float, + default=None, + help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.", + ) + parser.add_argument( + "--num_epochs", + type=int, + default=3, + help="Number of train epochs.", + ) + parser.add_argument( + "--add_pad_token", + type=bool, + default=False, + help="To add pad token if not exists.", + ) + parser.add_argument( + "--tp_plan", + type=str, + default=None, + help="pass 'auto' to use TP", + ) + parser.add_argument( + "--tp_size", + type=int, + default=None, + help="TP size to be used to shard the model", + ) + args = parser.parse_args() + config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} + training_function(config, args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py new file mode 100644 index 0000000000000000000000000000000000000000..1dbd86c46b4a0c12df8ea4d736c7cd1e03f81813 --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py @@ -0,0 +1,117 @@ +# Copyright 2024 The HuggingFace Inc. team. 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. +import torch +from transformers import ( + BertConfig, + BertForMaskedLM, + GPT2Config, + GPT2ForSequenceClassification, +) + +from accelerate import PartialState +from accelerate.inference import prepare_pippy +from accelerate.test_utils import torch_device +from accelerate.utils import DistributedType, set_seed + + +model_to_config = { + "bert": (BertForMaskedLM, BertConfig, 512), + "gpt2": (GPT2ForSequenceClassification, GPT2Config, 1024), +} + + +def get_model_and_data_for_text(model_name, device, num_processes: int = 2): + initializer, config, seq_len = model_to_config[model_name] + config_args = {} + # Eventually needed for batch inference tests on gpt-2 when bs != 1 + # if model_name == "gpt2": + # config_args["pad_token_id"] = 0 + model_config = config(**config_args) + model = initializer(model_config) + kwargs = dict(low=0, high=model_config.vocab_size, device=device, dtype=torch.int64, requires_grad=False) + trace_input = torch.randint(size=(1, seq_len), **kwargs) + inference_inputs = torch.randint(size=(num_processes, seq_len), **kwargs) + return model, trace_input, inference_inputs + + +def test_bert(batch_size: int = 2): + set_seed(42) + state = PartialState() + model, trace_input, inference_inputs = get_model_and_data_for_text("bert", "cpu", batch_size) + model = prepare_pippy(model, example_args=(trace_input,), no_split_module_classes=model._no_split_modules) + # For inference args need to be a tuple + inputs = inference_inputs.to(torch_device) + with torch.no_grad(): + output = model(inputs) + # Zach: Check that we just grab the real outputs we need at the end + if not state.is_last_process: + assert output is None, "Output was not generated on just the last process!" + else: + assert output is not None, "Output was not generated in the last process!" + + +def test_gpt2(batch_size: int = 2): + set_seed(42) + state = PartialState() + model, trace_input, inference_inputs = get_model_and_data_for_text("gpt2", "cpu", batch_size) + model = prepare_pippy(model, example_args=(trace_input,), no_split_module_classes=model._no_split_modules) + # For inference args need to be a tuple + inputs = inference_inputs.to(torch_device) + with torch.no_grad(): + output = model(inputs) + # Zach: Check that we just grab the real outputs we need at the end + if not state.is_last_process: + assert output is None, "Output was not generated on just the last process!" + else: + assert output is not None, "Output was not generated in the last process!" + + +# Currently disabled, enable again once PyTorch pippy interface can trace a resnet34 +# def test_resnet(batch_size: int = 2): +# set_seed(42) +# state = PartialState() +# model = resnet34() +# input_tensor = torch.rand(1, 3, 224, 224) +# model = prepare_pippy( +# model, +# example_args=(input_tensor,), +# ) +# inference_inputs = torch.rand(batch_size, 3, 224, 224) +# inputs = send_to_device(inference_inputs, torch_device) +# with torch.no_grad(): +# output = model(inputs) +# # Zach: Check that we just grab the real outputs we need at the end +# if not state.is_last_process: +# assert output is None, "Output was not generated on just the last process!" +# else: +# assert output is not None, "Output was not generated in the last process!" + + +if __name__ == "__main__": + state = PartialState() + state.print("Testing pippy integration...") + try: + if state.distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_XPU, DistributedType.MULTI_HPU]: + state.print("Testing GPT2...") + test_gpt2() + # Issue: When modifying the tokenizer for batch GPT2 inference, there's an issue + # due to references + # NameError: cannot access free variable 'chunk_args_list' where it is not associated with a value in enclosing scope + # test_gpt2(3) + state.print("Testing BERT...") + test_bert() + else: + print("Less than two GPUs found, not running tests!") + finally: + state.destroy_process_group() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py new file mode 100644 index 0000000000000000000000000000000000000000..f5352b19980288115d2442229620105b8440d03d --- /dev/null +++ b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py @@ -0,0 +1,59 @@ +# Copyright 2024 The HuggingFace Team. 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. + +import torch.distributed + +from accelerate.test_utils import require_huggingface_suite, torch_device +from accelerate.utils import is_transformers_available + + +if is_transformers_available(): + from transformers import AutoModel, TrainingArguments + + +GPT2_TINY = "sshleifer/tiny-gpt2" + + +@require_huggingface_suite +def init_torch_dist_then_launch_deepspeed(): + if torch_device == "xpu": + backend = "ccl" + elif torch_device == "hpu": + backend = "hccl" + else: + backend = "nccl" + + torch.distributed.init_process_group(backend=backend) + deepspeed_config = { + "zero_optimization": { + "stage": 3, + }, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + } + train_args = TrainingArguments( + output_dir="./", + deepspeed=deepspeed_config, + ) + model = AutoModel.from_pretrained(GPT2_TINY) + assert train_args is not None + assert model is not None + + +def main(): + init_torch_dist_then_launch_deepspeed() + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/METADATA b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..592ac50132955ef28992b625cdd3d2a091e9d167 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/METADATA @@ -0,0 +1,443 @@ +Metadata-Version: 2.4 +Name: deepspeed +Version: 0.17.2 +Summary: DeepSpeed library +Home-page: http://deepspeed.ai +Author: DeepSpeed Team +Author-email: deepspeed-info@microsoft.com +License: Apache Software License 2.0 +Project-URL: Documentation, https://deepspeed.readthedocs.io +Project-URL: Source, https://github.com/deepspeedai/DeepSpeed +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Description-Content-Type: text/markdown +License-File: LICENSE +Requires-Dist: einops +Requires-Dist: hjson +Requires-Dist: msgpack +Requires-Dist: ninja +Requires-Dist: numpy +Requires-Dist: packaging>=20.0 +Requires-Dist: psutil +Requires-Dist: py-cpuinfo +Requires-Dist: pydantic>=2.0.0 +Requires-Dist: torch +Requires-Dist: tqdm +Provides-Extra: 1bit +Provides-Extra: 1bit-mpi +Requires-Dist: mpi4py; extra == "1bit-mpi" +Provides-Extra: readthedocs +Requires-Dist: autodoc_pydantic>=2.0.0; extra == "readthedocs" +Requires-Dist: docutils<0.18; extra == "readthedocs" +Requires-Dist: hjson; extra == "readthedocs" +Requires-Dist: packaging; extra == "readthedocs" +Requires-Dist: psutil; extra == "readthedocs" +Requires-Dist: py-cpuinfo; extra == "readthedocs" +Requires-Dist: pydantic>=2.0.0; extra == "readthedocs" +Requires-Dist: recommonmark; extra == "readthedocs" +Requires-Dist: sphinx_rtd_theme; extra == "readthedocs" +Requires-Dist: torch; extra == "readthedocs" +Requires-Dist: tqdm; extra == "readthedocs" +Provides-Extra: dev +Requires-Dist: accelerate; extra == "dev" +Requires-Dist: clang-format==18.1.3; extra == "dev" +Requires-Dist: comet_ml>=3.41.0; extra == "dev" +Requires-Dist: deepspeed-kernels; sys_platform == "linux" and extra == "dev" +Requires-Dist: docutils<0.18; extra == "dev" +Requires-Dist: future; extra == "dev" +Requires-Dist: importlib-metadata>=4; extra == "dev" +Requires-Dist: mup; extra == "dev" +Requires-Dist: pre-commit>=3.2.0; extra == "dev" +Requires-Dist: pytest<8.4.0,>=7.2.0; extra == "dev" +Requires-Dist: pytest-forked; extra == "dev" +Requires-Dist: pytest-randomly; extra == "dev" +Requires-Dist: pytest-xdist; extra == "dev" +Requires-Dist: qtorch==0.3.0; extra == "dev" +Requires-Dist: recommonmark; extra == "dev" +Requires-Dist: sphinx; extra == "dev" +Requires-Dist: sphinx-rtd-theme; extra == "dev" +Requires-Dist: tensorboard; extra == "dev" +Requires-Dist: torchvision; extra == "dev" +Requires-Dist: transformers>=4.51.3; extra == "dev" +Requires-Dist: wandb; extra == "dev" +Provides-Extra: autotuning +Requires-Dist: tabulate; extra == "autotuning" +Provides-Extra: autotuning-ml +Requires-Dist: hjson; extra == "autotuning-ml" +Requires-Dist: tabulate; extra == "autotuning-ml" +Requires-Dist: xgboost; extra == "autotuning-ml" +Provides-Extra: sparse-attn +Requires-Dist: triton==1.0.0; extra == "sparse-attn" +Provides-Extra: sparse +Requires-Dist: neural-compressor==2.1.0; extra == "sparse" +Provides-Extra: inf +Requires-Dist: google; extra == "inf" +Requires-Dist: lm-eval==0.3.0; extra == "inf" +Requires-Dist: protobuf; extra == "inf" +Requires-Dist: qtorch; extra == "inf" +Requires-Dist: safetensors; extra == "inf" +Requires-Dist: sentencepiece; extra == "inf" +Requires-Dist: transformers>=4.32.1; extra == "inf" +Provides-Extra: sd +Requires-Dist: diffusers>=0.25.0; extra == "sd" +Requires-Dist: triton>=2.1.0; extra == "sd" +Provides-Extra: triton +Requires-Dist: triton==2.1.0; extra == "triton" +Provides-Extra: deepcompile +Requires-Dist: scipy; extra == "deepcompile" +Provides-Extra: all +Requires-Dist: clang-format==18.1.3; extra == "all" +Requires-Dist: xgboost; extra == "all" +Requires-Dist: pydantic>=2.0.0; extra == "all" +Requires-Dist: comet_ml>=3.41.0; extra == "all" +Requires-Dist: deepspeed-kernels; sys_platform == "linux" and extra == "all" +Requires-Dist: sentencepiece; extra == "all" +Requires-Dist: qtorch; extra == "all" +Requires-Dist: psutil; extra == "all" +Requires-Dist: triton==2.1.0; extra == "all" +Requires-Dist: torch; extra == "all" +Requires-Dist: accelerate; extra == "all" +Requires-Dist: tabulate; extra == "all" +Requires-Dist: pytest-xdist; extra == "all" +Requires-Dist: tqdm; extra == "all" +Requires-Dist: importlib-metadata>=4; extra == "all" +Requires-Dist: py-cpuinfo; extra == "all" +Requires-Dist: pre-commit>=3.2.0; extra == "all" +Requires-Dist: triton==1.0.0; extra == "all" +Requires-Dist: protobuf; extra == "all" +Requires-Dist: lm-eval==0.3.0; extra == "all" +Requires-Dist: wandb; extra == "all" +Requires-Dist: triton>=2.1.0; extra == "all" +Requires-Dist: qtorch==0.3.0; extra == "all" +Requires-Dist: transformers>=4.51.3; extra == "all" +Requires-Dist: sphinx_rtd_theme; extra == "all" +Requires-Dist: safetensors; extra == "all" +Requires-Dist: torchvision; extra == "all" +Requires-Dist: pytest-randomly; extra == "all" +Requires-Dist: mup; extra == "all" +Requires-Dist: sphinx; extra == "all" +Requires-Dist: sphinx-rtd-theme; extra == "all" +Requires-Dist: hjson; extra == "all" +Requires-Dist: docutils<0.18; extra == "all" +Requires-Dist: future; extra == "all" +Requires-Dist: transformers>=4.32.1; extra == "all" +Requires-Dist: diffusers>=0.25.0; extra == "all" +Requires-Dist: pytest<8.4.0,>=7.2.0; extra == "all" +Requires-Dist: neural-compressor==2.1.0; extra == "all" +Requires-Dist: tensorboard; extra == "all" +Requires-Dist: recommonmark; extra == "all" +Requires-Dist: pytest-forked; extra == "all" +Requires-Dist: google; extra == "all" +Requires-Dist: scipy; extra == "all" +Requires-Dist: mpi4py; extra == "all" +Requires-Dist: autodoc_pydantic>=2.0.0; extra == "all" +Requires-Dist: packaging; extra == "all" +Dynamic: author +Dynamic: author-email +Dynamic: classifier +Dynamic: description +Dynamic: description-content-type +Dynamic: home-page +Dynamic: license +Dynamic: license-file +Dynamic: project-url +Dynamic: provides-extra +Dynamic: requires-dist +Dynamic: summary + +[![License Apache 2.0](https://badgen.net/badge/license/apache2.0/blue)](https://github.com/deepspeedai/DeepSpeed/blob/master/LICENSE) +[![PyPI version](https://badge.fury.io/py/deepspeed.svg)](https://pypi.org/project/deepspeed/) +[![Downloads](https://static.pepy.tech/badge/deepspeed)](https://pepy.tech/project/deepspeed) +[![Build](https://badgen.net/badge/build/check-status/blue)](#build-pipeline-status) +[![OpenSSF Best Practices](https://www.bestpractices.dev/projects/9530/badge)](https://www.bestpractices.dev/projects/9530) +[![Twitter](https://img.shields.io/twitter/follow/DeepSpeedAI)](https://twitter.com/intent/follow?screen_name=DeepSpeedAI) +[![Japanese Twitter](https://img.shields.io/badge/%E6%97%A5%E6%9C%AC%E8%AA%9ETwitter-%40DeepSpeedAI_JP-blue)](https://twitter.com/DeepSpeedAI_JP) +[![Chinese Zhihu](https://img.shields.io/badge/%E7%9F%A5%E4%B9%8E-%E5%BE%AE%E8%BD%AFDeepSpeed-blue)](https://www.zhihu.com/people/deepspeed) + + +
+ + +
+ +## Latest News + DeepSpeed empowers ChatGPT-like model training with a single click, offering 15x speedup over SOTA RLHF systems with unprecedented cost reduction at all scales; [learn how](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/deepspeed-chat). +* [2025/06] [Arctic Long Sequence Training (ALST) with DeepSpeed: Scalable And Efficient Training For Multi-Million Token Sequences](https://www.snowflake.com/en/engineering-blog/arctic-long-sequence-training-multi-million-token-ai/) +* [2025/04] [DeepCompile: Unlocking Compiler Optimization for Distributed Training](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/deepcompile/README.md) +* [2025/03] [DeepSpeed-AutoTP: Automatic Tensor Parallel Training of Hugging Face models](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/huggingface-tp/README.md) +* [2024/12] [Ulysses-Offload: Democratizing Long Context LLM Training ](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/ulysses-offload/README.md) +* [2024/12] [DeepSpeed-Domino: Communication-Free LLM Training Engine](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/deepspeed-domino/README.md) +* [2024/08] [DeepSpeed on Windows](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/README.md) [[日本語](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/japanese/README.md)] [[中文](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/chinese/README.md)] + + +
+ More news + +
+ +--- + +# Extreme Speed and Scale for DL Training and Inference + +***[DeepSpeed](https://www.deepspeed.ai/) enabled the world's most powerful language models (at the time of this writing) such as [MT-530B](https://www.microsoft.com/en-us/research/blog/using-deepspeed-and-megatron-to-train-megatron-turing-nlg-530b-the-worlds-largest-and-most-powerful-generative-language-model/) and [BLOOM](https://huggingface.co/blog/bloom-megatron-deepspeed)***. It is an easy-to-use deep learning optimization software suite that powers unprecedented scale and speed for both training and inference. With DeepSpeed you can: + +* Train/Inference dense or sparse models with billions or trillions of parameters +* Achieve excellent system throughput and efficiently scale to thousands of GPUs +* Train/Inference on resource constrained GPU systems +* Achieve unprecedented low latency and high throughput for inference +* Achieve extreme compression for an unparalleled inference latency and model size reduction with low costs + +--- + +# DeepSpeed's four innovation pillars + + + + +## DeepSpeed-Training + +DeepSpeed offers a confluence of system innovations, that has made large scale DL training effective, and efficient, greatly improved ease of use, and redefined the DL training landscape in terms of scale that is possible. These innovations such as ZeRO, 3D-Parallelism, DeepSpeed-MoE, ZeRO-Infinity, etc. fall under the training pillar. Learn more: [DeepSpeed-Training](https://www.deepspeed.ai/training/) + +## DeepSpeed-Inference + +DeepSpeed brings together innovations in parallelism technology such as tensor, pipeline, expert and ZeRO-parallelism, and combines them with high performance custom inference kernels, communication optimizations and heterogeneous memory technologies to enable inference at an unprecedented scale, while achieving unparalleled latency, throughput and cost reduction. This systematic composition of system technologies for inference falls under the inference pillar. Learn more: [DeepSpeed-Inference](https://www.deepspeed.ai/inference) + + +## DeepSpeed-Compression + +To further increase the inference efficiency, DeepSpeed offers easy-to-use and flexible-to-compose compression techniques for researchers and practitioners to compress their models while delivering faster speed, smaller model size, and significantly reduced compression cost. Moreover, SoTA innovations on compression like ZeroQuant and XTC are included under the compression pillar. Learn more: [DeepSpeed-Compression](https://www.deepspeed.ai/compression) + +## DeepSpeed4Science + +In line with Microsoft's mission to solve humanity's most pressing challenges, the DeepSpeed team at Microsoft is responding to this opportunity by launching a new initiative called *DeepSpeed4Science*, aiming to build unique capabilities through AI system technology innovations to help domain experts to unlock today's biggest science mysteries. Learn more: [tutorials](https://www.deepspeed.ai/deepspeed4science/) + +--- + +# DeepSpeed Software Suite + +## DeepSpeed Library + + The [DeepSpeed](https://github.com/deepspeedai/deepspeed) library (this repository) implements and packages the innovations and technologies in DeepSpeed Training, Inference and Compression Pillars into a single easy-to-use, open-sourced repository. It allows for easy composition of multitude of features within a single training, inference or compression pipeline. The DeepSpeed Library is heavily adopted by the DL community, and has been used to enable some of the most powerful models (see [DeepSpeed Adoption](#deepspeed-adoption)). + +## Model Implementations for Inference (MII) + + [Model Implementations for Inference (MII)](https://github.com/deepspeedai/deepspeed-mii) is an open-sourced repository for making low-latency and high-throughput inference accessible to all data scientists by alleviating the need to apply complex system optimization techniques themselves. Out-of-box, MII offers support for thousands of widely used DL models, optimized using DeepSpeed-Inference, that can be deployed with a few lines of code, while achieving significant latency reduction compared to their vanilla open-sourced versions. + +## DeepSpeed on Azure + + DeepSpeed users are diverse and have access to different environments. We recommend to try DeepSpeed on Azure as it is the simplest and easiest method. The recommended method to try DeepSpeed on Azure is through AzureML [recipes](https://github.com/Azure/azureml-examples/tree/main/v1/python-sdk/workflows/train/deepspeed). The job submission and data preparation scripts have been made available [here](https://github.com/deepspeedai/Megatron-DeepSpeed/tree/main/examples_deepspeed/azureml). For more details on how to use DeepSpeed on Azure, please follow the [Azure tutorial](https://www.deepspeed.ai/tutorials/azure/). + +--- + +# DeepSpeed Adoption + +DeepSpeed was an important part of Microsoft’s +[AI at Scale](https://www.microsoft.com/en-us/research/project/ai-at-scale/) +initiative to enable next-generation AI capabilities at scale, where you can find more +information [here](https://innovation.microsoft.com/en-us/exploring-ai-at-scale). + +DeepSpeed has been used to train many different large-scale models, below is a list of several examples that we are aware of (if you'd like to include your model please submit a PR): + + * [Megatron-Turing NLG (530B)](https://www.microsoft.com/en-us/research/blog/using-deepspeed-and-megatron-to-train-megatron-turing-nlg-530b-the-worlds-largest-and-most-powerful-generative-language-model/) + * [Jurassic-1 (178B)](https://uploads-ssl.webflow.com/60fd4503684b466578c0d307/61138924626a6981ee09caf6_jurassic_tech_paper.pdf) + * [BLOOM (176B)](https://huggingface.co/blog/bloom-megatron-deepspeed) + * [GLM (130B)](https://github.com/THUDM/GLM-130B) + * [xTrimoPGLM (100B)](https://www.biorxiv.org/content/10.1101/2023.07.05.547496v2) + * [YaLM (100B)](https://github.com/yandex/YaLM-100B) + * [GPT-NeoX (20B)](https://github.com/EleutherAI/gpt-neox) + * [AlexaTM (20B)](https://www.amazon.science/blog/20b-parameter-alexa-model-sets-new-marks-in-few-shot-learning) + * [Turing NLG (17B)](https://www.microsoft.com/en-us/research/blog/turing-nlg-a-17-billion-parameter-language-model-by-microsoft/) + * [METRO-LM (5.4B)](https://arxiv.org/pdf/2204.06644.pdf) + +DeepSpeed has been integrated with several different popular open-source DL frameworks such as: + +| | Documentation | +| ---------------------------------------------------------------------------------------------- | -------------------------------------------- | + | [Transformers with DeepSpeed](https://huggingface.co/docs/transformers/deepspeed) | +| | [Accelerate with DeepSpeed](https://huggingface.co/docs/accelerate/usage_guides/deepspeed) | +| | [Lightning with DeepSpeed](https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html#deepspeed) | +| | [MosaicML with DeepSpeed](https://docs.mosaicml.com/projects/composer/en/latest/trainer/using_the_trainer.html?highlight=deepspeed#deepspeed-integration) | +| | [Determined with DeepSpeed](https://docs.determined.ai/latest/training/apis-howto/deepspeed/overview.html) | +| | [MMEngine with DeepSpeed](https://mmengine.readthedocs.io/en/latest/common_usage/large_model_training.html#deepspeed) | + +--- + +# Build Pipeline Status + +| Description | Status | +| ----------- | ------ | +| NVIDIA | [![nv-torch110-p40](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-p40.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-p40.yml) [![nv-torch110-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-v100.yml) [![nv-torch-latest-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-latest-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-latest-v100.yml) [![nv-h100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-h100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-h100.yml) [![nv-inference](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-inference.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-inference.yml) [![nv-nightly](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-nightly.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-nightly.yml) | +| AMD | [![amd-mi200](https://github.com/deepspeedai/DeepSpeed/actions/workflows/amd-mi200.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/amd-mi200.yml) | +| CPU | [![torch-latest-cpu](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-torch-latest.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-torch-latest.yml) [![cpu-inference](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-inference.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-inference.yml) | +| Intel Gaudi | [![hpu-gaudi2](https://github.com/deepspeedai/DeepSpeed/actions/workflows/hpu-gaudi2.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/hpu-gaudi2.yml) | +| Intel XPU | [![xpu-max1100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/xpu-max1100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/xpu-max1100.yml) | +| PyTorch Nightly | [![nv-torch-nightly-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-nightly-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-nightly-v100.yml) | +| Integrations | [![nv-transformers-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-transformers-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-transformers-v100.yml) [![nv-lightning-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-lightning-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-lightning-v100.yml) [![nv-accelerate-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-accelerate-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-accelerate-v100.yml) [![nv-mii](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-mii.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-mii.yml) [![nv-ds-chat](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-ds-chat.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-ds-chat.yml) [![nv-sd](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-sd.yml/badge.svg)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-sd.yml) | +| Misc | [![Formatting](https://github.com/deepspeedai/DeepSpeed/actions/workflows/formatting.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/formatting.yml) [![pages-build-deployment](https://github.com/deepspeedai/DeepSpeed/actions/workflows/pages/pages-build-deployment/badge.svg)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/pages/pages-build-deployment) [![Documentation Status](https://readthedocs.org/projects/deepspeed/badge/?version=latest)](https://deepspeed.readthedocs.io/en/latest/?badge=latest)[![python](https://github.com/deepspeedai/DeepSpeed/actions/workflows/python.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/python.yml) | +| Huawei Ascend NPU | [![Huawei Ascend NPU](https://github.com/Ascend/Ascend-CI/actions/workflows/deepspeed.yaml/badge.svg?branch=main)](https://github.com/Ascend/Ascend-CI/actions/workflows/deepspeed.yaml) | + +# Installation + +The quickest way to get started with DeepSpeed is via pip, this will install +the latest release of DeepSpeed which is not tied to specific PyTorch or CUDA +versions. DeepSpeed includes several C++/CUDA extensions that we commonly refer +to as our 'ops'. By default, all of these extensions/ops will be built +just-in-time (JIT) using [torch's JIT C++ extension loader that relies on +ninja](https://pytorch.org/docs/stable/cpp_extension.html) to build and +dynamically link them at runtime. + +## Requirements +* [PyTorch](https://pytorch.org/) must be installed _before_ installing DeepSpeed. +* For full feature support we recommend a version of PyTorch that is >= 1.9 and ideally the latest PyTorch stable release. +* A CUDA or ROCm compiler such as [nvcc](https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/#introduction) or [hipcc](https://github.com/ROCm-Developer-Tools/HIPCC) used to compile C++/CUDA/HIP extensions. +* Specific GPUs we develop and test against are listed below, this doesn't mean your GPU will not work if it doesn't fall into this category it's just DeepSpeed is most well tested on the following: + * NVIDIA: Pascal, Volta, Ampere, and Hopper architectures + * AMD: MI100 and MI200 + +## Contributed HW support +* DeepSpeed now support various HW accelerators. + +| Contributor | Hardware | Accelerator Name | Contributor validated | Upstream validated | +|-------------|-------------------------------------|------------------| --------------------- |--------------------| +| Huawei | Huawei Ascend NPU | npu | Yes | No | +| Intel | Intel(R) Gaudi(R) 2 AI accelerator | hpu | Yes | Yes | +| Intel | Intel(R) Xeon(R) Processors | cpu | Yes | Yes | +| Intel | Intel(R) Data Center GPU Max series | xpu | Yes | Yes | +| Tecorigin | Scalable Data Analytics Accelerator | sdaa | Yes | No | + +## PyPI +We regularly push releases to [PyPI](https://pypi.org/project/deepspeed/) and encourage users to install from there in most cases. + +```bash +pip install deepspeed +``` + +After installation, you can validate your install and see which extensions/ops +your machine is compatible with via the DeepSpeed environment report. + +```bash +ds_report +``` + +If you would like to pre-install any of the DeepSpeed extensions/ops (instead +of JIT compiling) or install pre-compiled ops via PyPI please see our [advanced +installation instructions](https://www.deepspeed.ai/tutorials/advanced-install/). + +## Windows +Many DeepSpeed features are supported on Windows for both training and inference. You can read more about this in the original blog post [here](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/README.md). Among features that are currently not supported are async io (AIO) and GDS (which does not support Windows). +1. Install PyTorch, such as pytorch 2.3+cu121. +2. Install Visual C++ build tools, such as VS2022 C++ x64/x86 build tools. +3. Launch Cmd console with Administrator permissions for creating required symlink folders and ensure MSVC tools are added to your PATH or launch the Developer Command Prompt for Visual Studio 2022 with administrator permissions. +4. Run `build_win.bat` to build wheel in `dist` folder. + +# Features + +Please checkout [DeepSpeed-Training](https://www.deepspeed.ai/training), [DeepSpeed-Inference](https://www.deepspeed.ai/inference) and [DeepSpeed-Compression](https://www.deepspeed.ai/compression) pages for full set of features offered along each of these three pillars. + +# Further Reading + +All DeepSpeed documentation, tutorials, and blogs can be found on our website: [deepspeed.ai](https://www.deepspeed.ai/) + + +| | Description | +| ---------------------------------------------------------------------------------------------- | -------------------------------------------- | +| [Getting Started](https://www.deepspeed.ai/getting-started/) | First steps with DeepSpeed | +| [DeepSpeed JSON Configuration](https://www.deepspeed.ai/docs/config-json/) | Configuring DeepSpeed | +| [API Documentation](https://deepspeed.readthedocs.io/en/latest/) | Generated DeepSpeed API documentation | +| [Tutorials](https://www.deepspeed.ai/tutorials/) | Tutorials | +| [Blogs](https://www.deepspeed.ai/posts/) | Blogs | + + +# Contributing +DeepSpeed welcomes your contributions! Please see our +[contributing](CONTRIBUTING.md) guide for more details on formatting, testing, +etc.
+Thanks so much to all of our amazing contributors! + + + + + +## Contributor License Agreement +This project welcomes contributions and suggestions. Most contributions require you to +agree to a Contributor License Agreement (CLA) declaring that you have the right to, and +actually do, grant us the rights to use your contribution. For details, visit +https://cla.opensource.microsoft.com. + +When you submit a pull request, a CLA bot will automatically determine whether you need +to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply +follow the instructions provided by the bot. You will only need to do this once across +all repos using our CLA. + +## Code of Conduct +This project has adopted the [Microsoft Open Source Code of +Conduct](https://opensource.microsoft.com/codeofconduct/). For more information see the +[Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or contact +[opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. + +# Publications +1. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He. (2019) ZeRO: memory optimizations toward training trillion parameter models. [arXiv:1910.02054](https://arxiv.org/abs/1910.02054) and [In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC '20)](https://dl.acm.org/doi/10.5555/3433701.3433727). +2. Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. (2020) DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters. [In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (KDD '20, Tutorial)](https://dl.acm.org/doi/10.1145/3394486.3406703). +3. Minjia Zhang, Yuxiong He. (2020) Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping. [arXiv:2010.13369](https://arxiv.org/abs/2010.13369) and [NeurIPS 2020](https://proceedings.neurips.cc/paper/2020/hash/a1140a3d0df1c81e24ae954d935e8926-Abstract.html). +4. Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, Yuxiong He. (2021) ZeRO-Offload: Democratizing Billion-Scale Model Training. [arXiv:2101.06840](https://arxiv.org/abs/2101.06840) and [USENIX ATC 2021](https://www.usenix.org/conference/atc21/presentation/ren-jie). [[paper]](https://arxiv.org/abs/2101.06840) [[slides]](https://www.usenix.org/system/files/atc21_slides_ren-jie.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/) +5. Hanlin Tang, Shaoduo Gan, Ammar Ahmad Awan, Samyam Rajbhandari, Conglong Li, Xiangru Lian, Ji Liu, Ce Zhang, Yuxiong He. (2021) 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed. [arXiv:2102.02888](https://arxiv.org/abs/2102.02888) and [ICML 2021](http://proceedings.mlr.press/v139/tang21a.html). +6. Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, Yuxiong He. (2021) ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning. [arXiv:2104.07857](https://arxiv.org/abs/2104.07857) and [SC 2021](https://dl.acm.org/doi/abs/10.1145/3458817.3476205). [[paper]](https://arxiv.org/abs/2104.07857) [[slides]](docs/assets/files/SC21-ZeRO-Infinity.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/zero-infinity-and-deepspeed-unlocking-unprecedented-model-scale-for-deep-learning-training/) +7. Conglong Li, Ammar Ahmad Awan, Hanlin Tang, Samyam Rajbhandari, Yuxiong He. (2021) 1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed. [arXiv:2104.06069](https://arxiv.org/abs/2104.06069) and [HiPC 2022](https://hipc.org/advance-program/). +8. Conglong Li, Minjia Zhang, Yuxiong He. (2021) The Stability-Efficiency Dilemma: Investigating Sequence Length Warmup for Training GPT Models. [arXiv:2108.06084](https://arxiv.org/abs/2108.06084) and [NeurIPS 2022](https://openreview.net/forum?id=JpZ5du_Kdh). +9. Yucheng Lu, Conglong Li, Minjia Zhang, Christopher De Sa, Yuxiong He. (2022) Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam. [arXiv:2202.06009](https://arxiv.org/abs/2202.06009). +10. Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He. (2022) DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale [arXiv:2201.05596](https://arxiv.org/abs/2201.05596) and [ICML 2022](https://proceedings.mlr.press/v162/rajbhandari22a.html). [[pdf]](https://arxiv.org/abs/2201.05596) [[slides]](docs/assets/files/ICML-5mins.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-advancing-moe-inference-and-training-to-power-next-generation-ai-scale/) +11. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, Bryan Catanzaro. (2022) Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model [arXiv:2201.11990](https://arxiv.org/abs/2201.11990). +12. Xiaoxia Wu, Zhewei Yao, Minjia Zhang, Conglong Li, Yuxiong He. (2022) Extreme Compression for Pre-trained Transformers Made Simple and Efficient. [arXiv:2206.01859](https://arxiv.org/abs/2206.01859) and [NeurIPS 2022](https://openreview.net/forum?id=xNeAhc2CNAl). +13. Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, Yuxiong He. (2022) ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers. [arXiv:2206.01861](https://arxiv.org/abs/2206.01861) and [NeurIPS 2022](https://openreview.net/forum?id=f-fVCElZ-G1) [[slides]](docs/assets/files/zeroquant_series.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-compression-a-composable-library-for-extreme-compression-and-zero-cost-quantization/) +14. Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, Yuxiong He. (2022) DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale. [arXiv:2207.00032](https://arxiv.org/abs/2207.00032) and [SC 2022](https://dl.acm.org/doi/abs/10.5555/3571885.3571946). [[paper]](https://arxiv.org/abs/2207.00032) [[slides]](docs/assets/files/sc22-ds-inference.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-accelerating-large-scale-model-inference-and-training-via-system-optimizations-and-compression/) +15. Zhewei Yao, Xiaoxia Wu, Conglong Li, Connor Holmes, Minjia Zhang, Cheng Li, Yuxiong He. (2022) Random-LTD: Random and Layerwise Token Dropping Brings Efficient Training for Large-scale Transformers. [arXiv:2211.11586](https://arxiv.org/abs/2211.11586). +16. Conglong Li, Zhewei Yao, Xiaoxia Wu, Minjia Zhang, Yuxiong He. (2022) DeepSpeed Data Efficiency: Improving Deep Learning Model Quality and Training Efficiency via Efficient Data Sampling and Routing. [arXiv:2212.03597](https://arxiv.org/abs/2212.03597) [ENLSP2023 Workshop at NeurIPS2023](https://neurips2023-enlsp.github.io/) +17. Xiaoxia Wu, Cheng Li, Reza Yazdani Aminabadi, Zhewei Yao, Yuxiong He. (2023) Understanding INT4 Quantization for Transformer Models: Latency Speedup, Composability, and Failure Cases. [arXiv:2301.12017](https://arxiv.org/abs/2301.12017) and [ICML2023](https://icml.cc/Conferences/2023). +18. Syed Zawad, Cheng Li, Zhewei Yao, Elton Zheng, Yuxiong He, Feng Yan. (2023) DySR: Adaptive Super-Resolution via Algorithm and System Co-design. [ICLR:2023](https://openreview.net/forum?id=Pgtn4l6eKjv). +19. Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, Yuxiong He. (2023) Scaling Vision-Language Models with Sparse Mixture of Experts. [arXiv:2303.07226](https://arxiv.org/abs/2303.07226) and [Finding at EMNLP2023](https://2023.emnlp.org/). +20. Quentin Anthony, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He, Aamir Shafi, Mustafa Abduljabbar, Hari Subramoni, Dhabaleswar Panda. (2023) MCR-DL: Mix-and-Match Communication Runtime for Deep Learning [arXiv:2303.08374](https://arxiv.org/abs/2303.08374) and will appear at IPDPS 2023. +21. Siddharth Singh, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He, Abhinav Bhatele. (2023) A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training [arXiv:2303.06318](https://arxiv.org/abs/2303.06318) and [ICS 2023](https://dl.acm.org/doi/10.1145/3577193.3593704). +22. Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Xiaoxia Wu, Connor Holmes, Zhewei Yao, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, Yuxiong He. (2023) ZeRO++: Extremely Efficient Collective Communication for Giant Model Training [arXiv:2306.10209](https://arxiv.org/abs/2306.10209) and [ML for Sys Workshop at NeurIPS2023](http://mlforsystems.org/) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-zero-a-leap-in-speed-for-llm-and-chat-model-training-with-4x-less-communication/) +23. Zhewei Yao, Xiaoxia Wu, Cheng Li, Stephen Youn, Yuxiong He. (2023) ZeroQuant-V2: Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation [arXiv:2303.08302](https://arxiv.org/abs/2303.08302) and [ENLSP2023 Workshop at NeurIPS2023](https://neurips2023-enlsp.github.io/) [[slides]](docs/assets/files/zeroquant_series.pdf) +24. Pareesa Ameneh Golnari, Zhewei Yao, Yuxiong He. (2023) Selective Guidance: Are All the Denoising Steps of Guided Diffusion Important? [arXiv:2305.09847](https://arxiv.org/abs/2305.09847) +25. Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari, Xiaoxia Wu, Ammar Ahmad Awan, Jeff Rasley, Minjia Zhang, Conglong Li, Connor Holmes, Zhongzhu Zhou, Michael Wyatt, Molly Smith, Lev Kurilenko, Heyang Qin, Masahiro Tanaka, Shuai Che, Shuaiwen Leon Song, Yuxiong He. (2023) DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales [arXiv:2308.01320](https://arxiv.org/abs/2308.01320). +26. Xiaoxia Wu, Zhewei Yao, Yuxiong He. (2023) ZeroQuant-FP: A Leap Forward in LLMs Post-Training W4A8 Quantization Using Floating-Point Formats [arXiv:2307.09782](https://arxiv.org/abs/2307.09782) and [ENLSP2023 Workshop at NeurIPS2023](https://neurips2023-enlsp.github.io/) [[slides]](docs/assets/files/zeroquant_series.pdf) +27. Zhewei Yao, Xiaoxia Wu, Conglong Li, Minjia Zhang, Heyang Qin, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He. (2023) DeepSpeed-VisualChat: Multi-Round Multi-Image Interleave Chat via Multi-Modal Causal Attention [arXiv:2309.14327](https://arxiv.org/pdf/2309.14327.pdf) +28. Shuaiwen Leon Song, Bonnie Kruft, Minjia Zhang, Conglong Li, Shiyang Chen, Chengming Zhang, Masahiro Tanaka, Xiaoxia Wu, Jeff Rasley, Ammar Ahmad Awan, Connor Holmes, Martin Cai, Adam Ghanem, Zhongzhu Zhou, Yuxiong He, et al. (2023) DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies [arXiv:2310.04610](https://arxiv.org/abs/2310.04610) [[blog]](https://www.microsoft.com/en-us/research/blog/announcing-the-deepspeed4science-initiative-enabling-large-scale-scientific-discovery-through-sophisticated-ai-system-technologies/) +29. Zhewei Yao, Reza Yazdani Aminabadi, Stephen Youn, Xiaoxia Wu, Elton Zheng, Yuxiong He. (2023) ZeroQuant-HERO: Hardware-Enhanced Robust Optimized Post-Training Quantization Framework for W8A8 Transformers [arXiv:2310.17723](https://arxiv.org/abs/2310.17723) + +30. Xiaoxia Wu, Haojun Xia, Stephen Youn, Zhen Zheng, Shiyang Chen, Arash Bakhtiari, Michael Wyatt, Reza Yazdani Aminabadi, Yuxiong He, Olatunji Ruwase, Leon Song, Zhewei Yao (2023) ZeroQuant(4+2): Redefining LLMs Quantization with a New FP6-Centric Strategy for Diverse Generative Tasks [arXiv:2312.08583](https://arxiv.org/abs/2312.08583) + +31. Haojun Xia, Zhen Zheng, Xiaoxia Wu, Shiyang Chen, Zhewei Yao, Stephen Youn, Arash Bakhtiari, Michael Wyatt, Donglin Zhuang, Zhongzhu Zhou, Olatunji Ruwase, Yuxiong He, Shuaiwen Leon Song. (2024) FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design [arXiv:2401.14112](https://arxiv.org/abs/2401.14112) +32. Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Reza Yazdani Aminadabi, Shuaiwen Leon Song, Samyam Rajbhandari, Yuxiong He. (2024) [System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://dl.acm.org/doi/10.1145/3662158.3662806) +33. Xinyu Lian, Sam Ade Jacobs, Lev Kurilenko, Masahiro Tanaka, Stas Bekman, Olatunji Ruwase, Minjia Zhang. (2024) Universal Checkpointing: Efficient and Flexible Checkpointing for Large Scale Distributed Training [arXiv:2406.18820](https://arxiv.org/abs/2406.18820) +34. Stas Bekman, Samyam Rajbhandari, Michael Wyatt, Jeff Rasley, Tunji Ruwase, Zhewei Yao, Aurick Qiao, Yuxiong He. (2025) Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences [arXiv:2506.13996](https://arxiv.org/abs/2506.13996) + + +# Videos +1. DeepSpeed KDD 2020 Tutorial + 1. [Overview](https://www.youtube.com/watch?v=CaseqC45DNc&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=29) + 2. [ZeRO + large model training](https://www.youtube.com/watch?v=y4_bCiAsIAk&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=28) + 3. [17B T-NLG demo](https://www.youtube.com/watch?v=9V-ZbP92drg&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=27) + 4. [Fastest BERT training + RScan tuning](https://www.youtube.com/watch?v=o1K-ZG9F6u0&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=26) + 5. DeepSpeed hands on deep dive: [part 1](https://www.youtube.com/watch?v=_NOk-mBwDYg&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=92), [part 2](https://www.youtube.com/watch?v=sG6_c4VXLww&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=94), [part 3](https://www.youtube.com/watch?v=k9yPkBTayos&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=93) + 6. [FAQ](https://www.youtube.com/watch?v=nsHu6vEgPew&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=24) +2. Microsoft Research Webinar + * Registration is free and all videos are available on-demand. + * [ZeRO & Fastest BERT: Increasing the scale and speed of deep learning training in DeepSpeed](https://note.microsoft.com/MSR-Webinar-DeepSpeed-Registration-On-Demand.html). +3. [DeepSpeed on AzureML](https://youtu.be/yBVXR8G8Bg8) +4. [Large Model Training and Inference with DeepSpeed // Samyam Rajbhandari // LLMs in Prod Conference](https://www.youtube.com/watch?v=cntxC3g22oU) [[slides]](docs/assets/files/presentation-mlops.pdf) +5. Community Tutorials + * [DeepSpeed: All the tricks to scale to gigantic models (Mark Saroufim)](https://www.youtube.com/watch?v=pDGI668pNg0) + * [Turing-NLG, DeepSpeed and the ZeRO optimizer (Yannic Kilcher)](https://www.youtube.com/watch?v=tC01FRB0M7w) + * [Ultimate Guide To Scaling ML Models (The AI Epiphany)](https://www.youtube.com/watch?v=hc0u4avAkuM) diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/RECORD b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..0eb7d39a8a0d89ad563c4f08aefeb2e0c1511303 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/RECORD @@ -0,0 +1,1580 @@ +../../../bin/deepspeed,sha256=fVvGbO4f9l-ArUn5Ic-zdG_onz7cAAiR6zp4IMhwP0w,120 +../../../bin/deepspeed.pt,sha256=fVvGbO4f9l-ArUn5Ic-zdG_onz7cAAiR6zp4IMhwP0w,120 +../../../bin/ds,sha256=fVvGbO4f9l-ArUn5Ic-zdG_onz7cAAiR6zp4IMhwP0w,120 +../../../bin/ds_bench,sha256=NCeQ5VdNlkSZRWIC2aQcrjIhd4Wz5OEEA9P-ibHbbec,870 +../../../bin/ds_elastic,sha256=9KRAtw2socyGIPD2PyG1hyHxeLAejtiJF0RIKVoAERA,1883 +../../../bin/ds_io,sha256=SSadr7J4RwFc3drWWJ6Eyvt5fon_WwNff5mwunNAN0Q,121 +../../../bin/ds_nvme_tune,sha256=4MaRWLsmGkRMji03UQB4eK5inNdzlfjrtEhGVxTB6Ec,299 +../../../bin/ds_report,sha256=SA4zydPPRa4xRemY_nRR1yqbFjSoYKdUF7_vQuksZBU,123 +../../../bin/ds_ssh,sha256=BT6cdZ47ceJiKPK40eVf9NCXF5RLjayEwWlyZ7viWbE,680 +../../../bin/dsr,sha256=SA4zydPPRa4xRemY_nRR1yqbFjSoYKdUF7_vQuksZBU,123 +deepspeed-0.17.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +deepspeed-0.17.2.dist-info/METADATA,sha256=cq8vc6YI5SGW6twYehsWxaN1IauzMrlCSpIx1nsE-z4,43003 +deepspeed-0.17.2.dist-info/RECORD,, +deepspeed-0.17.2.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +deepspeed-0.17.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91 +deepspeed-0.17.2.dist-info/entry_points.txt,sha256=btWVtkuERrrqyxQojUVMJZOeaoNPz0pjiZKaJ_Q4MJ8,84 +deepspeed-0.17.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357 +deepspeed-0.17.2.dist-info/top_level.txt,sha256=VK6WJ48PiB2ke4oVIo2mLzqXYqGXdu4MPE2GP-vfvA8,10 +deepspeed/__init__.py,sha256=tcS3eXx2HvMLr38_ofg-2RC9l9k1s3AsYX-TyJEBrJU,16832 +deepspeed/__pycache__/__init__.cpython-310.pyc,, +deepspeed/__pycache__/constants.cpython-310.pyc,, +deepspeed/__pycache__/env_report.cpython-310.pyc,, +deepspeed/__pycache__/git_version_info.cpython-310.pyc,, +deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc,, +deepspeed/accelerator/__init__.py,sha256=1pNnvGX9h418ekjY_Xnpb363tLo-5Szw3icQf9hRvk0,248 +deepspeed/accelerator/__pycache__/__init__.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/abstract_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/cpu_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/cuda_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/hpu_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/mlu_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/mps_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/npu_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/real_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/sdaa_accelerator.cpython-310.pyc,, +deepspeed/accelerator/__pycache__/xpu_accelerator.cpython-310.pyc,, +deepspeed/accelerator/abstract_accelerator.py,sha256=mF-Km8bgg_yGqwcoPcSheKRlYoyFP4i1cFzxUHRL72o,5840 +deepspeed/accelerator/cpu_accelerator.py,sha256=Mn58Mq3eZJr5QWmQliXOY3-D5c1ne-uzmYgeiz1oexM,10553 +deepspeed/accelerator/cuda_accelerator.py,sha256=mwj2tI4736nG1XD3C3p0OqCIpUKt1SzEir1PmDaPVtA,12920 +deepspeed/accelerator/hpu_accelerator.py,sha256=U8MLbGI_qP6PnTwP2F7ibzUalB0-MxnhsJr9Ld0F3BA,10971 +deepspeed/accelerator/mlu_accelerator.py,sha256=1tUxv8mnAHOfo4AY4BxiQ4rCD8KULFKZVeNRULOQNB4,9337 +deepspeed/accelerator/mps_accelerator.py,sha256=XkqWL7P86KFp2ebtGSNF6pPsR5OPsFm1Jl5UPEcR6C4,7006 +deepspeed/accelerator/npu_accelerator.py,sha256=FwTKiHZeaWW5PngtqZYFS49w5hKFmQwhsOPR1KreX2s,9020 +deepspeed/accelerator/real_accelerator.py,sha256=CJS2JocHzQCeGJRGQqxpGdzgLJBHqzR31PDB_SaFxN0,12791 +deepspeed/accelerator/sdaa_accelerator.py,sha256=YXdjlEwIUYuocPsu-0_etZH-2ak4yCNQYAFGL-_DKgM,10813 +deepspeed/accelerator/xpu_accelerator.py,sha256=AopBKrqyTu4_GLr1LpTeH075PUQ2iparLBU56ceDng8,10936 +deepspeed/autotuning/__init__.py,sha256=y0O9XDcX76E6dmhNFmRhy9TKfl0ywJrdkXmM2JuzErU,129 +deepspeed/autotuning/__pycache__/__init__.cpython-310.pyc,, +deepspeed/autotuning/__pycache__/autotuner.cpython-310.pyc,, +deepspeed/autotuning/__pycache__/config.cpython-310.pyc,, +deepspeed/autotuning/__pycache__/constants.cpython-310.pyc,, +deepspeed/autotuning/__pycache__/scheduler.cpython-310.pyc,, +deepspeed/autotuning/__pycache__/utils.cpython-310.pyc,, +deepspeed/autotuning/autotuner.py,sha256=_ucGsohYkJ1MFsz1JuB9DJyky-brugx_IDz0mbhqeY0,54335 +deepspeed/autotuning/config.py,sha256=pg5OeryVqrg32xak1NLEWaFiN1MYujdwlSgTuAz3Xl0,4633 +deepspeed/autotuning/config_templates/template_zero0.json,sha256=hR1baTMA5HzfTD1JlXjiTYjMtB9C8CFfxIBzqZeW2Sw,48 +deepspeed/autotuning/config_templates/template_zero1.json,sha256=AAIox-1dviiYg-Z7L-02bxItSnaggAlHkR4TekXXEIw,113 +deepspeed/autotuning/config_templates/template_zero2.json,sha256=FfrPKZHzMJuHn_TVJ7UZ-2s9yJ5IzBIKgT7nLYnIHh4,237 +deepspeed/autotuning/config_templates/template_zero3.json,sha256=WvsoMMp5MDgdVLq1RX1fhKsKbGSHQtpplZzGJikZvzs,485 +deepspeed/autotuning/constants.py,sha256=sr02-8B5ORdimK5VZ1N9FK1luWadmECFLSlMnM-uIxA,5943 +deepspeed/autotuning/scheduler.py,sha256=4IjLO56K-ZXr-HK1FuSuw5-v4tUdAWRZfH8R02WNx3U,15703 +deepspeed/autotuning/tuner/__init__.py,sha256=c9ImdL2iEc89lFltZ2PPe2EOUCuB3Olq3hRP1Q-yUMo,235 +deepspeed/autotuning/tuner/__pycache__/__init__.cpython-310.pyc,, +deepspeed/autotuning/tuner/__pycache__/base_tuner.cpython-310.pyc,, +deepspeed/autotuning/tuner/__pycache__/cost_model.cpython-310.pyc,, +deepspeed/autotuning/tuner/__pycache__/index_based_tuner.cpython-310.pyc,, +deepspeed/autotuning/tuner/__pycache__/model_based_tuner.cpython-310.pyc,, +deepspeed/autotuning/tuner/__pycache__/utils.cpython-310.pyc,, +deepspeed/autotuning/tuner/base_tuner.py,sha256=psA1I4-AEfqGrGLvYYUl3EjM8jG8AlhUXz-AnrFZ7G4,2754 +deepspeed/autotuning/tuner/cost_model.py,sha256=Uu9jD65cvdUK6aozm6Du3XJQZ_W-yOyUrSWFLdOT0do,1820 +deepspeed/autotuning/tuner/index_based_tuner.py,sha256=AEkTByT3XXCyYolAqBiMquV6XHN-ntGmYKMLZKPCRYE,1158 +deepspeed/autotuning/tuner/model_based_tuner.py,sha256=f_CrbgZNONRjdMuRhrJkQbhtih6GOg4LdkXqRDRDzaM,5614 +deepspeed/autotuning/tuner/utils.py,sha256=o5nD51Z6LBylJDRnDNIiOARlBHnOgINBVjBos6ypLEg,2329 +deepspeed/autotuning/utils.py,sha256=d4ePhVxVtekB7gIEK1dIOs6bhxOpmFF8yhKPC_rNRDc,15045 +deepspeed/checkpoint/__init__.py,sha256=3qhX6swvx0-zo72xLl99eZzBr_KC7PguCVCgemYyAGw,576 +deepspeed/checkpoint/__pycache__/__init__.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/constants.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/deepspeed_checkpoint.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/ds_to_universal.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/reshape_3d_utils.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/reshape_meg_2d.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/reshape_utils.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/universal_checkpoint.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/utils.cpython-310.pyc,, +deepspeed/checkpoint/__pycache__/zero_checkpoint.cpython-310.pyc,, +deepspeed/checkpoint/constants.py,sha256=FQ55iGvLxiXitjitVfepx3-vLrPx7J-btyfLUGyPGCc,3510 +deepspeed/checkpoint/deepspeed_checkpoint.py,sha256=mrL3zQ6pvWBQIZyBtzsXTNEIt2jbyemunrweHGE-Cl8,13095 +deepspeed/checkpoint/ds_to_universal.py,sha256=IbP-9e1A6Y95I6AWLolK1OCHjJ78Ndrp9i07zaLqJN0,22788 +deepspeed/checkpoint/reshape_3d_utils.py,sha256=pmnOtdI3srRS46R9sXLh9kijXPxT3bFJ_c6lWdygbwc,4674 +deepspeed/checkpoint/reshape_meg_2d.py,sha256=o5dpF3CbrbtEt1twz0XsL3xbqCfLfEse7O2ke-WR9_s,7885 +deepspeed/checkpoint/reshape_utils.py,sha256=n3k8xjXn2--nRzcs7jo7V9lFlma7LTLXjkk9S3IgUhM,3458 +deepspeed/checkpoint/universal_checkpoint.py,sha256=EF1VbI2eeHdHXmdgYbHCYVOIf8sjSTcGqY82qjtCtmQ,6788 +deepspeed/checkpoint/utils.py,sha256=Wog7wLoR_F88UOSlLRAX8XspdjA9mF3CiJ5iOYVWh9E,2693 +deepspeed/checkpoint/zero_checkpoint.py,sha256=OsB42DLdUAeuqFbynEHaJGAo3UTbuog1UicXg8dVJAU,5441 +deepspeed/comm/__init__.py,sha256=N31DOMdAdlfSnQd64slK89dpW_y1Vc9oPVZabLArcYw,137 +deepspeed/comm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/comm/__pycache__/backend.cpython-310.pyc,, +deepspeed/comm/__pycache__/ccl.cpython-310.pyc,, +deepspeed/comm/__pycache__/comm.cpython-310.pyc,, +deepspeed/comm/__pycache__/config.cpython-310.pyc,, +deepspeed/comm/__pycache__/constants.cpython-310.pyc,, +deepspeed/comm/__pycache__/reduce_op.cpython-310.pyc,, +deepspeed/comm/__pycache__/torch.cpython-310.pyc,, +deepspeed/comm/__pycache__/utils.cpython-310.pyc,, +deepspeed/comm/backend.py,sha256=TPFMx6ViU4n_NcfwWS56_Yp35PeCi-0Kh9za7UELeyA,1416 +deepspeed/comm/ccl.py,sha256=wGGn3sFUvbzblOsA4uMfKCltk6giyI2dkG6SkmMFzVc,8614 +deepspeed/comm/comm.py,sha256=XFedePFBQqfJngaz1WWhC9dLVEenUd-bHpR0M2JQWlQ,31135 +deepspeed/comm/config.py,sha256=vd3v2i0X638bgqgVENb03GddaJfZVJYV17k3dRHdmKM,728 +deepspeed/comm/constants.py,sha256=Adtnb5LCEjRpylnwgsMAFRmWgPvxb8CEEDvn4-D39lw,1298 +deepspeed/comm/reduce_op.py,sha256=A_KxQtpfUrNun6Rn8SX6iQiuoVgfYM7bTtAYc-YTENE,259 +deepspeed/comm/torch.py,sha256=IOBv9fuZjBdKK_vQvoOnHbUAzlYKnJSgzTCzKVByjs8,19481 +deepspeed/comm/utils.py,sha256=I7hJ2Jq5HEPaJDD09eEI9xV5VJZeDdag0ovTeAc3Ex8,3842 +deepspeed/compile/__init__.py,sha256=-DiusbPt_BUtHIDjxVtMiLEuFDUb6ANm671qfLuUXVY,126 +deepspeed/compile/__pycache__/__init__.cpython-310.pyc,, +deepspeed/compile/__pycache__/backend.cpython-310.pyc,, +deepspeed/compile/__pycache__/config.cpython-310.pyc,, +deepspeed/compile/__pycache__/fx.cpython-310.pyc,, +deepspeed/compile/__pycache__/graph_param.cpython-310.pyc,, +deepspeed/compile/__pycache__/inductor.cpython-310.pyc,, +deepspeed/compile/__pycache__/init_z1.cpython-310.pyc,, +deepspeed/compile/__pycache__/init_z3.cpython-310.pyc,, +deepspeed/compile/__pycache__/input_storage.cpython-310.pyc,, +deepspeed/compile/__pycache__/list_schedule.cpython-310.pyc,, +deepspeed/compile/__pycache__/partitioner.cpython-310.pyc,, +deepspeed/compile/__pycache__/patch_compiled_func.cpython-310.pyc,, +deepspeed/compile/__pycache__/patch_fake_tensor.cpython-310.pyc,, +deepspeed/compile/__pycache__/util.cpython-310.pyc,, +deepspeed/compile/backend.py,sha256=7J3IuJJwMKDuH5hg5Ot4nwJhjDbesQxGRwT9-SV3pKc,15885 +deepspeed/compile/config.py,sha256=n-0B63IyyOlK6j-bfUSfVlCxdvJHMDdjPmadKYrb2MQ,1521 +deepspeed/compile/fx.py,sha256=WY5Shuq80GsEj1r3vNd20azPZNfcbMbGkr-c6JfMLVY,4717 +deepspeed/compile/graph_param.py,sha256=KBsOiXs6KfQ_p3UAakfOoTFQfvKXaEc-zNcnRh33V3w,3802 +deepspeed/compile/inductor.py,sha256=dfQC6RSYFZbR2_IXtfkrhNl88t31QZjosp159k477q0,8933 +deepspeed/compile/init_z1.py,sha256=slw23gEhnglbkN1UciCkrgClPmY1cCIf7JHWABwf-n0,3488 +deepspeed/compile/init_z3.py,sha256=59J3fEVt0J1cS_hAjE13tj59SO4RUyE2xv72khpVJOg,3905 +deepspeed/compile/input_storage.py,sha256=u71Xm3YBa5TV7in8SvqzF7PJA4iVO0Fjyw9YlRibwfc,7835 +deepspeed/compile/list_schedule.py,sha256=S2d9OolXfOwpyUkE68QHslWUPOmGWCBc8ijx2qCAsjc,15870 +deepspeed/compile/partitioner.py,sha256=fOtJtoRgoULNvYImDDgAnYli0OH3khT3W0GI98dM2Q8,6902 +deepspeed/compile/passes/__init__.py,sha256=YPYE70D8HDgY-Wo0z8HgOF7ikgUN0F29deid37kdF8k,1470 +deepspeed/compile/passes/__pycache__/__init__.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/offload_activation.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/offload_adam_states.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/offload_parameters.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/prefetch.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/selective_gather.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/zero1_compile.cpython-310.pyc,, +deepspeed/compile/passes/__pycache__/zero3_compile.cpython-310.pyc,, +deepspeed/compile/passes/offload_activation.py,sha256=fPCEG4OtfyMdLEtVCzZYrlvkeZkAeiVX29qh7tlrt_0,3849 +deepspeed/compile/passes/offload_adam_states.py,sha256=vprSD81ZODjObxQFOpkAoawQoyD1nL_X2DVYeKX8YpY,20605 +deepspeed/compile/passes/offload_parameters.py,sha256=zOOIpmbra727IUfSb-Rihuco-PVcTF2cJMfys626ZbM,2046 +deepspeed/compile/passes/prefetch.py,sha256=6NW_QLnIw6VzV5edhmb6sonp0Wu-TSqRPh7KhMO4Ho0,7696 +deepspeed/compile/passes/selective_gather.py,sha256=icQWdK73H61VaD54xZdq-R-9hHvLF7-spofiwdeK3Jw,5466 +deepspeed/compile/passes/zero1_compile.py,sha256=UHM5rtq1JX6Rjj0asbF7Nt9EUy25pvalseM4xXd7QsY,2323 +deepspeed/compile/passes/zero3_compile.py,sha256=PzrtTW-oDyEWQb_J38F3cAzIsRvU0w100SWv-FnB8L8,7645 +deepspeed/compile/patch_compiled_func.py,sha256=ijs5Yo7IkLve0j5Hpwe3QxiiTL1AEhOrlGCW7GiAjVo,2856 +deepspeed/compile/patch_fake_tensor.py,sha256=1d-sXFu63HF2vCPqA7Wg3rVvUupeQjezrRlqeMOaYxU,2010 +deepspeed/compile/profilers/__init__.py,sha256=OSidjoyeb6NwY7vBzdvPEJHDPpS08WRUbR-m8HFLyXE,957 +deepspeed/compile/profilers/__pycache__/__init__.cpython-310.pyc,, +deepspeed/compile/profilers/__pycache__/comm_profile.cpython-310.pyc,, +deepspeed/compile/profilers/__pycache__/graph_profile.cpython-310.pyc,, +deepspeed/compile/profilers/comm_profile.py,sha256=xM3vbFB7gYPdyVqTicuEqOCB4mW89TCR81j6kRVVbew,5103 +deepspeed/compile/profilers/graph_profile.py,sha256=UDzJGK7ENwoDzpB2LepW0Aa--CHmF4I3Ow5D5rdoSrk,11243 +deepspeed/compile/util.py,sha256=Zirqj_4obWpWc-wPi_Hrn5SU7UZ11vBIiSALu4WHKCE,15586 +deepspeed/compression/__init__.py,sha256=s3fNEEJJdLa4D4rEwtraP5axoSqpxxmh6LAIvyzfUTM,243 +deepspeed/compression/__pycache__/__init__.cpython-310.pyc,, +deepspeed/compression/__pycache__/basic_layer.cpython-310.pyc,, +deepspeed/compression/__pycache__/compress.cpython-310.pyc,, +deepspeed/compression/__pycache__/config.cpython-310.pyc,, +deepspeed/compression/__pycache__/constants.cpython-310.pyc,, +deepspeed/compression/__pycache__/helper.cpython-310.pyc,, +deepspeed/compression/__pycache__/scheduler.cpython-310.pyc,, +deepspeed/compression/__pycache__/utils.cpython-310.pyc,, +deepspeed/compression/basic_layer.py,sha256=oZvNusdxqnlhM-RsaALf5Q-4m2kQLcKAUUMuIKBtXec,36047 +deepspeed/compression/compress.py,sha256=DueechYg3bDNAIn2XW0L9nv8qN81acTPwimvD4usE_8,11886 +deepspeed/compression/config.py,sha256=5TTZNLB_GDULXEf704-dbNJo0ILL47DNVWHWiMGio8Q,25067 +deepspeed/compression/constants.py,sha256=Fc2681jhRfm0ADzNjFKjhYIPSXAI29gqeQSGW0FxLoA,5569 +deepspeed/compression/helper.py,sha256=lEFpurhyU7-yKWgO-wp1RojHECkqR8qJ16KrT4Hl06U,14637 +deepspeed/compression/scheduler.py,sha256=QwnyBtRhkEkufjyhns-Riz_-GBdUePzuI8UEQJOu7Ho,8112 +deepspeed/compression/utils.py,sha256=vLWAEOQHa1w61zmMyFQ3dwxbm-5kM_S-zOI_eQhYx5E,7818 +deepspeed/constants.py,sha256=K5l2jzMTXYfqyUBbO9QIMbe_1TFMxdJ9LJDEQ6HJLRY,867 +deepspeed/elasticity/__init__.py,sha256=Cmi3gEN48VYdPSg4zLOJ1KGKP2G3Xeruz9TykgC13m4,383 +deepspeed/elasticity/__pycache__/__init__.cpython-310.pyc,, +deepspeed/elasticity/__pycache__/config.cpython-310.pyc,, +deepspeed/elasticity/__pycache__/constants.cpython-310.pyc,, +deepspeed/elasticity/__pycache__/elastic_agent.cpython-310.pyc,, +deepspeed/elasticity/__pycache__/elasticity.cpython-310.pyc,, +deepspeed/elasticity/__pycache__/utils.cpython-310.pyc,, +deepspeed/elasticity/config.py,sha256=OCmF8fP8_dL3LkKiEJDtOREnJB3NYQuCaK48qRs-eFA,4703 +deepspeed/elasticity/constants.py,sha256=5Ynz57XbIFxeFkhC3JwanR-Dw0-U7qrsqKIH9oQo-Hw,2454 +deepspeed/elasticity/elastic_agent.py,sha256=KrPIykg6Q4K6BrjkgWmj0PRkgGHNuOb5vvgNjf7XZwk,8009 +deepspeed/elasticity/elasticity.py,sha256=oGJXhEj-VVdT2iUa-ei_gHXAfhP5e_VgtUlw5N500mU,17361 +deepspeed/elasticity/utils.py,sha256=-mbP_VX-SSstZKrEp0hkFwezql9pkAEvxx3BxMO2-ic,347 +deepspeed/env_report.py,sha256=0ZVv7RaHSaQIOJKWG5dqs6OfFalx6adq3NFZwB3nVIs,7491 +deepspeed/git_version_info.py,sha256=IZlPdw3IO5LiSJQ5ECQBVNnIj9TL3ezz-EEYfSR54ag,1075 +deepspeed/git_version_info_installed.py,sha256=XabhZhUkeGWlt88HDwsiXiUqQgraSOvvYg27wpRyljU,369 +deepspeed/inference/__init__.py,sha256=ZV8FRXVTJ3IGIDnyXJgJjRlCbB2jH_KJ1hWoTS4l_y0,267 +deepspeed/inference/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/__pycache__/config.cpython-310.pyc,, +deepspeed/inference/__pycache__/engine.cpython-310.pyc,, +deepspeed/inference/config.py,sha256=zBiw9UvgAQ79FfauGJPLpXwXgtxrrBMUwUdokM-TEww,11283 +deepspeed/inference/engine.py,sha256=2ndGZLN0dHv2RDDeaWj0do5o5OLrzutfqZMPDO6DsIA,30250 +deepspeed/inference/quantization/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 +deepspeed/inference/quantization/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/quantization/__pycache__/layers.cpython-310.pyc,, +deepspeed/inference/quantization/__pycache__/quantization.cpython-310.pyc,, +deepspeed/inference/quantization/__pycache__/quantization_context.cpython-310.pyc,, +deepspeed/inference/quantization/__pycache__/utils.cpython-310.pyc,, +deepspeed/inference/quantization/layers.py,sha256=s62KAh9pysdpMzV3KuaFeOdrezxLSsPcZi-Fmsj0w84,5729 +deepspeed/inference/quantization/quantization.py,sha256=-r1gbQC0AZe3umfsjla_2lp6opWrfI-0jfs6KtOaJ8o,4396 +deepspeed/inference/quantization/quantization_context.py,sha256=0IV04DTAk8PiOi07TGgABxFMkrKM4QFOqzRaT-RZO7w,514 +deepspeed/inference/quantization/utils.py,sha256=qu0HJHOSz020DT1UuMUo45c1IyARecLl_cCoMcDuf9c,11933 +deepspeed/inference/v2/__init__.py,sha256=9LrBSo2xdUl33MPzLedqmI13rnQOJJO-eT6fVsrt61k,283 +deepspeed/inference/v2/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/allocator.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/config_v2.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/engine_factory.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/engine_v2.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/inference_parameter.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/inference_utils.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/logging.cpython-310.pyc,, +deepspeed/inference/v2/__pycache__/scheduling_utils.cpython-310.pyc,, +deepspeed/inference/v2/allocator.py,sha256=A-yTKojaNr_O8COkDIFagTl4dd9aYJRIMtR7g8TIvJ4,1182 +deepspeed/inference/v2/checkpoint/__init__.py,sha256=6y_8HbNnkff-I3LYNyzZQ3DVDcy8CwjQAYykp_JiCw8,252 +deepspeed/inference/v2/checkpoint/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/checkpoint/__pycache__/base_engine.cpython-310.pyc,, +deepspeed/inference/v2/checkpoint/__pycache__/huggingface_engine.cpython-310.pyc,, +deepspeed/inference/v2/checkpoint/__pycache__/in_memory_engine.cpython-310.pyc,, +deepspeed/inference/v2/checkpoint/base_engine.py,sha256=9EmyYEwt1ZE861yuS-ayx9ykG2OzRi9p7_ua4hECMEs,1391 +deepspeed/inference/v2/checkpoint/huggingface_engine.py,sha256=sn6nmRLiOLEVxTCMHRExUbbIjMVavCAFyerbdxssz-M,5997 +deepspeed/inference/v2/checkpoint/in_memory_engine.py,sha256=4tEQQPAWl4BG4FJFdcgnoa8jahMQUtrXtcssTwRUSeo,1487 +deepspeed/inference/v2/config_v2.py,sha256=HT1FCeys_QnuUTNteBLQhJ8Y6Cg4WVW5B1HwAwQSUX4,1389 +deepspeed/inference/v2/engine_factory.py,sha256=ap6CwbdaOvD3IcCOCaqfhAvIRhp0lu0efZCLHTEMST8,6028 +deepspeed/inference/v2/engine_v2.py,sha256=2IiT1buZLPYL66oC8OQdpANhRSFjN-xjvnnCJh9Wqts,10491 +deepspeed/inference/v2/inference_parameter.py,sha256=rPeDJ4h58L4weNeIiFDIUrpVgDLpkQtzihUl_L-rkbg,2788 +deepspeed/inference/v2/inference_utils.py,sha256=V5gryIBXY9kVVkwPWukMMzQKbn0vH2hyDWnDhQ4jWQk,2378 +deepspeed/inference/v2/kernels/__init__.py,sha256=KHPoSEGybfJdU54HoIOCZzeIovB3EsfzaZuHv_i36V4,132 +deepspeed/inference/v2/kernels/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/__pycache__/ds_kernel.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/__init__.py,sha256=72TCKbfeUY6LbxtwVeNKljxXvxFrY-zobANphLq7nVQ,276 +deepspeed/inference/v2/kernels/core_ops/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/bias_activations/__init__.py,sha256=BpRq2zT5tKcsx4e7f8zcnDA5F6wEedmQXSI6idSEoTo,127 +deepspeed/inference/v2/kernels/core_ops/bias_activations/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/bias_activations/__pycache__/bias_activation.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.cpp,sha256=PdSuObELmrSu3lsYgfVEOO9tEEmexqO0an3_pP4nU3k,3084 +deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.h,sha256=q3vJjn6I7s_raM1Gy2TNnrY3nbskDDP9QyWwcDjpk6E,688 +deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.py,sha256=6uu2cWl65xBVWVHfZR9TfLiqrovUJOxOm50y2iZ9Y28,2437 +deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation_cuda.cu,sha256=dq_Bwd-UGv8bK3wWA1llNGVUEXO_rAMZInl5PcfpO4Y,4820 +deepspeed/inference/v2/kernels/core_ops/blas_kernels/__init__.py,sha256=ND9hDC4WOLfDXnlqe7FySprn8RkCXvzrwQKwYG8ZFRQ,123 +deepspeed/inference/v2/kernels/core_ops/blas_kernels/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/blas_kernels/__pycache__/blas_linear.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas.h,sha256=YtINP8OfWWiyZEx44zb1yylsXBHGBLdfxcb3xjBUNeI,5164 +deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas_linear.py,sha256=AQ4eOYQ3ejZspCeBHm87Q32f_Lf_ooauZIoVobdD_G0,2023 +deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas_utils.h,sha256=IGCBmqS3jASIg3Trt70BRfl2KIxtF2T8fQ2HxIOUJNs,11400 +deepspeed/inference/v2/kernels/core_ops/core_ops.cpp,sha256=4HbbHAOUmkaolemLrP_nWOdXSIbAP7KGgl1jTU1LMfI,1514 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__init__.py,sha256=tUarONsuAuUxyfogdH03MC3yebDySSrUEusyIEWNjrg,174 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_fp_ln_base.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_ln.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_post_ln.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_pre_ln.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_fp_ln_base.py,sha256=CmnqS4_iIbPWKbzNCKEv3wnng8UnRtyhSJiEz08cfPI,1302 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_ln.py,sha256=bp4CvANVf8hQQzyKOIScnR1OU7kB41uBNK5ANIydCvs,854 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_post_ln.py,sha256=cIJ9l9ey_fAdjuX00ZjooepUgvNEg_bc2ikSiR0MWK0,990 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_pre_ln.py,sha256=meb0fP8vvu1lFGlMR8iEus7YKlGP6FMlad9FJ7ZPg8A,1225 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.cpp,sha256=cWe1Z6kgQuPsrSTCOO1mWBtFY_CYqXHUt1HqhfiOyNs,4199 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.h,sha256=SCuiMaChAEgmE0t2kwfbqaYXQ1ycU1QXDVDJSHxtQ6o,2157 +deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu,sha256=NLCh_1U7k5iA5EOmmqISVMzoa47ypMYYZ5UAk_8tvMo,19730 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/__init__.py,sha256=0J6ih4FAYwgMMo0HSRWLTiBimI9hdeAFh9Ui5iwx2NU,123 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_linear/__pycache__/cuda_linear.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py,sha256=_2OxKHsY9P5mKOW3sT1ZTIy72d3BGpEjt7_Kb2bNsD4,6827 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/configs.h,sha256=XPl9PmMyD-OVFksGXN-eK67a-aoxJvjxNEBIXm3q6Fs,5044 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/kernel_matmul.cuh,sha256=e92X51umAswaFGMz9ntC58TRD1W33CjVZPS2eVeZFOk,14700 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/kernel_reduction.cuh,sha256=GQix8cchlsjHeTajB6oM08OEoEC1Mcp62LEsq3XKnAU,1494 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/ptx_cp.async.cuh,sha256=kdZEfq2Akuw1rSJMdR6DQ0k23mbYhRuQ4cKffK9_wUc,2357 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/ptx_mma.cuh,sha256=_-OJfWV4dhi9or_B1orbEkxBdZCNiXnmjW5-eNHtmno,5139 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_core.cuh,sha256=9RDNZ3HQFo0rBbFgaBJpM7JzxXQ4fYiUNdnlsZOfM84,10251 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_gmem.cuh,sha256=pyOFBGMAlcDFZx5iSS-HPnXA0Pqwx9oCvrUFV_EkioY,3207 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_paralleldequant.cuh,sha256=bqJJIaioGtHe1feD1vIjjWS1OpKGTrau1h03cCdkiaI,4057 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/weight_prepacking.h,sha256=JTrsVJdRovts8Q7bG7MlCzsVyKL_R76moU-zyakMPGI,8896 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels.cpp,sha256=XIFk-nXPumydZOOpupiRPrqgN6wbPRMr86u8_AAqcnM,8892 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels.h,sha256=i97f3j3gOIiFK4Ijoiu05U9aj0S4p2NfhBIIYe-hAbM,828 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.cu,sha256=l5rmu_QnbM1VNmao97uVySwP2sOGjjF919pX18_KXwA,13947 +deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.h,sha256=ON5qqFROLBv392C3nwM0adOlKtCHsjO59yYUgGQwz_A,1669 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__init__.py,sha256=S3Q1i1iwyLp6kWWJjD2kTinhJQ_ND4gKHejhuxQ8-74,171 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/rms_norm.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/rms_norm_base.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/rms_pre_norm.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.cpp,sha256=XqvW2ObN0dkRQIsqtWRGl_JYxkFjKAi-zjFMeEZVa3M,5873 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.h,sha256=FUf2hMir5TnFA6iU1Rx48DlnSYX__Yl0jkiNWARdREA,978 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.py,sha256=OkC86DHLQ03zshmFNUQO9r1mZ4mWD3gN0505IvzC1NE,770 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_base.py,sha256=cbZXvZSXRaD1MoTokRSRuw1anp_FXdGzg0moAxasQ1I,1311 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu,sha256=GrWivwyCSZ8gh1zaiJoKHh8sTcq8aot5m7XKx-Auw1A,10216 +deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py,sha256=qzME0FbrgV9IXbGCqIA0IP91rfFyRiTdiXwwQ39SufM,1204 +deepspeed/inference/v2/kernels/core_ops/gated_activations/__init__.py,sha256=2mIcBil301zqLsisa-B2C2EnTgQ2Zaz2QFRxPTD8eUE,128 +deepspeed/inference/v2/kernels/core_ops/gated_activations/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/gated_activations/__pycache__/gated_activation.cpython-310.pyc,, +deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation.py,sha256=XQXwTY3PN0qWIdZMBWh8De2lEGtN8syw7ItH_uN319g,2755 +deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.cpp,sha256=aog8VhJGL0F4bwUOZTR232KTgRJRavFFk3Xj_55shLk,3521 +deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.h,sha256=DpAhhx6lO3TNA8eid6Tct9qfDrNFweHRypmn4serMW0,776 +deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels_cuda.cu,sha256=0uLVwa_H3oSJczKK61AZCbBIUZF1mts1QGIM3yKUI6k,6080 +deepspeed/inference/v2/kernels/cutlass_ops/__init__.py,sha256=hpTBFM7draPD0dwxTHvGYgAdic-ysFMG6_XEr8x6V5c,146 +deepspeed/inference/v2/kernels/cutlass_ops/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/cutlass_ops/cutlass_ops.cpp,sha256=v6-PzCFZnmxFPXJm-pooeUISCq2V0S4E4_PAvg7oA6w,483 +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__init__.py,sha256=leIZbZ0995YwKwql6mU54HN_DJ1iBe67Cq_GdlXSMAY,122 +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__pycache__/mixed_gemm.cpython-310.pyc,, +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.cu,sha256=fDh2z9IwBffz5AI1NltOo4tgC2a2c0bwZ-leLahnR2k,4131 +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.h,sha256=laUi-GFh6qbLo0Ywiz_h7TtNIxroXSFwPQKc6yBGhE8,410 +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.py,sha256=gZ7ljMs2f20rcL1624vIhiLg_yemxkrYYFDoX-K1M3Y,2674 +deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm_api.h,sha256=soDC6vGW7UCYrnfw1EftpmOz6BuOi27Ihyi8J2Cfw2Q,1570 +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__init__.py,sha256=RRc_fJh0o2_l9wVb7Czf2GqTxLzvS6Rj_zeYUGzRjjI,150 +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__pycache__/mixed_moe_gemm.cpython-310.pyc,, +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__pycache__/moe_gemm.cpython-310.pyc,, +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/mixed_moe_gemm.py,sha256=pAaowuHqgwtCqv3tyof9jpbPxZxgLdcj1vPM2NgDcx4,2996 +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.cu,sha256=FRZSfEKLH3lHjvdbTbuEr1qGAi1spT6lhTLuIJ_OGuI,8594 +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.h,sha256=hdqzk6YOF68wO1DBVWDdtX5z-MIGHdtekijLU8z1cHM,740 +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.py,sha256=zX6btY3JLxP-6mgcjlNLSHlmgHJ2z3ZMH6PV3DQPfeE,2572 +deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm_api.h,sha256=RWmGN2LrcQukrsLoYaqykymCQJ1onmubdqXrHHgx2Is,1839 +deepspeed/inference/v2/kernels/cutlass_ops/shared_resources/weight_variant.h,sha256=1Mr74mw0SDFt0BsMgWiLvYXXaMWmhIgCUz1MZlWXYMs,278 +deepspeed/inference/v2/kernels/ds_kernel.py,sha256=R0n_uBCLCGAGlJa7lvep4cNF0DBxkD6TDKDFNAWpBpU,961 +deepspeed/inference/v2/kernels/includes/activation_type.h,sha256=WkFfkL5mgLyYAZimLUouThCkyA2aiF9LQlh1q206L9o,264 +deepspeed/inference/v2/kernels/includes/conversion_utils.h,sha256=augq8Zf-t46MaxWh0ryFb46_Guxud0k4kIK-h93rI8Y,12379 +deepspeed/inference/v2/kernels/includes/ds_kernel_utils.h,sha256=OKOMT46XItRnkrzqKHHK9n_ewxHh6AN-Uhv4aKixBm4,1297 +deepspeed/inference/v2/kernels/includes/memory_access_utils.h,sha256=wqpAqnfmSD5Fx15qwsnY63VCL6tYyJaChCqLxakM10w,33966 +deepspeed/inference/v2/kernels/includes/reduction_utils.h,sha256=189h93WF77CblgWonfW5eQqoTeJnb-Rt2Ov2fZ6mO-E,22352 +deepspeed/inference/v2/kernels/ragged_ops/__init__.py,sha256=8C6TomumccM_-p-iiyPyjZnDLPjTdQS8E4ydGQGNLu4,324 +deepspeed/inference/v2/kernels/ragged_ops/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__init__.py,sha256=djM0DTjGkr4vvNuhRIWV-2OI7lroepRMBiHIJtrcqTA,124 +deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__pycache__/atom_builder.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.cpp,sha256=AjIgvgoS7PUYGt4MW7oFu8QEcQvm3J16tyerW-dx8cI,2051 +deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.h,sha256=i232uHXwHjh2EaAF3709YNcTQ29uVDI2GQOs6GnYAsQ,732 +deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.py,sha256=QgVnamgI1qG8xjzu-9i5sdBccOuhgbTeAtufUtLNdog,1744 +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__init__.py,sha256=696rc8-r9nWCnfSG8UZ_fagI28Nerwq-nZCDvIIXPRs,125 +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__pycache__/blocked_flash.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/attention_atom.h,sha256=-ITylnAQx_20rXiWdPPldb149cGG-fNQUqOg8iHVugU,1062 +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.cpp,sha256=qaNovmlpqw7Cw_Y1ArGmH-7fB1nPQd0idvmC4nDR9pw,3912 +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.h,sha256=io7L6krERzOgiBIKNk3zwEppMDSvWtOqCY9nH35dB_U,459 +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.py,sha256=-aU8A7FbwrRqJPSHZEOef6Y3_PRndGKfCFT1c-3S1DQ,3830 +deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/flash.h,sha256=VJKtKOCufGTGt8ZGuPXfsA-dzsi87cqNTYX7pbqHKH4,1931 +deepspeed/inference/v2/kernels/ragged_ops/embed/__init__.py,sha256=y4HCIAaLuHexvTBcNppgiXqo76vHb_pAahL75kQBKY0,137 +deepspeed/inference/v2/kernels/ragged_ops/embed/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/embed/__pycache__/embed.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/embed/embed.cpp,sha256=iI-rjkQlCGnujjVzY1h8e6ccO24G7eVxwPLY5t_HKxY,4840 +deepspeed/inference/v2/kernels/ragged_ops/embed/embed.cuh,sha256=C5ZS1W9lnEDPBmWsuoJpRP-pgCYwLxqTK_OI-DQrT04,956 +deepspeed/inference/v2/kernels/ragged_ops/embed/embed.h,sha256=XbqCgqlx16bRtRLljUstzjssRJIGx7AixbbaofhILl8,700 +deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py,sha256=6lxDG2ymTyx2yGZ_9jDFyvcchbtD1iwKr2icLgHtmxE,3011 +deepspeed/inference/v2/kernels/ragged_ops/embed/embed_cuda.cu,sha256=5JfvhMBcKAPXQbF4WvE3Ljmolom8ArjvkNtprx1SaKo,6055 +deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h,sha256=T0LJvEtOulSPO1r0YA4jEHW4l-3XSYyS8nuI6Ok2F6E,782 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__init__.py,sha256=ou-mEvP9FMj89HvxswZbtTvS_T8Gcy3zf16jPcb_Urw,208 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/blocked_kv_rotary.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/blocked_trained_kv_rotary.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/linear_blocked_kv_copy.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cpp,sha256=qUdVHb-2UMnlH8Q6zZJLUW69T75AMwfqEZjco2EXbdY,8110 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cuh,sha256=X9w1u75CiGNEOV0uaXAmjX8usEiTXFTZG7EK3ulvZDc,1676 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.h,sha256=5RfNKDify3eHIV6GjfqNozT6_lc2shcwNidgrQ6AJD0,2406 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py,sha256=spVrPPi6tMh31Zjb3l4UooFgJ2YTFq9UJI2f50Mziuk,3142 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu,sha256=ztJXnG_ik8nCCJQw2d3P0ScuNmTUPvCKQT1I3L28EzA,19275 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_trained_kv_rotary.py,sha256=FGU4a0Wcb4NtsbJiLswLnd6-M89yE800UW19v5kFp9Y,3220 +deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/linear_blocked_kv_copy.py,sha256=mAa-291Bzc06yNGZVgsYPhrpq7dVoYDkWDHFUr2i1fY,2957 +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/__init__.py,sha256=-6rVFXUnZ3U4MBA1VNuKQfhcdjlWgc_xnufXpnFfGyg,125 +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/__pycache__/logits_gather.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cpp,sha256=wItdeUW_dSJzfePXbRKjYeggcfGG_l8yO7OAPXF4IeU,1800 +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cuh,sha256=OQlrhwuHJIwGfhbozWlM4NKhYoKBtAL-3nc0vqATdZY,624 +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.h,sha256=bRw1hkiDAAqkP9FkNwsL1QZ0UeGSx1_pYHOY4yoTF5k,571 +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.py,sha256=kGPLmxKKqOh2x35QWtCpeXLXxoG15kxmE_Z_-cd7Yv0,2121 +deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather_cuda.cu,sha256=Dd2HJJHpPlDkt2C91667Fn1GVuc4jonsIW8s1vPKhHk,3312 +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/__init__.py,sha256=K3BZ4_CDJGkNmvFCeYvIEBY2DRMWqq5CzCL7HAStX8Q,122 +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/__pycache__/moe_gather.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cpp,sha256=zWULPg_6e15LzTk5q2DJJcsgwT4zoLhTNIc0ClJf73w,2392 +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cuh,sha256=orang8iBcPNQ0BjFTeiBWNX09UjDq82Sh8tyye-5xyU,706 +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.h,sha256=MfJuJ2rg_ba6KvIm1LNbTTHMBH8MYQHgd7JeYJSG9Q8,561 +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.py,sha256=31zC0Sb3caA12JdhT7qJoawCUTkKB5PjZpFupQI6FvE,2265 +deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather_cuda.cu,sha256=oy_bzQhyMfmjwiXag0u4YKjyKHaLXMDWYHlJQnQZNWM,6679 +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/__init__.py,sha256=e3cxpdOKqIBYpxvNmMnj6Ms883dqdsl4B_fXdqTnlnM,123 +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/__pycache__/moe_scatter.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cpp,sha256=ywodf_RYumTot90MnTLb4o1h7HDWgiFuVTuhHwxNooE,2858 +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cuh,sha256=M6t_ydSLYKsekCym5m07AGDu8-U_qVG0hLuVH_nrO-A,772 +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.h,sha256=pBdy99sow7TKoB3eV3gdPbCB4_pGNx48K-Eo8C8XjHs,701 +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.py,sha256=7c_jFxhoJINWl57LIhFZwvlr08W8Nbdk7ULSFmC8Yqc,2504 +deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter_cuda.cu,sha256=US5_PjVcIqYcqR2DDKXUIJAMfdXKI0ElfbNCUmLscbw,8070 +deepspeed/inference/v2/kernels/ragged_ops/ragged_helpers/ragged_dtypes.h,sha256=OVEob4D9LFPkyxA4PscOrGV5BzBalfiMrTrkcifCvL0,940 +deepspeed/inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.cpp,sha256=QaWqjTz4kXXR804xXlgNBTOxtOtLSP3NcGDNHuFqqQI,1149 +deepspeed/inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.h,sha256=_7Xb4l1mfYbSz7i2HAZ8TqQSLJy6TBBC51OMaarJXfM,563 +deepspeed/inference/v2/kernels/ragged_ops/ragged_ops.cpp,sha256=Bnku0tlvBSyPIwOtx7-kvIGg3OPGmsgjMO72rLABEkQ,1472 +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/__init__.py,sha256=N0qSPddtzqj_UkISW7rxEVXpdk8a5JItbXBjNAY_FBE,139 +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/__pycache__/top_k_gating.cpython-310.pyc,, +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cpp,sha256=EBxWZ_Nm2ygSbFtnzI7kGMsl3nZxECXXefAVjMDDVx4,2435 +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cuh,sha256=1U1wNpJPmIhCiSwRi7xeRSR_zCcGnuXbFwBOz_gQplw,751 +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.h,sha256=FdryRq5nESxX9HgzKIhnUgwfaoOzLir5O361gf-4g_M,556 +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py,sha256=-xgNLLmp6ea_QIMkGtuYDgt-fLfvq6wPVKgMWQpIHfA,2578 +deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating_cuda.cu,sha256=LCuHhCZhWYGT69EW4mq3sPYr9JafW9Ef1YIScTNCQC0,5130 +deepspeed/inference/v2/logging.py,sha256=H4Dp4vKNi42HZaNgXOAOG-Vki2oC3U3APVQwq3ay8kE,780 +deepspeed/inference/v2/model_implementations/__init__.py,sha256=g4p-wyfxcks_kDpGydBi2Sho_dI46rEfin8-xAydRd4,577 +deepspeed/inference/v2/model_implementations/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/__pycache__/flat_model_helpers.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/__pycache__/inference_model_base.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/__pycache__/inference_policy_base.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/__pycache__/inference_transformer_base.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/__pycache__/layer_container_base.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/__pycache__/parameter_base.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__init__.py,sha256=_mp6NDvxO-3TM9LiiOTrMBRIweeiRjU6s6L-GpTvX9U,359 +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/attn_output_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/embedding_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/invfreq_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/mlp_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/moe_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/norm_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/qkv_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/unembed_parameters.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/common_parameters/attn_output_parameters.py,sha256=VwOj8hRbIQQnWfoHVq-vbRK_QLt9f-uimARsTHfot4w,763 +deepspeed/inference/v2/model_implementations/common_parameters/embedding_parameters.py,sha256=k7CMnQ06UchamvQXsdeiTfxAkjpor3qZjTU0cD20hZU,605 +deepspeed/inference/v2/model_implementations/common_parameters/invfreq_parameters.py,sha256=Srz4b0mFK73OBzsw3EK3Aydj21Z1gZjLuyC8qfry6QI,401 +deepspeed/inference/v2/model_implementations/common_parameters/mlp_parameters.py,sha256=fMdVvya2UlX-DJ52QT0Mv8YKFukIAWx4F833ZoI0-IA,3363 +deepspeed/inference/v2/model_implementations/common_parameters/moe_parameters.py,sha256=hTKZ8zIERsvQ6i6zQRWe47qUcJmMAGOHfJinwpY3r04,2548 +deepspeed/inference/v2/model_implementations/common_parameters/norm_parameters.py,sha256=xTaLT5H81-Et8y9FivEaY8N9VQ0XDsvtlV4v1wnFUG4,454 +deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py,sha256=pxBm6EBMNLvbJTPAY9ABjdQ68-oxbY_nWb5CFGDdvdA,4086 +deepspeed/inference/v2/model_implementations/common_parameters/unembed_parameters.py,sha256=-gdVA8UMMvFsd-lSAvCbQp5mQ25j1sT5f7KC2mSVY_A,651 +deepspeed/inference/v2/model_implementations/falcon/__init__.py,sha256=IVA-W-wyht49JDQt7MTTsBMN4dHAAQxWPEULVVwkoVY,129 +deepspeed/inference/v2/model_implementations/falcon/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/falcon/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/falcon/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/falcon/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/falcon/container.py,sha256=TNdqevl_yh4g06l80i3RJC88dExK2bfd_22NSA-yz0M,4479 +deepspeed/inference/v2/model_implementations/falcon/model.py,sha256=ELmk5fmti6cUmUqYRcsnykS746An6eEibPf0G3IQsgs,7768 +deepspeed/inference/v2/model_implementations/falcon/policy.py,sha256=oFyIvfzFEistRJCzXS-wPtFfUXDPqt9ANhPa7ywrKSQ,1342 +deepspeed/inference/v2/model_implementations/flat_model_helpers.py,sha256=XT6oMqJAfZhYLCr5KrPBxSFVW-WKnFjOi3FfEElshP4,9772 +deepspeed/inference/v2/model_implementations/inference_model_base.py,sha256=3YVtQafVIrZMI-kZadRKmQklafI5fjc38GTdTZeM3jQ,10083 +deepspeed/inference/v2/model_implementations/inference_policy_base.py,sha256=XKH3GcOtP8WAeAopwuQACflJA3cvX4F2MGybrIqcBPc,9481 +deepspeed/inference/v2/model_implementations/inference_transformer_base.py,sha256=SgXRyFYWj31DJMbDpgo06nEyAJLD-w3qKJF1xbAm9Jo,23592 +deepspeed/inference/v2/model_implementations/layer_container_base.py,sha256=qvvGjZ0aL6Ah4MDtjpowV76DKSpI8VuZyY_4hZRbB1o,15471 +deepspeed/inference/v2/model_implementations/llama_v2/__init__.py,sha256=J8aLmbt1KDVL_LOtSqaP-xACSxlgIDjeaIA9R-Grsgs,129 +deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/llama_v2/container.py,sha256=blAioiSrdB9kH0zHFV0ZO1kWChv6-geaZW1ivEaSQx8,2664 +deepspeed/inference/v2/model_implementations/llama_v2/model.py,sha256=YmauqGYxtx2o_apo28aYvEhxi3SmYYvLm1IoEE7tGgs,7515 +deepspeed/inference/v2/model_implementations/llama_v2/policy.py,sha256=A1z_j19nj8yOURNNC29DeSyMatAjzuPUtcnIz09zZII,1148 +deepspeed/inference/v2/model_implementations/mistral/__init__.py,sha256=rzXGZBulBvo8ZNt7E5z1mSi8OlN1edNDTFlLCMCqo3M,130 +deepspeed/inference/v2/model_implementations/mistral/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mistral/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mistral/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mistral/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mistral/container.py,sha256=0THLABZ6g4W4H_Kf2INGnyVrXHroT2uveKr-F2deP88,2784 +deepspeed/inference/v2/model_implementations/mistral/model.py,sha256=QhDoAoi2aMguJf4GJXqNyivqRKAjJLUwf5EsVgYHl38,7350 +deepspeed/inference/v2/model_implementations/mistral/policy.py,sha256=ceBkx-EOWH6UmcosmmhcE7anO5cDMaPGhLInK-x3PBA,1056 +deepspeed/inference/v2/model_implementations/mixtral/__init__.py,sha256=ywKUfNNkXZuDaMa9NT05H3MkYJQgiCWl3eM8kZMkhJE,130 +deepspeed/inference/v2/model_implementations/mixtral/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mixtral/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mixtral/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mixtral/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/mixtral/container.py,sha256=YMdy5flaitAQHdy71VBCi7n0vQBfWcPRyXTfZeABchA,1658 +deepspeed/inference/v2/model_implementations/mixtral/model.py,sha256=smAqT1nkiFmn2ipNNvJqd2BiIHHprI5O7iWKOFsKM4M,9096 +deepspeed/inference/v2/model_implementations/mixtral/policy.py,sha256=u1Am_kWG7UvEt8A5V4lTn_BjUZ4eSlnTB32R5agIYy4,1057 +deepspeed/inference/v2/model_implementations/opt/__init__.py,sha256=WeJqbtuxqlXaqif4F_4KpZnGBCwXwKEn0OwUZOcxmpo,126 +deepspeed/inference/v2/model_implementations/opt/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/opt/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/opt/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/opt/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/opt/container.py,sha256=1jIU2Icv-lSEi_XMhaS4jGUVWFfbP-B2t2xf52eQVwY,3415 +deepspeed/inference/v2/model_implementations/opt/model.py,sha256=3FFAoq-iBo_F5ep0guXKXk6d5BnNOx0QZFiv2VILScs,7262 +deepspeed/inference/v2/model_implementations/opt/policy.py,sha256=9cqEPs_pM9gdzsBegbXJ7GHEgHfmD6n_AaRpjMH1yGE,1066 +deepspeed/inference/v2/model_implementations/parameter_base.py,sha256=FioZn7e4Gw2MUq_ov4popyXxD-dle2X0r7dYagAXmvg,9026 +deepspeed/inference/v2/model_implementations/phi/__init__.py,sha256=jKL7LDi4qrR3X3S6rcFpkls2K0wXVV6Qb3mIkiOtRKk,126 +deepspeed/inference/v2/model_implementations/phi/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi/__pycache__/containers.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi/containers.py,sha256=L1S3qe2iHCv73EtWzb-IPn_dg8rxzwUn27HPXMatmCg,3197 +deepspeed/inference/v2/model_implementations/phi/model.py,sha256=b12uXjTpbHy_0EklvawD67KxwkrV-ooYSSdjNxGbz58,7251 +deepspeed/inference/v2/model_implementations/phi/policy.py,sha256=FIcCwm7nwUQOoFT1yGV_OTFd2xldUWmL_TqmpjtjUgU,1175 +deepspeed/inference/v2/model_implementations/phi3/__init__.py,sha256=97ojikdkWIzS9unY8jgler8jl3vaOBjGXpifKCtPXxY,127 +deepspeed/inference/v2/model_implementations/phi3/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi3/__pycache__/containers.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi3/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi3/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/phi3/containers.py,sha256=hn-rEguXEgYbcu1GSYH7Pc4e6bOFTBQeEbGomMhmG0Y,2464 +deepspeed/inference/v2/model_implementations/phi3/model.py,sha256=0MtLieFL5uPP_eHcKdsAe9Q48ONQIWtPEYuB6OND5wQ,7296 +deepspeed/inference/v2/model_implementations/phi3/policy.py,sha256=ytHLbsKAqQN27Lku715CNU5bMg71EohF5W2-4P7v0lU,1033 +deepspeed/inference/v2/model_implementations/qwen/__init__.py,sha256=1TgzvNCP0loWbohqWRMH9NXCKKjXVVagBtatz_NCHc8,127 +deepspeed/inference/v2/model_implementations/qwen/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen/container.py,sha256=4fnYCNY92cDw116OyEVhmp_fuFtfeeyq4ZG4T5NHVi0,2362 +deepspeed/inference/v2/model_implementations/qwen/model.py,sha256=cCoVWHR2t_9LrLSwYE2LTCJa6nxXj9-l-Qr5q819Iq0,8112 +deepspeed/inference/v2/model_implementations/qwen/policy.py,sha256=0nP7dboH8FSnh6xvImRKK0gVWkw2iOTImwCJw73V-OA,1066 +deepspeed/inference/v2/model_implementations/qwen_v2/__init__.py,sha256=s0OcwPN9W24-8vYSQ8kztzky9jEwOvER9aMR60Xejuk,128 +deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2/container.py,sha256=HpDpzng1cMtevhB-XykEKvmPPXENaDRV6BiLVYM7wjo,2817 +deepspeed/inference/v2/model_implementations/qwen_v2/model.py,sha256=EENwSZm-DZRQrv64UCLoEHop_U8xZ9yVejtboPzri_8,8029 +deepspeed/inference/v2/model_implementations/qwen_v2/policy.py,sha256=U9rMVvegwkKMqw6SoFXGZXobYTrjfpfySUryrpF0TDs,1140 +deepspeed/inference/v2/model_implementations/qwen_v2_moe/__init__.py,sha256=nxhwmjwG8giEHatkVAzPoQDEzgLgXCOcxluBG7Y4cn8,131 +deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/container.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/model.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/policy.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py,sha256=QK62BZaTWjQACeHUJklh_qVJgRbX_N9WGhaOAtY-8dQ,4087 +deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py,sha256=bcL2sE2G5mHIPk1EZHT3vwFqNEiaBNmosSsKm1d1bjY,13516 +deepspeed/inference/v2/model_implementations/qwen_v2_moe/policy.py,sha256=k_HNyspBw8UG9mjCnCilY8khGf5dp8TI_74Du7eSwCM,1064 +deepspeed/inference/v2/model_implementations/sharding/__init__.py,sha256=M0NsMjUqO0Rj0PAPm2lyRepH5lc3uFQswNM5Gr3QAcI,247 +deepspeed/inference/v2/model_implementations/sharding/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn_out.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/embedding.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/mlp.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/qkv.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/types.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/unembed.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/__pycache__/utils.cpython-310.pyc,, +deepspeed/inference/v2/model_implementations/sharding/attn.py,sha256=QAapGaB4Bp1k9xI53-QhhyKFoaeIPandzI1qnPT2rxE,2376 +deepspeed/inference/v2/model_implementations/sharding/attn_out.py,sha256=MWf6O0P55Am0x947474ZqGcOEuZn2nGptKcNPCxfS0c,4909 +deepspeed/inference/v2/model_implementations/sharding/embedding.py,sha256=mJgZR289D6_ODuTZm4pVsedyJEHaXH2a4ncI8_7YdG0,1270 +deepspeed/inference/v2/model_implementations/sharding/mlp.py,sha256=UIpBwlbU9xDNv1iDbIfcxqPgcAf_WvO42NxtMGyZ1QQ,2895 +deepspeed/inference/v2/model_implementations/sharding/qkv.py,sha256=pQfXbEPeqEZ9eCJwcdSvJsid5jeBJ2jptWhqVaPJSPM,7609 +deepspeed/inference/v2/model_implementations/sharding/types.py,sha256=YqDHnyDLGhy2sf37n-a0eHIIXqcQ8OjHKHQ-rcKAl8E,575 +deepspeed/inference/v2/model_implementations/sharding/unembed.py,sha256=1AmukReQ6SAfFlq2YdiQwJrrvPLVx7fhel7KKCGXXjQ,1625 +deepspeed/inference/v2/model_implementations/sharding/utils.py,sha256=wKlbKIwSyblinL1xupsHDXodws2zfWhxt-flmac1grc,5066 +deepspeed/inference/v2/modules/__init__.py,sha256=oZhP5-VwkgPMtHVFO_9MefXPnSEfcHdQeIXaRg3VOeY,193 +deepspeed/inference/v2/modules/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/__pycache__/ds_module.cpython-310.pyc,, +deepspeed/inference/v2/modules/__pycache__/heuristics.cpython-310.pyc,, +deepspeed/inference/v2/modules/__pycache__/module_registry.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__init__.py,sha256=IJ5d0k84imLnM7jh0UcZ9k3_ITPU7YI_6ZoGRpEJcqs,449 +deepspeed/inference/v2/modules/configs/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__pycache__/attention_configs.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__pycache__/embedding_config.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__pycache__/linear_config.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__pycache__/moe_config.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__pycache__/norm_config.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/__pycache__/unembed_config.cpython-310.pyc,, +deepspeed/inference/v2/modules/configs/attention_configs.py,sha256=IWbmcmCc3CdL1E4FEUhrN9xU9wfbEaswUD8a0cM_D3Y,2834 +deepspeed/inference/v2/modules/configs/embedding_config.py,sha256=PuRrpOWLwtqjZRg9jumeMHuqzzsTcrbCjRxKL2L6xww,1812 +deepspeed/inference/v2/modules/configs/linear_config.py,sha256=fz9m8RIL6NVjbkq06-0WSP4-7IxIdSVQi1nRY4WQdTc,1301 +deepspeed/inference/v2/modules/configs/moe_config.py,sha256=EqX1OLq_QGt4IrSGhAJUXfVvy0MqjZk2E3toeUkb7kA,1289 +deepspeed/inference/v2/modules/configs/norm_config.py,sha256=iTKxfBxWcHs2JiiuxNZU356DGJR-a0OgDW3p5oshedk,839 +deepspeed/inference/v2/modules/configs/unembed_config.py,sha256=2E9XPh8VSRJckwng_DwbQksLMXtWZMFiXUjo5GWhzqI,822 +deepspeed/inference/v2/modules/ds_module.py,sha256=1h8X2TvYvDsUyEUjllPlziT7XOqwRxP02GBagzJw7FU,1856 +deepspeed/inference/v2/modules/heuristics.py,sha256=eRitZ_RX2yg4KOgmjiJaTEDZfH-zauMDIaf3WVQ2rGc,7978 +deepspeed/inference/v2/modules/implementations/__init__.py,sha256=0rEv4jixmtnw_sIe_Y5UKaIEcG7qwALUSzx1GNORTdk,289 +deepspeed/inference/v2/modules/implementations/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/attention/__init__.py,sha256=tNj2_fpavmUWaI5VAkUF0I7k0OxmIlZRNZZl7JIqrR4,157 +deepspeed/inference/v2/modules/implementations/attention/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/attention/__pycache__/dense_blocked_attention.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/attention/dense_blocked_attention.py,sha256=JgS7E2TzKmWylwRYWiYR-F7ldLdg8qhGr8Q9bXSPoCA,7735 +deepspeed/inference/v2/modules/implementations/embedding/__init__.py,sha256=CUiHnEi77BVJmH3aoJOanq5mTJIvEcIo3cHcbkddDHs,144 +deepspeed/inference/v2/modules/implementations/embedding/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/embedding/__pycache__/ragged_embedding.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/embedding/ragged_embedding.py,sha256=zddQlHcpjLe2_dKQUTp_5VfJwTP6L-BRk8HUHD2vPRo,2740 +deepspeed/inference/v2/modules/implementations/linear/__init__.py,sha256=FXT-bP-U4YlbMbBDrgwkpMPTHps3NiztMczMulIgmPU,203 +deepspeed/inference/v2/modules/implementations/linear/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/linear/__pycache__/blas_fp_linear.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/linear/__pycache__/quantized_linear.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/linear/blas_fp_linear.py,sha256=nHZHiLKVcppWyn0hNrQq4sHrS9UYKlHunrwL92DUVkQ,3629 +deepspeed/inference/v2/modules/implementations/linear/quantized_linear.py,sha256=5CQEpEBNq-Tzyt5dH0OK3RzZJUWjAxBieqzuC2OFJgw,7772 +deepspeed/inference/v2/modules/implementations/moe/__init__.py,sha256=XP3Ddz7eBtLvpIGum9u3gzUhl61w0W_jNfopcMI3QRE,143 +deepspeed/inference/v2/modules/implementations/moe/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/moe/__pycache__/cutlass_multi_gemm.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py,sha256=OzMYYHO7YIT5zOrAqfndOohQgJGwM4oHiPT_1mJU4EU,11126 +deepspeed/inference/v2/modules/implementations/post_norm/__init__.py,sha256=4hE4KqEcl74cOWm1kvexksZwHMK8rk-OQMzJYkdmgQo,141 +deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/cuda_post_ln.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/post_norm/cuda_post_ln.py,sha256=DNHGpJ2wNC8P6vyyk8M8hDpuXpA7k-ni78W3olRKO18,2086 +deepspeed/inference/v2/modules/implementations/pre_norm/__init__.py,sha256=_yJ5_1apmG5KriWHdy9klvL5H0vFzmLyWnZf1VIrmrA,184 +deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_ln.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_rms.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_ln.py,sha256=VvFlm3Xc6U7VVyfXnFUeA5tsLw4TLhK_LVlKJkN4SmY,2744 +deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py,sha256=WxZeP0B1ghhpOiMIHig4vb0ou0q7jkiLZFfbbdREXtQ,3245 +deepspeed/inference/v2/modules/implementations/unembed/__init__.py,sha256=ve_166eLGPGfIncTIvEvqoFGKi2pvtgvSLVVvTPXpHQ,140 +deepspeed/inference/v2/modules/implementations/unembed/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/unembed/__pycache__/ragged_unembed.cpython-310.pyc,, +deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py,sha256=cAd3By6SBykoCHJI3oxwE0FXHcMP5HEF5CSrWJyBjjg,4921 +deepspeed/inference/v2/modules/interfaces/__init__.py,sha256=2djwkw0m2KAAuBV-LuWY0W0ncwIrCxDX5uwf4oc8E-U,519 +deepspeed/inference/v2/modules/interfaces/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/attention_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/embedding_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/linear_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/moe_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/post_norm_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/pre_norm_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/__pycache__/unembed_base.cpython-310.pyc,, +deepspeed/inference/v2/modules/interfaces/attention_base.py,sha256=Kl8aPbpzyhP4bF21LZ_OcEwHFoI1KFirOiukiV4YajE,3613 +deepspeed/inference/v2/modules/interfaces/embedding_base.py,sha256=3pXB9F4IlWMqFGV8ZYE3k3bG2Je4B9uvVnVS0pRzLeU,3071 +deepspeed/inference/v2/modules/interfaces/linear_base.py,sha256=8g8or7xUMtGX_3uQj6zt0YfUBS9SwSbRY02lrr3gUVE,2146 +deepspeed/inference/v2/modules/interfaces/moe_base.py,sha256=5BCoKWZplUfxMj7mOPA58b4aBGk7APM-WG-ByxUIUyA,2850 +deepspeed/inference/v2/modules/interfaces/post_norm_base.py,sha256=PktjhFD5Pi3y2zt_VEwB9L-yNVyIo4dHpJgdVwgbcVc,2123 +deepspeed/inference/v2/modules/interfaces/pre_norm_base.py,sha256=QlYoS3N4N3ZKfUJk6C9u_vCT10p4-ACVgAEydLquHnY,2106 +deepspeed/inference/v2/modules/interfaces/unembed_base.py,sha256=W5ToOG46DZwH5e1VnFXxXp_f9-EvnO2iIw6zqZkL38w,2055 +deepspeed/inference/v2/modules/module_registry.py,sha256=UyDd-5FYQwbEJ7ihPZ60GkEiQ7M6Z2XMWnEAaLZN_WQ,2073 +deepspeed/inference/v2/ragged/__init__.py,sha256=2fFLZ35zp9WgcDaQ3XaTti8Hzi99zt8EgueCaaOqZs4,418 +deepspeed/inference/v2/ragged/__pycache__/__init__.cpython-310.pyc,, +deepspeed/inference/v2/ragged/__pycache__/blocked_allocator.cpython-310.pyc,, +deepspeed/inference/v2/ragged/__pycache__/kv_cache.cpython-310.pyc,, +deepspeed/inference/v2/ragged/__pycache__/manager_configs.cpython-310.pyc,, +deepspeed/inference/v2/ragged/__pycache__/ragged_manager.cpython-310.pyc,, +deepspeed/inference/v2/ragged/__pycache__/ragged_wrapper.cpython-310.pyc,, +deepspeed/inference/v2/ragged/__pycache__/sequence_descriptor.cpython-310.pyc,, +deepspeed/inference/v2/ragged/blocked_allocator.py,sha256=8rwqli63wrLAbzOuLoOgVqLi5UlpPlyxXtyRS1gU6zM,3661 +deepspeed/inference/v2/ragged/csrc/fast_host_buffer.cu,sha256=8U10U6VZk12oGwXCniQsqduta_cEcVBFuN9Fq3LO1JQ,501 +deepspeed/inference/v2/ragged/csrc/ragged_ops.cpp,sha256=Z6na7zRDPD7PrYbZ4lBhBsGyMatZ9yUNYw_PBzRGHw8,2726 +deepspeed/inference/v2/ragged/includes/fast_host_buffer.h,sha256=N73g7w1uE1gEsyWzrEnEibGeU8nBJPDeL-SWJEWK_0M,302 +deepspeed/inference/v2/ragged/kv_cache.py,sha256=eekwPz2OMNgQJHYXcxIrW86yK0guS419GeNqXJ38wIg,8559 +deepspeed/inference/v2/ragged/manager_configs.py,sha256=_SDSRv5VC5ZdmN6VDckjHtuC4HdJK6KjH9Vec839WcM,5965 +deepspeed/inference/v2/ragged/ragged_manager.py,sha256=Mk09dkqkEj5CMA9kWvTWuQeiw7fWiO9xn_nMIs-1Hho,7334 +deepspeed/inference/v2/ragged/ragged_wrapper.py,sha256=RJJpqSVz6vsElbtfePx99TDYZ2Q_Hao7qt_2Jt6RP2s,12701 +deepspeed/inference/v2/ragged/sequence_descriptor.py,sha256=oo1kf0upitTJkbZB6ME7GY3_MoOnGABXcw1coq49ljM,10891 +deepspeed/inference/v2/scheduling_utils.py,sha256=S1MPmx8zU8qHmDd_eDLhKP6hh71N9lo_Q8kJxO1AyOY,1363 +deepspeed/io/__init__.py,sha256=948yAWx0rWW6wzxU7f2JTmwF88rNisJdS1FoPs_JNZc,249 +deepspeed/io/__pycache__/__init__.cpython-310.pyc,, +deepspeed/io/__pycache__/base_file_writer.cpython-310.pyc,, +deepspeed/io/__pycache__/base_io_buffer.cpython-310.pyc,, +deepspeed/io/__pycache__/constants.cpython-310.pyc,, +deepspeed/io/__pycache__/double_io_buffer.cpython-310.pyc,, +deepspeed/io/__pycache__/fast_file_writer.cpython-310.pyc,, +deepspeed/io/__pycache__/mock_file_writer.cpython-310.pyc,, +deepspeed/io/__pycache__/py_file_writer.cpython-310.pyc,, +deepspeed/io/__pycache__/single_io_buffer.cpython-310.pyc,, +deepspeed/io/__pycache__/utils.cpython-310.pyc,, +deepspeed/io/base_file_writer.py,sha256=nXR4OuhLiHga6-Colhu578n97LP8sEkXCwKfFKOZ5rw,1007 +deepspeed/io/base_io_buffer.py,sha256=C-FVSF4w6wc7lCZ9VXkRhwBnW75TUE-n_tTko4YplnU,1841 +deepspeed/io/constants.py,sha256=9nxuv0_QNaB_1CN9d9FXvlEPrZ2fEHFL3yuo8eoshZg,806 +deepspeed/io/double_io_buffer.py,sha256=6qRY3Je1o2onzvXIP50KZdfJfzraYzBCBlRJoC-K0xM,3028 +deepspeed/io/fast_file_writer.py,sha256=lh8RGJc9JAKKNrzKkyOh28DCpLoh_DHqUatNbC-WJIU,11049 +deepspeed/io/mock_file_writer.py,sha256=pWXxx5o7DyV1xKsAOiWoSUH9ZIYDIlvRtlS8JHTkj-4,1638 +deepspeed/io/py_file_writer.py,sha256=knJ292prjk4AdfbbSiDEK0DGoBsZU6rMG_RYqo7KcAU,906 +deepspeed/io/single_io_buffer.py,sha256=CeUIj3GO-CnfTWHA9wrqe89W_PRjOF7GjXJYuCE36i4,1324 +deepspeed/io/utils.py,sha256=D0IOzZmhByrqOBQxP5m3KvA2HGanM3tV0OxZXQyMyas,1464 +deepspeed/launcher/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/launcher/__pycache__/__init__.cpython-310.pyc,, +deepspeed/launcher/__pycache__/constants.cpython-310.pyc,, +deepspeed/launcher/__pycache__/launch.cpython-310.pyc,, +deepspeed/launcher/__pycache__/launcher_helper.cpython-310.pyc,, +deepspeed/launcher/__pycache__/multinode_runner.cpython-310.pyc,, +deepspeed/launcher/__pycache__/runner.cpython-310.pyc,, +deepspeed/launcher/constants.py,sha256=X-6bI3aC4NXSZgrMJ6LaIILrsWu_TNteNBwgKOhNrgI,375 +deepspeed/launcher/launch.py,sha256=aHutEg7kj_kY5VdH_CQLclyYOIZoh2YSRLD-wXvM_BI,14898 +deepspeed/launcher/launcher_helper.py,sha256=tpvx6ItaLbBNYnUwP8XoKcn8P0mFXRffpB1Z_uLB_U0,3931 +deepspeed/launcher/multinode_runner.py,sha256=QcTvKOnkP94gmjLrK9om9Upw8FKClAmSLyCEtKQFKjU,18797 +deepspeed/launcher/runner.py,sha256=3u_dE4Q1e0bTmFjwk85JVRRWwbRmnzYmnloS0crJS_0,26306 +deepspeed/linear/__init__.py,sha256=fwQSoiw98NsTIwp6BuhtNMykJsaJOpJWd5QAUpitfcc,238 +deepspeed/linear/__pycache__/__init__.cpython-310.pyc,, +deepspeed/linear/__pycache__/config.cpython-310.pyc,, +deepspeed/linear/__pycache__/context_manager.cpython-310.pyc,, +deepspeed/linear/__pycache__/optimized_linear.cpython-310.pyc,, +deepspeed/linear/__pycache__/quantization.cpython-310.pyc,, +deepspeed/linear/config.py,sha256=U4yr-hEialc28KiEp4qRfm-8H_SF5xsLij2DROOTGSQ,2178 +deepspeed/linear/context_manager.py,sha256=L_aADkH1e59C7LCpoS7XyxC3rJ3QBSBqXucrVOjy1ss,3574 +deepspeed/linear/optimized_linear.py,sha256=RcpH8q8QpB7hE85DqU_evn29xVxKcWxOYzvx8aiUaUQ,10474 +deepspeed/linear/quantization.py,sha256=Lpvwk_-yb3EQpsNQC1sFT9fGvlj6dPdQx5fg0Rxb0s8,6352 +deepspeed/model_implementations/__init__.py,sha256=jJCJcnbvmFc9eQxIUIb7Gd6Ysd8ijSuuk-aaqFrDkK4,220 +deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc,, +deepspeed/model_implementations/diffusers/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-310.pyc,, +deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-310.pyc,, +deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-310.pyc,, +deepspeed/model_implementations/diffusers/unet.py,sha256=W2_9-c-5Tz-nymu7kK2du2rzpV8-eC8JZEpQGRhGMek,3056 +deepspeed/model_implementations/diffusers/vae.py,sha256=xbjGaYsXeJA4nu49gQaGO5GrEazCHwcTUVqPsxoOTn4,6244 +deepspeed/model_implementations/features/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/model_implementations/features/__pycache__/__init__.cpython-310.pyc,, +deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-310.pyc,, +deepspeed/model_implementations/features/cuda_graph.py,sha256=-KgILcht5qw_ayzI0CqSPGGFoqZKHG8E5CTfD8ACgw8,563 +deepspeed/model_implementations/transformers/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc,, +deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc,, +deepspeed/model_implementations/transformers/clip_encoder.py,sha256=BTc_do6u4A79uhoNwQYxKrbqSnk2C63BO6xN0Ez-fz0,3096 +deepspeed/model_implementations/transformers/ds_base.py,sha256=feBwoCtVctmdYb_KZruLVCJr7WUBsfNZtVUxEKtO5MQ,388 +deepspeed/model_implementations/transformers/ds_bert.py,sha256=fckfXCn3zf9DsvqiyVAb0pwn25GA9H4ZceaCI4gLLSI,667 +deepspeed/model_implementations/transformers/ds_bloom.py,sha256=_YHlkBOUBTT27ZRWs7x6DxKz_5gILGBAe2YNmlkCij0,669 +deepspeed/model_implementations/transformers/ds_gpt.py,sha256=PkhQU0iHFkSGRxaEb_j1UDvRWruwrOPF8MLZfAo55ZY,665 +deepspeed/model_implementations/transformers/ds_llama2.py,sha256=geGAyCXgRNzMlk0xhvqu3FFmHfFgotVXT6aVV4ZlsYo,2128 +deepspeed/model_implementations/transformers/ds_megatron_gpt.py,sha256=kVtyp2Uss9Z_MRGwgAvJI_2tfWlPgzMpywIPS0xeOIo,682 +deepspeed/model_implementations/transformers/ds_opt.py,sha256=oHS210T27R5zOtg8ik_QwV22F8BIJxKVWPiPQaDXzjY,665 +deepspeed/model_implementations/transformers/ds_transformer.py,sha256=y981LDa0stEBTe8BW_FnQ_1WxHoJSCHNt7FrLiodBzo,8659 +deepspeed/module_inject/__init__.py,sha256=AechzL-WB6riK3mtss4gUD4HD79MhxZzwmPpZzBJSSk,461 +deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/inject.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/layers.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/policy.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc,, +deepspeed/module_inject/__pycache__/utils.cpython-310.pyc,, +deepspeed/module_inject/auto_tp.py,sha256=KAyjZgGCmaQ5_HRgUQvjBgUikmecpnXZw9q62eYQFyw,24809 +deepspeed/module_inject/auto_tp_model_utils.py,sha256=G_tPVAaT7b422Wa8J785rEOcM742uXnoW-5fvd1cinU,5910 +deepspeed/module_inject/containers/__init__.py,sha256=AQbuSXpxEyLpLN3RoAkQeyvtioX_8ShH6C8cqzcQ2Jw,1015 +deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc,, +deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc,, +deepspeed/module_inject/containers/base.py,sha256=kw-culwCQ480zVCgeY7K9LH30BSuKWuPtm2lBMh-L_o,13762 +deepspeed/module_inject/containers/base_moe.py,sha256=lfyM4AhAXmYp5dxa60XoEbcJDQVj9uB-pRlzIDR2ypI,5756 +deepspeed/module_inject/containers/bert.py,sha256=f0VMB8R5mCHxNgIVhQSYeTgVsxo1IN4BV6vzSIOgLZc,3769 +deepspeed/module_inject/containers/bloom.py,sha256=fV_vVFyjB0BB95kP7TrjwqrqruKbEwWQH6ttduzfcwg,6256 +deepspeed/module_inject/containers/clip.py,sha256=KR3s7ev3v-ccvrFhqfm2pKjApQeYeCfhm4aZZ0_Ou_A,2822 +deepspeed/module_inject/containers/distil_bert.py,sha256=j8ufi0mXYe862fLYMtY1J0Rh2uY-xzkXKDKrhK-EsgI,3188 +deepspeed/module_inject/containers/features/__init__.py,sha256=GOivP3ix552VHgGPZ70UehyLQ_YZ8TLRc0-rgTCO07c,275 +deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc,, +deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc,, +deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc,, +deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc,, +deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc,, +deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc,, +deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc,, +deepspeed/module_inject/containers/features/gated_mlp.py,sha256=ep8fnRs-06C7oL75w5UBAaLyCb9BwAx2v3c-sBhkB2k,5397 +deepspeed/module_inject/containers/features/hybrid_engine.py,sha256=ki_k_0SstM39XfN0SNWwiSO3qzIBs2qKxxAyrGxpfUo,8374 +deepspeed/module_inject/containers/features/hybrid_megatron.py,sha256=qp04fT0ManTJwzVdABFP6OWvLQsD6pWDaJAxNJU4-ZQ,4129 +deepspeed/module_inject/containers/features/megatron.py,sha256=BrbmtBC9ZKFL4ESbIYRJLzoEnEp3QOhkj1VcoqlcEek,1200 +deepspeed/module_inject/containers/features/meta_tensor.py,sha256=oem92mPilw6_UOBzkE-7VMoi9vpA9kvA1lemDcCSPwQ,2928 +deepspeed/module_inject/containers/features/split_qkv.py,sha256=totYVBg1M6DcsdC3762Vhoq7gpGJ1BIioTfDFdFwkF0,7136 +deepspeed/module_inject/containers/gpt2.py,sha256=98yUKFAGrfyyCv5yK7V4Zr8oWf6m4qTXgGIzUbRG0kM,2221 +deepspeed/module_inject/containers/gptj.py,sha256=OJ4Z-G5sjd-92yrNDJ0D53Jz_o0n2tfkg9fS9Mll1ss,5074 +deepspeed/module_inject/containers/gptneo.py,sha256=NakhEgR8qh-U9rZ6JD6nIyY_r1mH2obkrgLat4eKpMo,5790 +deepspeed/module_inject/containers/gptneox.py,sha256=zM8hYq0e1FKrDd0n0m1oUPKKMDNpwxVmXQHm8jlv1CA,5897 +deepspeed/module_inject/containers/internlm.py,sha256=3JYddGlL-87jC35MSKvYw9XzZFHYIIYC4Un0Ek5jsL4,7761 +deepspeed/module_inject/containers/llama.py,sha256=II9F76RipDydYbh6R-KYtDekTTQvVcAAHzgizkjT5Z4,6566 +deepspeed/module_inject/containers/llama2.py,sha256=5WbtiVVAL1dVseHiO5kvx2qOwVAAvq-cKBKo0JZ7YjA,6359 +deepspeed/module_inject/containers/megatron_gpt.py,sha256=PiD58-Ul61ZJE4QAkCKCtGwjg44wMaSPOkqnWcRudyo,5417 +deepspeed/module_inject/containers/megatron_gpt_moe.py,sha256=xFOff3OOIKMz2YAFRrqLHxpkzLgB-chA4L69HXiey98,3936 +deepspeed/module_inject/containers/opt.py,sha256=exIuEYIu-KCjwPL3L-ykz1ud477q-F4WBwbXIL0Hef4,6905 +deepspeed/module_inject/containers/unet.py,sha256=fWteUHNx0S8u6C3xpiU2TLK4snlTr7Bv81yCNlVIDY8,1862 +deepspeed/module_inject/containers/vae.py,sha256=yLb5XWwZcaVOXVbUjGW1n0hUr3wrL4rqzqFkvBbD1Y4,1505 +deepspeed/module_inject/fusedqkv_utils.py,sha256=GvycFptz45D6k8eCPqJ-jJs4Vui8QzoyqK0CLVqbICs,10188 +deepspeed/module_inject/inject.py,sha256=KkpC_LLybQeiFzIhr6pz-OlKOl92XNFf3Vh7GqD7jsY,4719 +deepspeed/module_inject/layers.py,sha256=TcMmFqGcapbLQLiZUSlpBFuavpo3PB8V6TCbY7JZ5Bc,31475 +deepspeed/module_inject/load_checkpoint.py,sha256=BwyqWbY18JGZJv709sKG3_ZSUP1lGUjFnP8KYDpESgM,15434 +deepspeed/module_inject/module_quantize.py,sha256=NrOwvSfPWWOnlUVe82sSOZhay76JzqkgZ-j4NP-yEOY,3144 +deepspeed/module_inject/policy.py,sha256=kfyIzx4R4INZve94RPBPZu0jnnQKYpWtou-nRvZQ0Z0,8259 +deepspeed/module_inject/replace_module.py,sha256=_pfrNWVj286Qv4gQYglPWSiMvumKjBj2Tkam7CH4cL4,33068 +deepspeed/module_inject/replace_policy.py,sha256=jlIhG7pbXcjD8JPjWJCU-Gs_WSTG_q7K-IHmewu04pw,1119 +deepspeed/module_inject/tp_shard.py,sha256=n7n9vJYwPMuJVJjN5UaMiDwiPI-xTDOi8oaJMRNKOxg,2126 +deepspeed/module_inject/utils.py,sha256=q5gZWY7YK_-BI7ce16-G8dvU9Qgx7Qbezz9Hwz7oKS0,1995 +deepspeed/moe/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/moe/__pycache__/__init__.cpython-310.pyc,, +deepspeed/moe/__pycache__/experts.cpython-310.pyc,, +deepspeed/moe/__pycache__/layer.cpython-310.pyc,, +deepspeed/moe/__pycache__/mappings.cpython-310.pyc,, +deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc,, +deepspeed/moe/__pycache__/utils.cpython-310.pyc,, +deepspeed/moe/experts.py,sha256=7xFMeuKcBmfyVdaVKpUw5uyGzHYbtxXTiKqcUuBjrQ8,1315 +deepspeed/moe/layer.py,sha256=TayzIRGqEOakrq4TthCiBpqdwwdm5sGqISWMZG13KMs,6905 +deepspeed/moe/mappings.py,sha256=S-o5bwyGoZPmKL0i2E2TrWakyri_P7zARLgOY_4a8u8,4090 +deepspeed/moe/sharded_moe.py,sha256=7z2v3inspO0O8YSGWTll036Xxc3_TiB-vECG2wlKy9E,27332 +deepspeed/moe/utils.py,sha256=YGnqMJD26e2q3ezp-A0srHyb2sdWaHUEg1ld_I6qPgU,6699 +deepspeed/monitor/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/monitor/__pycache__/__init__.cpython-310.pyc,, +deepspeed/monitor/__pycache__/comet.cpython-310.pyc,, +deepspeed/monitor/__pycache__/config.cpython-310.pyc,, +deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc,, +deepspeed/monitor/__pycache__/monitor.cpython-310.pyc,, +deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc,, +deepspeed/monitor/__pycache__/utils.cpython-310.pyc,, +deepspeed/monitor/__pycache__/wandb.cpython-310.pyc,, +deepspeed/monitor/comet.py,sha256=16xIXzwylfuGJyt4p3Lu46JXemCe89_2XCi_cLoWDlY,2903 +deepspeed/monitor/config.py,sha256=vdDaZ7qM0RLXIStjW6S2C4vJimCx0g0IQIi6Hv_sn4k,5258 +deepspeed/monitor/csv_monitor.py,sha256=E9bWWx7QKhex8GWOxMOLjOsrAyBmXqCkaxNUyVirSY0,2907 +deepspeed/monitor/monitor.py,sha256=d_Ke7YvSsXpXeyLmx61KaqkYpgG026sSSbMel0-lfVg,1894 +deepspeed/monitor/tensorboard.py,sha256=9crE_YD9adHS65QmRxAuSuezRQUpL_VHrJMZmPhmuiE,2227 +deepspeed/monitor/utils.py,sha256=MCyfYHyZL-F1a4jbEqG5_d2-1CvX81-gZC6JPlkIblY,1226 +deepspeed/monitor/wandb.py,sha256=G-2BnrWGjiuWG9Tg5-bzkBsjRLwlsyINAYryxOmMdvE,1128 +deepspeed/nebula/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/nebula/__pycache__/__init__.cpython-310.pyc,, +deepspeed/nebula/__pycache__/config.cpython-310.pyc,, +deepspeed/nebula/__pycache__/constants.cpython-310.pyc,, +deepspeed/nebula/config.py,sha256=THO1Mwto2utxnBWRfn0vIVWevOpT2EPFJFZkDCJF-ck,1764 +deepspeed/nebula/constants.py,sha256=0HQOkViV_lRv13vtczIhBJG_YAH876vhh0R4BqHAj68,2786 +deepspeed/nvme/__init__.py,sha256=iPQdRvD-GrL7bbRWcJtX0TFD7kz-OZVTmHDQa1mG3Z0,241 +deepspeed/nvme/__pycache__/__init__.cpython-310.pyc,, +deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc,, +deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc,, +deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc,, +deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc,, +deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc,, +deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc,, +deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc,, +deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc,, +deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc,, +deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc,, +deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc,, +deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc,, +deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc,, +deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc,, +deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc,, +deepspeed/nvme/ds_aio_args.py,sha256=gNjCXHv3VyyioJAWrMoAx4HJYmrN68VNNdq9pwly7SE,8458 +deepspeed/nvme/ds_aio_basic.py,sha256=WXImsGkd3E5qwuMOjZ6wHhCHQM7tzt7nKyg-LzeUpVg,2410 +deepspeed/nvme/ds_aio_constants.py,sha256=Af0dLTpRrw8VXmMby28zy60_tEWx61i15mpUSNzk51E,493 +deepspeed/nvme/ds_aio_handle.py,sha256=HS-sNFfUgwJn6T8o0pZjdK2gMQt4_xLzzHXu-M7PXtU,4944 +deepspeed/nvme/ds_aio_job.py,sha256=v-4U2oTjTqPZ7jjg1yIDtZOMlwfOQcOPMKCmzqTtAQQ,1276 +deepspeed/nvme/io_engine.py,sha256=nUJ1cXfLg6x3IQQ9g-zILY2I584beeigBaTHbsn2lPs,3847 +deepspeed/nvme/parse_nvme_stats.py,sha256=UjaSbCHcP0MifIngLT7mmWbf1BbkKRb-VYO8WbfBmqc,3835 +deepspeed/nvme/perf_generate_param.py,sha256=VPIleO_Ihe5JJRARe94DrXf_sWNCgC8RBPIwz7H_tYI,3289 +deepspeed/nvme/perf_run_sweep.py,sha256=PaLX4sfWHjbIgIJiG3W3t_Mgll_NfHCm6Qmp627Kwb4,9747 +deepspeed/nvme/perf_sweep_utils.py,sha256=pvreOkVAjNHwxSQ33bgXr4Yelwjzy-l9SEOcFTKGjhw,425 +deepspeed/nvme/test_ds_aio.py,sha256=0v_ZTRfro2gjANPa9CVoii1gz021Z5Eol1z9_hbLlR8,590 +deepspeed/nvme/test_ds_aio_utils.py,sha256=BORiau54UKygrVdLsqQLIq1e90ED3PsW1X6WI9DO2s8,3103 +deepspeed/nvme/torch_fastio_engine.py,sha256=zf4S1zsoE64jHUx9hrwYY-Ciy_6GSUC1O0x9rsL5DnQ,3509 +deepspeed/nvme/torch_io.py,sha256=gIq32rm0kRVYP_R-RFZphkmVVDv3TnJVYD8JjKx2gWg,2270 +deepspeed/nvme/validate_async_io.py,sha256=RsuUU_w27HuHWZAKsan5M0-pcllXRyf_H_VQcs-U4T8,302 +deepspeed/ops/__init__.py,sha256=jr72zRfZAdER7PTHnZc-9GziYpnxW80sBCyBthYWYhI,407 +deepspeed/ops/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/adagrad/__init__.py,sha256=aJHYZSC21yulH09u72ZYAe43nLgeMemqvmC9wNNh9CQ,141 +deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc,, +deepspeed/ops/adagrad/cpu_adagrad.py,sha256=nXZLRyBmJS4hcdt6XULXVrX7hQq09nrLFFajJTk5do8,4181 +deepspeed/ops/adam/__init__.py,sha256=ZVagnjkzHVw4akcDibUCfsArTcu8J9wsz9xKmbrJnIs,169 +deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc,, +deepspeed/ops/adam/cpu_adam.py,sha256=URu3D8FceMISyD_-49MgPYWRHJ2SS59deX7sGqnQAKQ,7520 +deepspeed/ops/adam/fused_adam.py,sha256=5_JBk_UvBFcnOL0_hPr5SVFATmmmKHFIlVSR1am-sBc,8767 +deepspeed/ops/adam/multi_tensor_apply.py,sha256=APt3UCnfw-nLjJOGWXC3izWRXNu-16TYhB-s213DpM0,429 +deepspeed/ops/aio/__init__.py,sha256=dHc8QXzdN4Cw0D7px-gbO_fzV-lKT21ySZJ67LtUJuI,136 +deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/compile/__init__.py,sha256=6GIOlNEFXSRcqWtTL1JcczbQQcLh2Ab1oBVcnlnii0I,140 +deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp,sha256=FPSi-AtroLc-zPnzN1yF2OnKkb-wCpKALqN_e8wKTPI,7455 +deepspeed/ops/csrc/adam/cpu_adam.cpp,sha256=9bS2_kdvipPCUoonzJNR6CLuCjNviTPcZJlzaNqGxGY,409 +deepspeed/ops/csrc/adam/cpu_adam_impl.cpp,sha256=QycbyNXh4VGNq4FihHwBysu-_oHSWf1d-STXBgtKZFE,8468 +deepspeed/ops/csrc/adam/fused_adam_frontend.cpp,sha256=tu6pBcAvqcdh1ti6SB7fj8HnqFDQg3Rqbaa4xsQ2v7M,880 +deepspeed/ops/csrc/adam/multi_tensor_adam.cu,sha256=13ixr7zAf-vH8NzAPKYiKXPmVRgsJxILT6cj75_92j8,8417 +deepspeed/ops/csrc/adam/multi_tensor_apply.cuh,sha256=tcBmATZa_Et2VH9hnq5Jf6LJNnLaWDNrSqHdRr2yCE4,5609 +deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp,sha256=hZ8EYNZ-Lcrll0VVnvmb-Gd7mr2CZ_P2a0F9_9bnexo,13274 +deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h,sha256=8fp1mBNfXIAgE48vL5CEtH_n6PeP8TeVs3UXkm-4lXs,1358 +deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp,sha256=vrR8UQ9EK-8EgHBUyQIMgPLxNSFtwNUQVeoClggfPnQ,2033 +deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h,sha256=pnFoIKThq3P_bHoo-jxZgL86SMbXlJZGfpghK7pQqBg,1402 +deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp,sha256=UihEk2VXEqQLCiT0u2ovXoQULBMGbZgSfNBwIx-csSQ,4674 +deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h,sha256=iCd_aS9H_DqJae7-O5qS9GASyE0gNYLLwwrJT-td_GY,2176 +deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp,sha256=G23cmsy1S5u6_vOynCYJnO3Q1B4iZkvwpMEzcprcqQ8,1215 +deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h,sha256=1n67p-zqla1BFNUPQ54Zrpq9OuWaKc7DuAUPKPDjScU,1065 +deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp,sha256=pUuqwGkCkFLDtXtU5mDbgSTW7BXoFlAzFsTyIcx88GI,1394 +deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h,sha256=fAX8as116Ecs7NObKX_LdjALHck0godLITqVkg-Byt4,885 +deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp,sha256=Wye-QMpyW8AxLqqouStbdUDRzikV20ajizyskDygd0g,3433 +deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h,sha256=lXYUW0DsMKdighnXSdKSMfmM4oZxpXCMJ71-Swtmf0c,1038 +deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp,sha256=MpYmRFffPgfkcjI_4tjw-nnEFimy1wgUxFF-HCNl5K4,1850 +deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h,sha256=oIJ0bbxcgE3XFwEMxSOx15XXk1rNbQFUNcYdHqFkQdU,861 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp,sha256=aZysoGV-1JU2wfAKQ3qYpyTGaAFPC8NDJGLQuucwRA0,4232 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h,sha256=z3O_AoPJXMHKjkwsrfrXIcIONXUNc2Yr3Pz1w3Q-99U,965 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp,sha256=YvU1K6pKxcMG9knW0BRkbUh3jr-_N4ySp_QhDjIyudY,903 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h,sha256=WgLYd05nnvNWO492PGBPNoIP8B0X2JI3axKFFSGLntQ,624 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp,sha256=xIqqhPWz-bqBD89HV9TItBnacCWr-JalPjqSaUodwrg,4414 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h,sha256=LxCbVzeTiKRmU2P_cbRZFHIInlW6HQWszpB5s_p2l0M,1169 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp,sha256=Frut6FvfUuJDMIPioBIJzXuuTR54Xjj6nPdBXNAclOo,12834 +deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h,sha256=1bDbRqqOuAOHlIO91k3MytydEQM4YLArmaCakBu5tig,3975 +deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp,sha256=KGjMfQDWaKmViwB6cmJ7uZCJBK5HANTOkufHehFwz3I,4772 +deepspeed/ops/csrc/aio/py_test/single_process_config.json,sha256=6dSFspQHi0nNbLClaIfQ_3Jzj4Nw7urGCly8ZEWnfIo,433 +deepspeed/ops/csrc/compile/deepcompile.cpp,sha256=lar-g-tVzTCXzkrUXy3bevvZ4H65Se3hIfiQF5MUeHU,5868 +deepspeed/ops/csrc/compile/init.cpp,sha256=YfdOvY9YZGyWlo1JohwrsJe5RZuF6EibZekMyphzsGI,4461 +deepspeed/ops/csrc/compile/util.cpp,sha256=v9lmBKZDKPWRWOyi7Azve5PYGa5QO89uLwBHqlwzCCI,3769 +deepspeed/ops/csrc/compile/z1.cpp,sha256=rUEojV-gw-92ygEdS3hDBK52vE1cOfEU9nfz9jqXpMA,4391 +deepspeed/ops/csrc/compile/z1.h,sha256=gn4YNfbL21uG7McpLUhbmml8R7BRLJL9FTpUiwUyZoQ,459 +deepspeed/ops/csrc/compile/z2.cpp,sha256=ACKjRQS_TzisIyohi2HXVbRi_oL5fEsWd8boNO3b7fk,3973 +deepspeed/ops/csrc/compile/z2.h,sha256=HsOGpkbegv3bhZ65vKreTMAD8G2ZS1WJ0ZT-Vx5z_GM,247 +deepspeed/ops/csrc/compile/z3.cpp,sha256=pVo-ktWIbEYygvphK7E8CpBP3H6sV-_rUgy2ftrgcTM,20140 +deepspeed/ops/csrc/compile/z3.h,sha256=eRBbP0wfO15pLQtk4abWFNluBOHNGmOXEGiPQaj1WbY,2347 +deepspeed/ops/csrc/cpu/adam/fused_adam.cpp,sha256=m45BQnPCTkrYHiYcdKNa791zqVqtbdwia4TFoXjnOK0,1433 +deepspeed/ops/csrc/cpu/comm/ccl.cpp,sha256=1U2qIchXfW7ypHD1DIG7mjPXR2levhB55Jz00Kv5Hnk,12342 +deepspeed/ops/csrc/cpu/comm/shm.cpp,sha256=bcPXkJacINgDy081QimwLTBnXXxeUexV3zb_oxStsmk,27403 +deepspeed/ops/csrc/cpu/comm/shm.h,sha256=iM952l2aMWnNP7uI44blK4b8sS80Am756zWpjhd7K3o,401 +deepspeed/ops/csrc/cpu/comm/shm_interface.cpp,sha256=C19jMwHp4HhTinkHeFMdV7NtdMrq2JwUsXIcQK8CwnI,6195 +deepspeed/ops/csrc/cpu/lion/fused_lion.cpp,sha256=-_QUZhwjkjsePW-YyZ4ishpUPCGmbgfDp4C2jNgRIA0,1230 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp,sha256=1d5msQb6IiUdDeAYpzsiKjFHEFPHNAvBXDkO-aY2W_g,2095 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu,sha256=10fKkFtrljPwiHjIAhoZJrBEcIRjdBP8JDagnFdWyOU,9971 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu,sha256=lfsGocib3dneilmyTKnEL9fdmALP8v_UhvwApmwgpVM,7509 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h,sha256=ziZZar1z41wsr6ia5HtX8aAjpAuomtq5cCgZMzYylWI,12185 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h,sha256=RvYmxAW4NN7kMc9pLfEmkXRO-nXMKR1BX3fhhRUG3SI,24898 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h,sha256=WWnsDQeR-XHcjKeEZixxhaTkv89tpyc1jaXoSAEiWt8,9606 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h,sha256=Q_QP5LK3AwDtt6GHxLdUg-ttVqtDkxN2gv4bXiDnlPw,6392 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h,sha256=rXA6J4g_TuQ1Ufm2c0Q2dMcuC1FV6X4MpkYpk7TeYp0,6270 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h,sha256=zn3uylR7vb9GNLKG3UjodjXyPO_XpIuK0bRBwILm4zo,6511 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h,sha256=WeVgJF-IIsVjd50TwDsL_oKwr2UMYUWdLEdhOiy_bfE,30281 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h,sha256=7MyvAD-eXwtNkfyjtd459IvbAbhW81Ys1RzFN8TasTM,15770 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h,sha256=l0wa_nFB9Vrz11bR_0lH-D4vaK-sWfuHm31Mg_Em3bc,9420 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h,sha256=G6LeFpLrmBPVVsnfiYCPjh0h4VdJ5su6jyYPLvPymqI,16280 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h,sha256=o_6L3rJlRtGmughlt9g0xLCzvTqBbUObl72lNxkbhNY,84216 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h,sha256=G02ENUZevpYkNI8FwYfxdiaY8BmFsZkIZh3ieGGWpzg,11661 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h,sha256=Vl83Qgukp3lN0EwAkZdftc1gxIfvwvFzjRw7M8wSQ9A,27084 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h,sha256=bl7iRKlZYnjqYwoST7SRHWwqygjsvYCUmLMBfbAxV7Q,4245 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h,sha256=Jg4qjZ0ukBy2ZIZ-2wxZ9u4mrIcdVAcil56e7CdCY8M,71866 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h,sha256=RQ1sW5cWwiPveRxIs8rVuOgaXBBhsgHU5mbiLh6rawc,31326 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h,sha256=aZy_jvL9zakb9oeFXkoihOabhLLhZjRSbLkf_6lUwS8,71626 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h,sha256=PtMp-E5HhdBmwgRiEbE2ZjlVL-b97euA80waCNO7jYQ,2536 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h,sha256=QvDXwcGimdRqSoaGSzN5aS3q37VSXeihlU_A6ZPee9M,10125 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h,sha256=lLlud-SKkjuOa9UcJkOUfSZRV3RQpPRf9J6qMtAODBQ,95042 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h,sha256=CTv6YV9cCkmzU7w49AqRhENGWOccHYz0d4YzmRhvVBM,47657 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h,sha256=7yLAnvgViEcY5qls5eEJyfQDNkToG0TfPKQhm8Qi0-E,5952 +deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h,sha256=Dr_uSpu3dH8GAfChWggDeon9rn7F4G8BFrjm5_Xyt0E,4476 +deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp,sha256=74zRV9Yj3aoyA6rYhQmZUu0tVa4TqC7g3uR8mbytux0,5789 +deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu,sha256=j17WEtr-J7xnSmOd4yt04OqKZawqPK-P7Aqdfk_omIM,25629 +deepspeed/ops/csrc/fp_quantizer/includes/context.h,sha256=n-WCCbWIdbfNBCeS8c1wdY8LNdGIkz8qwn8C9ImX_nM,1427 +deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h,sha256=UTTOIZa18Cz0wOip7Tfn6itcCeUWySA0DJLDRZLOUY8,5979 +deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp,sha256=bmF57wG0eiTrKQ1WVR3eavTPfVPBSnxnb1i9tcUKQns,5815 +deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h,sha256=pD5329Gd2a75fzfThYDYHIb1Fs6TAeT0Br-YPjunHGg,1131 +deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h,sha256=4abARh9qrz9TE9GkOUJOxZGzavRhQCubOFRuj-G3uLs,4790 +deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp,sha256=X337agbrlRfSF-jnSzCLv4z_CejhaUyjyZowl8EoC4Q,4309 +deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h,sha256=ZookPemckgcq6mTCuStAFxq47xmaRLnKGKeV7ca3BBA,1657 +deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp,sha256=oJb2OJ-RoTtlqHc04Yz6DCg_rG9XZsayURECIbXaqwM,5106 +deepspeed/ops/csrc/includes/StopWatch.h,sha256=DocXNohnOPWiZBqfrdI6AKV9IggYRB1Bylzxfke_i5s,1981 +deepspeed/ops/csrc/includes/Timer.h,sha256=4mYxcATDJElSy5GCj_xFDZ2AvHHmMrhJlW0jJbMP7nQ,1180 +deepspeed/ops/csrc/includes/activation_type.h,sha256=WkFfkL5mgLyYAZimLUouThCkyA2aiF9LQlh1q206L9o,264 +deepspeed/ops/csrc/includes/compat.h,sha256=5ivuTerIqiCaYu2HpRPzJcUtqMsAtRB7aOXel69HeZQ,336 +deepspeed/ops/csrc/includes/context.h,sha256=Cs6hRXBiFSjpKKQOAMCX9F0uP_Hh1IeZLloNJxHxk1E,7207 +deepspeed/ops/csrc/includes/conversion_utils.h,sha256=augq8Zf-t46MaxWh0ryFb46_Guxud0k4kIK-h93rI8Y,12379 +deepspeed/ops/csrc/includes/cpu_adagrad.h,sha256=zcXXyFLZyjNSFfiXgaoYS8ytfFtvKyKQk8hP-SrmLuY,4036 +deepspeed/ops/csrc/includes/cpu_adam.h,sha256=Xkulb32K2B35yGaYqKKeBWEOBrhGb2160UqO-D2C0ok,7365 +deepspeed/ops/csrc/includes/cpu_lion.h,sha256=mFmdjrrWqll4mKKsNGeJysPomGFKBD1e5F2zosWSAq8,5516 +deepspeed/ops/csrc/includes/cublas_wrappers.h,sha256=Gyu-fwvPndgvI1PqGjb-CaoAkeGW8YsRRBBpIjKUjLg,3840 +deepspeed/ops/csrc/includes/custom_cuda_layers.h,sha256=lqmn3AfbfSynF2BLu3YWcZZnwQxsO89LiOluP-J2efU,12964 +deepspeed/ops/csrc/includes/deepcompile.h,sha256=qR0lZVVJCLzBxU_am8OxEOQIg8s4QB7yyRMragxyFLE,21211 +deepspeed/ops/csrc/includes/dequantization_utils.h,sha256=uFpPmBIYd3gNkKXImQLidYnGDYYTXkpY9V4o-AkdYpk,7301 +deepspeed/ops/csrc/includes/dropout.h,sha256=l5JuSzyfxRsfufBKIebPwpghnmYR_E6ucCmGbLvZ1qc,2195 +deepspeed/ops/csrc/includes/ds_kernel_utils.h,sha256=OKOMT46XItRnkrzqKHHK9n_ewxHh6AN-Uhv4aKixBm4,1297 +deepspeed/ops/csrc/includes/ds_transformer_cuda.h,sha256=OwVCwN4tS1kMEeF9J6QV9GG8AvtYklaHRJrMPmji-zY,6164 +deepspeed/ops/csrc/includes/feed_forward.h,sha256=8IbZYyUWVFGMhkjve5ZdWGQc8WvYXVhNEAddT2Vvp6s,3560 +deepspeed/ops/csrc/includes/gelu.h,sha256=iCsuAdNFzhLeccvc5D7KnsAcOcvXN_QD4jlaFRPF340,1018 +deepspeed/ops/csrc/includes/gemm_test.h,sha256=X9oIVKaZpFOGU0yCw-kUz4m5XDq3V1leW5v5dMGQYnA,11546 +deepspeed/ops/csrc/includes/general_kernels.h,sha256=0ksHr78VwjPiAfYPAloHM5uMxriRfNPEZyk9OklLhhI,1507 +deepspeed/ops/csrc/includes/memory_access_utils.h,sha256=Lhpfu-OXn3TscCXK6VdWr8SEp-9rFrtnThrQGDU6aQI,34901 +deepspeed/ops/csrc/includes/normalize_layer.h,sha256=j2lJkrxENFILqdUpEZIMj1qGrFty5SNoAeUrifThvSI,7089 +deepspeed/ops/csrc/includes/quantization.h,sha256=vqhmpPNw0Um-TqEtR_KyO3Pu6eZV5lGpXbraNLIfEoQ,5889 +deepspeed/ops/csrc/includes/quantization_utils.h,sha256=ciuykcFNUW9BxCYTE_moCtN5gJYFd6dgibCleBX7NdM,17366 +deepspeed/ops/csrc/includes/quantizer.h,sha256=YIRNwwr1eZYMmXbs6AZT82shpXuWbn94P4OVMcA3tss,346 +deepspeed/ops/csrc/includes/reduction_utils.h,sha256=AeQx6X3WD4vExcIQRO9LuHdCAWa5vLBcf2QmcIDgePM,23380 +deepspeed/ops/csrc/includes/simd.h,sha256=Ti4NkJ6ciwgPqEuAYmIr0zGRNPRq0z8gK8YmsRSZDKE,10190 +deepspeed/ops/csrc/includes/softmax.h,sha256=UGZHziAfbGN6OdTIA24G5CWHhV8hAfQbq16NN2_0Wlo,1642 +deepspeed/ops/csrc/includes/strided_batch_gemm.h,sha256=vp3ZGYjoSj_xZFTvdIj6PMrn_mupI3YlA4rLFjhYE8I,7279 +deepspeed/ops/csrc/includes/type_shim.h,sha256=8VuX9zK5EsmleO7wuAq9D1eHZC7zKWaf0N97eaGg7x8,6388 +deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp,sha256=icBfzI6QpN8LqIZT4-AY91UXlg_zDqkptNZy10erbws,3995 +deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu,sha256=vjhQw1pj6v19tVE2otOCSStk53Ss0R82g9o1lXfJ5vw,15291 +deepspeed/ops/csrc/lion/cpu_lion.cpp,sha256=didiEmqBezyZ5C2yv92MVbPne8WQbDgkGN5JYV-e1ww,409 +deepspeed/ops/csrc/lion/cpu_lion_impl.cpp,sha256=xCsrkojmjl4uk-QSNa5TpONYRbAQi_V68yacdwvL5vs,7402 +deepspeed/ops/csrc/lion/fused_lion_frontend.cpp,sha256=RMiEOozd9gjXPidc4Ykg-kDIJek0OOASFVwmLIfFs0M,732 +deepspeed/ops/csrc/lion/multi_tensor_apply.cuh,sha256=tMXW3UwsOhNM4J1axQbNkn4gNfmig7_vXPiBLdFCCv4,5595 +deepspeed/ops/csrc/lion/multi_tensor_lion.cu,sha256=11RptKP9XmwMuvGBrmq3P1T8PAH9IvBJirB8gInMVoQ,4377 +deepspeed/ops/csrc/quantization/dequantize.cu,sha256=iibytiJox1QcqE8b57KCLPOaD610STv2VUmKoC7RVM8,3302 +deepspeed/ops/csrc/quantization/fake_quantizer.cu,sha256=CPMZnjOjJTKFibCBeYG1WrXnwEG-0l6POuZCRsBrHBo,37581 +deepspeed/ops/csrc/quantization/pt_binding.cpp,sha256=rcfjK9PldJDktjppS0_Rc9vFh_Rd_YSauGoyj6KzzAY,17434 +deepspeed/ops/csrc/quantization/quant_reduce.cu,sha256=mQ4rRfiPRDfgGJHgt5uFlgsBMav94aZqDYeui9JLBLs,25376 +deepspeed/ops/csrc/quantization/quantize.cu,sha256=IzmI5sq3V8B3tF92i84wlouV186djjVvdISREcepYRc,6988 +deepspeed/ops/csrc/quantization/quantize_intX.cu,sha256=SIEWB2qHtFkHINoTTNukiJnXnQS06BKVH3Wjtg-im1k,8514 +deepspeed/ops/csrc/quantization/swizzled_quantize.cu,sha256=dbNOMXtGQwBI7GnGZXKfVvR9h0iRC-aDU6jSeKrzoMo,19787 +deepspeed/ops/csrc/random_ltd/gather_scatter.cu,sha256=YFiLiYlgkq1RNYS8yRNI3CKYUrv_8ht0lsvZsUcsUYg,8408 +deepspeed/ops/csrc/random_ltd/pt_binding.cpp,sha256=7NAanMKS-j_hBputc6Ltv1yH6-c7xLwhDZJUfLiH--M,9590 +deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu,sha256=KqANgsPTAAHyys0cPh3023JaBF8b0kOKF-W13W39Q_4,5076 +deepspeed/ops/csrc/random_ltd/token_sort.cu,sha256=Tm3oqRHM6UF1PWMuVaghnV0NftUjv3SUEbtBejNLJNc,7044 +deepspeed/ops/csrc/sparse_attention/utils.cpp,sha256=HmNSJfE5WzsE3dYSWjOKOAxawUOcgZ2JwQpJSjvB3b4,4523 +deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu,sha256=8N-FAKh-6Fwu7byz6e_sjXpd8XCRwdFntbZzZL6H8x8,6273 +deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp,sha256=K2zwA7EXdtjM05P9Bl8YV_p35TQXrt4xrADEkq44Dlo,3863 +deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h,sha256=dPh5gJbsCh2gB4_g9PSzMb1zN6LJvAGLfD7Vw_08_mc,915 +deepspeed/ops/csrc/transformer/cublas_wrappers.cu,sha256=HntTKLKThdx3Gj_jBK6gfLD1pnbVE_-wt5Yd_JTQatk,21052 +deepspeed/ops/csrc/transformer/dropout_kernels.cu,sha256=u6U7RuMioRS_aixeRL0aGvU3r7exZ-EFeE7Etgg66IM,29835 +deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp,sha256=sztj1XsEX1rbeq2lovNbYf_jZ-ExLL9qB-TLVRP_1D4,47582 +deepspeed/ops/csrc/transformer/gelu_kernels.cu,sha256=f1emIwHWorXkuIqaB9kD93L4HjzSCcQ2JL7GlE6hh6s,12191 +deepspeed/ops/csrc/transformer/general_kernels.cu,sha256=TzoOK2TqstG0MX9swk0UptK2zBgkpYwo99PeY9j0Fns,14520 +deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu,sha256=7z2HGEwCY1N-JutgeAYGsJJtXyvi36D8NdlpNaqRjcE,8594 +deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu,sha256=z7GzRO9XUJFb0DNPr7luWML9cJTYOZrR5wAQGeuStqY,4916 +deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu,sha256=peyOpVxxboodFWAY0YB6vr3kakCcr8x0EJ4j7HiEY4U,28877 +deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu,sha256=79gLSVFW19gsgFQNezyITbBImha7_oXJo5NfBloOhzU,20883 +deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu,sha256=tJP7V172fEHn86xtO1H8lKzBS7CToDNofJBnOVhivEA,2476 +deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp,sha256=aIdgiEyTGJm_InXaO_eXQ4zH_2MMiRAx9u7vXEEnGkU,89443 +deepspeed/ops/csrc/transformer/inference/csrc/relu.cu,sha256=F_W6D8eeHtm_btn-FTyAV8FnLTDJ9H8mpkDBWpmxxDc,2318 +deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu,sha256=0vHKVYfQQ8w-qLUD-ihfCd0jWDQ_N03QuQu0rimHrdU,10251 +deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu,sha256=GY3hRUkY59BcQqe2BA0QCGw3tlN2UO4GQipRezSzJLM,27245 +deepspeed/ops/csrc/transformer/inference/csrc/transform.cu,sha256=6S1NPVL8Fw9y6wJ0rho2NVIXm64YxthNMKxSWTDsMPo,31887 +deepspeed/ops/csrc/transformer/inference/includes/inference_context.h,sha256=2-KQOUFuWdx1rOLR3rZtTmOpz9LUmY9Vs7a5qNybUzI,10531 +deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h,sha256=AXxx5E4YDDpbKkAxozTHsNJQpIN7dSWr8CiNnpJnBLE,21657 +deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h,sha256=TlXSYeyoKSWBaxuN_dDxrTMhutjQocPjKeGL_NPXIPo,9057 +deepspeed/ops/csrc/transformer/normalize_kernels.cu,sha256=u9fnXA77eqpoC1hri_w0IzC0xW7_OIG5Cs6F_UFnivs,74900 +deepspeed/ops/csrc/transformer/softmax_kernels.cu,sha256=GKmqpxpbAFlhM2p1SFMZ02phEpYMSGUFA2J9ukLovno,26758 +deepspeed/ops/csrc/transformer/transform_kernels.cu,sha256=Ob-Hae1q-KZmDxaji0Ci0nHyNuRYDvPFBF_yrnB8ZoE,22709 +deepspeed/ops/csrc/utils/py_ds_utils.cpp,sha256=jB0kJ0caZitZaw81e6LJ989yaCtCHP_tEFFXvXEfMLk,682 +deepspeed/ops/csrc/utils/tensor_cast.cpp,sha256=GZMlX8HcDDPI0ahjL4_Pf-xjurK1VlvI1dp_LUyhEXQ,797 +deepspeed/ops/csrc/utils/tensor_cast.h,sha256=OBCrxtt5tY0wmP6gd7SS8qlPVcOuYEba-PvC5MTsBSA,383 +deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp,sha256=-LiTifCTt4Pr_4miMyP5-8W84FHhORWZn8mx-VZDR58,6814 +deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp,sha256=tu6pBcAvqcdh1ti6SB7fj8HnqFDQg3Rqbaa4xsQ2v7M,880 +deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp,sha256=Du_BWrHEoYqoCfDb6qKxDeKXP69ZfdFEW1mt7ir5ah0,6595 +deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp,sha256=__3PLu5FDG8vFF9fOwP45KL_yWTXYxKiNdLUYNyqnNI,9482 +deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp,sha256=fSy_UoyFgDlpKWEDRtUqwrhaLMMUV1OG7eCVZfzcGKs,3525 +deepspeed/ops/csrc/xpu/includes/compat.h,sha256=5ivuTerIqiCaYu2HpRPzJcUtqMsAtRB7aOXel69HeZQ,336 +deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h,sha256=7yZNV7T38Yvnv2JLfcdm4KdBuXyPdvT-7Plzva_5BbQ,3951 +deepspeed/ops/csrc/xpu/includes/cpu_adam.h,sha256=3a0twZXzUUnAhrEymYF8gKFoVW594B-I-Snjg6paK9Q,7956 +deepspeed/ops/csrc/xpu/includes/simd.h,sha256=tRwYTW1Xm20rDzsew6c7dfEPbCBd0tsjY26mx1VnVYk,6588 +deepspeed/ops/csrc/xpu/includes/type_shim.h,sha256=qA0DdBShoh_Wute7v52S18jvHB81xi2ADhYvn8EXRVg,7834 +deepspeed/ops/csrc/xpu/packbits/packing.cpp,sha256=oUzgPsI6tXEK56YWrSxAvf5ei289d36We-vmqTXKPuU,3292 +deepspeed/ops/deepspeed4science/__init__.py,sha256=LRmYuJYPcnqiyIuSCwsvqpD61wcScSmS7An8zq0UY4M,175 +deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc,, +deepspeed/ops/deepspeed4science/evoformer_attn.py,sha256=50z8I1w4lJww3bRiBxq8PfEJxcacZw_qXk4DyTyWsSQ,4121 +deepspeed/ops/fp_quantizer/__init__.py,sha256=UxAI8a6mhBD9nD5dIfV6dAcO4qGlz8tMLBSEgVucflA,174 +deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc,, +deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc,, +deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc,, +deepspeed/ops/fp_quantizer/fp8_gemm.py,sha256=ZFcVM33qjxFg5utlPDtZ9GYMThcdFMgeqMAblhs8Fcs,939 +deepspeed/ops/fp_quantizer/fp8_gemm_triton.py,sha256=XDdslcyiSZUoOhPKKm7Hdzy60gqqkOdODqehj-iCYJY,7106 +deepspeed/ops/fp_quantizer/quantize.py,sha256=Cw19CwiaIk8kC_1dyjZSbMcYqVHEbdXNYybkVXN1xog,6481 +deepspeed/ops/gds/__init__.py,sha256=j3yUaB9pVgZKU8_d5i6-ypeWGZbpzKovc-r3ZpF5XLQ,132 +deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/lamb/__init__.py,sha256=HG9WeYIi4tTtKZBuLuJNIVpLSK75djGxjXDkrgfreEk,130 +deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc,, +deepspeed/ops/lamb/fused_lamb.py,sha256=gZbjO38b517Ii1tEJwT0O81YwJ5bXtflM2uq57NFXtI,7815 +deepspeed/ops/lion/__init__.py,sha256=EFF2vXYA9hYPNuqBsLsd1MTFCPiOFfW1ZWBrlbxfbD8,169 +deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc,, +deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc,, +deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc,, +deepspeed/ops/lion/cpu_lion.py,sha256=mJkqnY3EG7YEDi71Fg2qD0HFiSFFFkkz84Vcx8xfG2Y,5306 +deepspeed/ops/lion/fused_lion.py,sha256=QoYzriRhYkaa_bDxIY_YGviCUe-BcNZvvNtFA4g5kLc,5548 +deepspeed/ops/lion/multi_tensor_apply.py,sha256=APt3UCnfw-nLjJOGWXC3izWRXNu-16TYhB-s213DpM0,429 +deepspeed/ops/op_builder/__init__.py,sha256=HJ00U1SrhiXDftTIq9k5hZp-CuRPZ5GWOyQerqgC21U,2005 +deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc,, +deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc,, +deepspeed/ops/op_builder/all_ops.py,sha256=bymFEiEsBwsM_BlMOePCM-9xMGeQdk6u0ZXh2fPQo9w,1223 +deepspeed/ops/op_builder/async_io.py,sha256=EMeLJZ9GBIAjuD2AwDDsfW0GpKPCg-YrK12yqze_Y9w,4500 +deepspeed/ops/op_builder/builder.py,sha256=r3sljupH_zvXGHApweR1FuCGDEQ7hdnuePttpC6UL9E,35323 +deepspeed/ops/op_builder/cpu/__init__.py,sha256=fRbecWZV5CWo2W5bRzfweUx0C_8gjVwmj4Yb6YgIbzY,353 +deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc,, +deepspeed/ops/op_builder/cpu/async_io.py,sha256=aY_ndccYdKQYQ7adRJfl3NfwY8Dj6yqn1OkBr42qn9g,3550 +deepspeed/ops/op_builder/cpu/builder.py,sha256=Bi1gqcUHoDb12ci-5Urb2QRopaQFAy5E5uNELEd4Hzg,1460 +deepspeed/ops/op_builder/cpu/comm.py,sha256=e_bVQ0AnLnZSpsuY4du3NgcpYjhpRoQOVqJu1NQmqW4,2114 +deepspeed/ops/op_builder/cpu/cpu_adam.py,sha256=rGtsjwhzzP9iMDIphGxjh5TLUBtUiEwtrnf5x5BbG08,625 +deepspeed/ops/op_builder/cpu/fused_adam.py,sha256=5LAQSEL2H4Bcc1b7LArqKCcF2Ztuq2uaABmRmI7m5Ss,546 +deepspeed/ops/op_builder/cpu/no_impl.py,sha256=OIZ3aVWXK9W8iRtgNL0pAtl033WV0hEoOuPP35UuFUc,616 +deepspeed/ops/op_builder/cpu_adagrad.py,sha256=fB81OcyxFAMVzINEi1uY1CneERg5KKS8IafOM_ICzCo,622 +deepspeed/ops/op_builder/cpu_adam.py,sha256=aBvkhLFpMxELFvFEpB57C2Fmlq_-3O3azRNEI8XglGo,635 +deepspeed/ops/op_builder/cpu_lion.py,sha256=8l8o0abx4TCn30ps_Iy40HCnymQ9ve225kp_uLRwpWk,635 +deepspeed/ops/op_builder/dc.py,sha256=VL5DLUUq2NbCrJh3Y2WrIcqZmgA1WKj9AVvGfVFDY1U,1157 +deepspeed/ops/op_builder/evoformer_attn.py,sha256=dmGaZQv-G74ckTri8yAMRwjdwomqIjjDeiu2zyCPNgU,3602 +deepspeed/ops/op_builder/fp_quantizer.py,sha256=WUFftK3MJApYc2IZDF7Cm0fBD2PmYCYprbdA23dKSIU,4130 +deepspeed/ops/op_builder/fused_adam.py,sha256=IC87eCjvRPJ2ETVdQUv0-d_yxMXhDJvbPML1AEU9cSQ,1044 +deepspeed/ops/op_builder/fused_lamb.py,sha256=Ojx9euSMUeeU7DVIFxE11w_fHERqYsZxNNtaGxnAAag,1216 +deepspeed/ops/op_builder/fused_lion.py,sha256=1ZJh92_chzH0Jupt8RO8GUOnrm6ACKLECZ5mTiV7WI0,1044 +deepspeed/ops/op_builder/gds.py,sha256=zJPklZuTb6tbi8rnswfmBAC862pwVKv-E6gcv-0QT0M,2195 +deepspeed/ops/op_builder/hpu/__init__.py,sha256=E7-aI0c0S1kZjNCKsjp7XImSfdziWHt3psvwHc82YHQ,370 +deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc,, +deepspeed/ops/op_builder/hpu/builder.py,sha256=nVgIhaRZRgppzIeuYrqfvWcwpGShcK7JUDxHgosA5Cc,1391 +deepspeed/ops/op_builder/hpu/cpu_adam.py,sha256=FzHM2A3XH3dXSYy1DWA7H61swa7m4S9OLsLrbvL8Myk,681 +deepspeed/ops/op_builder/hpu/fp_quantizer.py,sha256=0278EpesoRE54UveUX1ZdlcVlqdmHkzhPk2bFg1bH-U,2947 +deepspeed/ops/op_builder/hpu/fused_adam.py,sha256=ZpP8fkXQ7KdEDtrW5XT9qQ5QlsiGQhONo5r2Mpn03SE,3498 +deepspeed/ops/op_builder/hpu/no_impl.py,sha256=cUJ87-4-I46CK3GmS-Ve52nzfMhcFayT_YZF5FWEmzY,616 +deepspeed/ops/op_builder/hpu/transformer_inference.py,sha256=dPHPLX_wdRFIbCbx4FbxeREn1AJoWOzRbO36dKmcYig,1341 +deepspeed/ops/op_builder/inference_core_ops.py,sha256=ZwfV2HtgoSUv055qWQ8xGSyjI3EUW7TiPy9Xd5DbgQ8,3835 +deepspeed/ops/op_builder/inference_cutlass_builder.py,sha256=OMvFKGBukkUTiYcNHu5CcjZ_oWNA7ifTVGlDu_wBWd0,3373 +deepspeed/ops/op_builder/mlu/__init__.py,sha256=oPAhp7_PS9EV3UCyKDFRfFifj7mJPV2YEMHlqTbVbD8,402 +deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc,, +deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc,, +deepspeed/ops/op_builder/mlu/builder.py,sha256=QdxnbQZxPI98OtPEvDl186lgOat6nQIoIJsIlxYC1aw,1283 +deepspeed/ops/op_builder/mlu/cpu_adagrad.py,sha256=FCrdOCfiYHdfs1HSwhYgohE2x27jmOcFpE9TtqZRa48,565 +deepspeed/ops/op_builder/mlu/cpu_adam.py,sha256=i-Oj1EP5jeuU_7JHu7iqrs16BONW93S7RIq_ImUfyZQ,669 +deepspeed/ops/op_builder/mlu/fused_adam.py,sha256=o574SNUxqSYNXVzGAsRkXdoXpINpwQ6-rysW6SRDVMo,1137 +deepspeed/ops/op_builder/mlu/no_impl.py,sha256=j0r-HhbzJuCVzna9vLWVaDp2wE21AuHDHBN8yy7vUQU,799 +deepspeed/ops/op_builder/npu/__init__.py,sha256=utWarSvvEuL_LUk8-i4Kx2oCHJRemddoGeGfzHOCsTE,419 +deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc,, +deepspeed/ops/op_builder/npu/async_io.py,sha256=D6N-B3SF9bpg8y8hyHxljAdSWIcpvjhXQ3J80DJnACQ,3770 +deepspeed/ops/op_builder/npu/builder.py,sha256=k21-MqhnSCYyMpuOxbRjDWbsIcnDnA6GxWPtsmtBuC8,3200 +deepspeed/ops/op_builder/npu/cpu_adagrad.py,sha256=CM4wGBKhOHayOyzvOGX2yDIh4PgOCu6AUHXF4GKq9kE,581 +deepspeed/ops/op_builder/npu/cpu_adam.py,sha256=S5V76jfADBzNNP-CyBoTXZSWrQMLYWuos6ACygoumDE,594 +deepspeed/ops/op_builder/npu/cpu_lion.py,sha256=rDkiA8vVCMolVFWUnvN9PkrLOJlc6P_rSRNYmOUrYV8,594 +deepspeed/ops/op_builder/npu/fused_adam.py,sha256=uKjA8NKcDUqjr34ToapiXrQHD661uClj6N9equFMh_g,2142 +deepspeed/ops/op_builder/npu/inference.py,sha256=-IsUpVW106uskkegCrYzyMoed5W-RhhffawBWsXYad8,15187 +deepspeed/ops/op_builder/npu/no_impl.py,sha256=qZnbySBzMLLUrMlZV656Pdh8TGfEZKNACKdmD6uT9ho,755 +deepspeed/ops/op_builder/quantizer.py,sha256=k0kTBN9_9Ri_dyzhW_ENxYpUVhiFOeHaIdVY1W50ZPE,1037 +deepspeed/ops/op_builder/ragged_ops.py,sha256=dpipyaz2IADHNLHYC9vavFlXLEMUBJhVwWgwxagKjbI,4912 +deepspeed/ops/op_builder/ragged_utils.py,sha256=CVQGORZRAxjymG9HCuhRtnKudI2sMXd3v3MRVQhi1fQ,2727 +deepspeed/ops/op_builder/random_ltd.py,sha256=zNBXnNjckT9dZqtj4pVYHcSUu4m2govven5qh5Qp9VU,879 +deepspeed/ops/op_builder/sdaa/__init__.py,sha256=N_LPrOGu5tIJVcRRXzDm86SBtEUInPAIJLWMJWqMr1s,1811 +deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc,, +deepspeed/ops/op_builder/sdaa/builder.py,sha256=59jhxPvqJLB-5qePUz0qXUryvSwskkFKKwaRPqRD0U4,2788 +deepspeed/ops/op_builder/sdaa/cpu_adam.py,sha256=DgbrhjECQzOkSyyfDlbrm1l1zxDKC9z3hJ0CrUVeXtQ,2175 +deepspeed/ops/op_builder/sdaa/fused_adam.py,sha256=4u_xou5pjkpTtiPMYhvmX5XuLS-RSQNwRAg9OCNvTlI,2647 +deepspeed/ops/op_builder/sdaa/no_impl.py,sha256=spQASQULrUtgzUL4Cr5TNlItdW_9TWCJwIvGIikLOUQ,2306 +deepspeed/ops/op_builder/sparse_attn.py,sha256=nled88Ns_XIbbbqEITlFF_ZgQujuuv3JKyZIEqV9GNc,3221 +deepspeed/ops/op_builder/spatial_inference.py,sha256=0c8zejdt2H05Q879i-kVCuUxnrsJo7ni0F532EGK5Gg,1599 +deepspeed/ops/op_builder/stochastic_transformer.py,sha256=nQPCaLrtuHVtJxMWIQp_8IT1iQ2XS9nQ1BJVwwAeIPI,565 +deepspeed/ops/op_builder/transformer.py,sha256=wyrj-FZTKPyXzwKOGPYDyic3w5hdtXh4dWJS_wfRdCU,1094 +deepspeed/ops/op_builder/transformer_inference.py,sha256=9R9XaLdkPxf89oZgSfZ9yTcJBpN7B4ic27vBDKTBHbY,2868 +deepspeed/ops/op_builder/utils.py,sha256=w2kP0Vg2PMV0RGbmy0erVVCYK4FN6V7A3YWhpSgJODo,455 +deepspeed/ops/op_builder/xpu/__init__.py,sha256=3WbwNA_gET5S9B3BiSglGVJB6325AKYgdsMXWYU1-7Y,421 +deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc,, +deepspeed/ops/op_builder/xpu/async_io.py,sha256=gFV9LPwhIHNGcSk9bisKe_Z3bCXCxsPyIBFaL-sEtc0,3738 +deepspeed/ops/op_builder/xpu/builder.py,sha256=Pul5G1bF8D5fb-RpXISLJhVoZyE5vH3GNeOA8qKW_NM,5816 +deepspeed/ops/op_builder/xpu/cpu_adagrad.py,sha256=bi9U6d3UmHRYw7qkIcRNRYTuasPYoSe38ngF71ZM1RQ,576 +deepspeed/ops/op_builder/xpu/cpu_adam.py,sha256=1mUDRDd_SV_ZFxCzbt39PajtYqEjRdzHAQPzCClx7FY,627 +deepspeed/ops/op_builder/xpu/flash_attn.py,sha256=YUkFzxNQ_c8MsriNLC1pdqjk-8op60DiDIaky8u52Jk,1615 +deepspeed/ops/op_builder/xpu/fused_adam.py,sha256=yEvt078k2gqqqDjPG_D5I5Ntuz_aAqEjYgtCgjpgAjg,701 +deepspeed/ops/op_builder/xpu/inference.py,sha256=_4-VhP2iIhY9qA_tYadGlKWWtDQ95NVPooVGcyiqzIc,1134 +deepspeed/ops/op_builder/xpu/no_impl.py,sha256=ODPPWw4qSX5ITgjXYT3LHWc-JSe92eBzAIQVKpgkZRA,752 +deepspeed/ops/op_builder/xpu/packbits.py,sha256=eDSmRqT3cFETG48aeHvSHtYg5oaenw2X7WQFW4WJcM0,626 +deepspeed/ops/quantizer/__init__.py,sha256=5IdLoKmcCi6MuZNNTGADQs88dgcCUdN50WofVXDcvpI,132 +deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc,, +deepspeed/ops/quantizer/quantizer.py,sha256=z3dh0MRUklHnR74y662k9H31FOMU0buQOGE9uPLcUZA,1193 +deepspeed/ops/random_ltd/__init__.py,sha256=MC02456CIFtrgIYKClskYz4kDnQ9X7zRyyKNp110l58,191 +deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc,, +deepspeed/ops/random_ltd/dropping_utils.py,sha256=tHMRvk6WIVDoPR58HWc0OwUNzXhfn8VYd7Xw13X-r5o,4902 +deepspeed/ops/sparse_attention/__init__.py,sha256=qhwanCYxLO9eaH1cx52E4LNe_An8Nrja_8WyyoYk42I,467 +deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc,, +deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc,, +deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc,, +deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc,, +deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc,, +deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc,, +deepspeed/ops/sparse_attention/bert_sparse_self_attention.py,sha256=kQM4sl9L-jNirgCWwA9U-4EWGo-AKSzQE7d2CPcCwgw,3465 +deepspeed/ops/sparse_attention/matmul.py,sha256=RYTMwGqaYRUTR7DC60Ek7PNWUc5n9sKJNgcTpzmiZfo,32948 +deepspeed/ops/sparse_attention/softmax.py,sha256=B_q4n2TP29iPfMLTnkRqgHxJ8y5cBjsTixPOO7zrrMw,11322 +deepspeed/ops/sparse_attention/sparse_attention_utils.py,sha256=XlNgnI3CdG8Kur-Ac1O5qJVQRX_lgexTzXtfarMAHb4,12300 +deepspeed/ops/sparse_attention/sparse_self_attention.py,sha256=Mgs_4NaoBoHIghomqaqnKtjbAgYs1xWlZOlZAoSaMao,6746 +deepspeed/ops/sparse_attention/sparsity_config.py,sha256=eUCYxt5QPHWXdot-Myu206H9Y5WV8HrqxvO7LPKHHKg,42463 +deepspeed/ops/sparse_attention/trsrc/__init__.py,sha256=2G0yT0H4-bH6beifS-DoGK8lQcfZ2aIqt7IFekVRxP0,1032 +deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/sparse_attention/trsrc/matmul.tr,sha256=nM6iQGzBXyKFzHqmCDjpr8xmykwEls4Te6PWVJg22Dg,6628 +deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr,sha256=C--cFHczPdwJA7wm0s66p9f_nkzDMatqu3-iBQoAlog,1923 +deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr,sha256=t4bvel9w4ilPsn24DQgzIqGNG08Pt4pfCBFua8rj9fc,4047 +deepspeed/ops/transformer/__init__.py,sha256=ADlxWPujgq5jT_XL8yuRtTg9jdQiBK-5ydzgnNwap-o,413 +deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__init__.py,sha256=ENKx17yjSD2dIqD6kPq-zxY-Xae-O84-ZPGGNeFq064,315 +deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc,, +deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc,, +deepspeed/ops/transformer/inference/bias_add.py,sha256=x1gk_iN4pmFNsSt6jwtBXil-3no4yDhTMbi_6yrs7gY,876 +deepspeed/ops/transformer/inference/config.py,sha256=g9IEbiLEc1kZUvJ3pKvYNQEvwljcIthhupUXpSX914c,6120 +deepspeed/ops/transformer/inference/diffusers_2d_transformer.py,sha256=JVSE9B60n6XHeg37zE26W3-jzKOhxiOtJkmkCiM1nDU,236 +deepspeed/ops/transformer/inference/diffusers_attention.py,sha256=BJa7R6fYdgeqB1gBNX38hD5CqLQdcj8R0g5m4sAp8_Q,9600 +deepspeed/ops/transformer/inference/diffusers_transformer_block.py,sha256=1WHyfxnq1-JMtCZavXr0mK0lPM3vm6QNCRAqVlxwzyc,4324 +deepspeed/ops/transformer/inference/ds_attention.py,sha256=Oiy0oW5ByvWArjqzqU3qUG3eWQmhYQRf_JBwQOyJPu4,15207 +deepspeed/ops/transformer/inference/ds_mlp.py,sha256=Lr5q6pk3MrAx8D36wXhTNDlV7MsFMkFFaHb7CiKzfzo,6294 +deepspeed/ops/transformer/inference/moe_inference.py,sha256=p3Ck2gBRtnY-km8EYnFfjhuyc63rN-ZUVdBBVs97s4A,17360 +deepspeed/ops/transformer/inference/op_binding/__init__.py,sha256=JhdbkMuKVwM62BEU4aI4OscG3GGkxFKsvnttFx9vc9k,382 +deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc,, +deepspeed/ops/transformer/inference/op_binding/base.py,sha256=YN1s4f-7BKDmvCyKkVW3DOaXz40Xxk8FrIeDU2BfmJA,536 +deepspeed/ops/transformer/inference/op_binding/bias_add.py,sha256=9zrXRrRgwfd21zT_jqYF0tJUFur3Dup3x4-WwD_4d_E,980 +deepspeed/ops/transformer/inference/op_binding/bias_gelu.py,sha256=dOuLD_SwIw0p0V45qQnosr82bCDBn1cXtMMzGUXzQ5g,1214 +deepspeed/ops/transformer/inference/op_binding/bias_relu.py,sha256=Z87V8pcsHv01FxuXk2qw-qBsEtG6CNJ0CjJ8MSBUQL0,1161 +deepspeed/ops/transformer/inference/op_binding/bias_residual.py,sha256=ZURDKvdWx8fHPr_2pe9V0d66C_8hbcPLAxgTqKeTClc,952 +deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py,sha256=P5OTQgxyxc73DLQvyEQzzYsH7NQde25awdmSMV9HhK4,938 +deepspeed/ops/transformer/inference/op_binding/gated_activation.py,sha256=xLe2XuWm0-myxEkamkjwliNAode8-GiiCMtcpSB7iLg,1549 +deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py,sha256=LZ3wUeuGrBG56n0jMFbv7-7X_TbzUnJQ6FyQpIGDHNc,2191 +deepspeed/ops/transformer/inference/op_binding/layer_norm.py,sha256=Fj2rFo3RV8Bqg1XpeX1hYVCvWiP2QReN_b9TUSypB60,2370 +deepspeed/ops/transformer/inference/op_binding/linear.py,sha256=vIYuAmj9CX0KERm9_IE7pdmf4XG4UPXkanBqqtlnhgc,2719 +deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py,sha256=ZTwwz6ptQ2hRQw1Mp1_0bheTc8X7Wc8m7e5Clw7yKZ8,5745 +deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py,sha256=2w4xeLtx3u7MdD4jZQ7f1DPeWPIvbOaHVMSICVa9bDw,1002 +deepspeed/ops/transformer/inference/op_binding/pad_transform.py,sha256=dlxmmFGmY0xnUUCONADu2lH4mMnHH4KWvm1odvEnTAc,899 +deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py,sha256=L9j-2d55UW8cqUyq4g1bVxVHT7doW4ACcFdhJXjhs64,1009 +deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py,sha256=vDrBmcbq1xnFcoQ7J6AGQS9X_2NHEaveC-6kPW3E1y0,4514 +deepspeed/ops/transformer/inference/op_binding/residual_add.py,sha256=1w8a4TEIRPtZKPg6z40PC-_dj0ssNDpVA-8JlfPgBFo,3084 +deepspeed/ops/transformer/inference/op_binding/rms_norm.py,sha256=S68R8QOSlT3qyzdmubRx0q6YkMC_qggDH1qPl_Hxmho,985 +deepspeed/ops/transformer/inference/op_binding/softmax.py,sha256=LE1zRTF5JvIN3YmAexhtmTiWdkk5zwdl8A8ZTkmfp1M,3105 +deepspeed/ops/transformer/inference/op_binding/softmax_context.py,sha256=A0L_4CG72sLDUK2n0dWdEVATqt5loB3UpqM66jHNVco,7527 +deepspeed/ops/transformer/inference/op_binding/vector_add.py,sha256=n_l4idWGoSjrFze9_9j_V2YnaDwnVJYzOoSCERnoVAI,882 +deepspeed/ops/transformer/inference/op_binding/vector_matmul.py,sha256=LhLI-FzBDCM9ivWOENA_NkvGAna-eqxESQBNVp7zu44,2640 +deepspeed/ops/transformer/inference/op_binding/workspace.py,sha256=-Of25xBArRF5u0zYYv2rK8SVrjWaQSDla0ztbaqv6Sw,9584 +deepspeed/ops/transformer/inference/triton/__init__.py,sha256=FVRc5cfjotwsFtcB7gkAm09lXFbNl01gaoet67x4B_k,350 +deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc,, +deepspeed/ops/transformer/inference/triton/attention.py,sha256=iLSFq5MOkMXhUc9uQvpmoGo5fzYSIQPJttGNAE6BG5E,16576 +deepspeed/ops/transformer/inference/triton/gelu.py,sha256=OAViD-qU6B50KhjQgg-Y3leuEKP6gW2OkX7MX_VGDko,1152 +deepspeed/ops/transformer/inference/triton/layer_norm.py,sha256=q5Xt2ov4z5DonYRBnWevaUYTFwJU5qIrJc6RVDtLaPw,7512 +deepspeed/ops/transformer/inference/triton/matmul_ext.py,sha256=6020-OGlpSZv8ytcKDnctmwggKCQ9XdLcOjjgkK3R-0,16028 +deepspeed/ops/transformer/inference/triton/mlp.py,sha256=c4yRQv3GPzEKcaxMFwkms4q2eF47ewemkYcnuhu0iec,4225 +deepspeed/ops/transformer/inference/triton/ops.py,sha256=5YLnOEbqn-GydhgR26Qw2z5zZRNDDxQQ_O3oNQyng5w,3804 +deepspeed/ops/transformer/inference/triton/residual_add.py,sha256=PYcAlDC8yz00LzJ8eXsJKoMmxjD_nMM0X5PbK8w_J70,3118 +deepspeed/ops/transformer/inference/triton/softmax.py,sha256=2q_qNmffFM4l2UgsAFKrJmAVrID95KK1rcg5dEzBmco,3208 +deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py,sha256=8pa5U6ZJphYxJ3oX2cPdxdYilmuKf8w_lvt1A8-pjbg,12570 +deepspeed/ops/transformer/inference/triton_ops.py,sha256=lR82VswoMZ9b2Axbck_eF5GCOczUmvKTDAT1W8bVH1w,5487 +deepspeed/ops/transformer/transformer.py,sha256=1UMs9JNjlK3mD5PJi3iER9QARUO2_a5h0sZAY0Rmo24,20600 +deepspeed/pipe/__init__.py,sha256=ddCnO6IbTpGzlAIgfEbGg6RPCz70HNN-MwTjCpTsqLI,164 +deepspeed/pipe/__pycache__/__init__.cpython-310.pyc,, +deepspeed/profiling/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/profiling/__pycache__/__init__.cpython-310.pyc,, +deepspeed/profiling/__pycache__/config.cpython-310.pyc,, +deepspeed/profiling/__pycache__/constants.cpython-310.pyc,, +deepspeed/profiling/config.py,sha256=OzJuTtxVvm62T1iA5eXocvP8xQQXHNtzcCdqISLTTKM,1959 +deepspeed/profiling/constants.py,sha256=QaZFzozkotv71aArhiwtQJ13DNDGDU1xC0Jm0KrVH08,1243 +deepspeed/profiling/flops_profiler/__init__.py,sha256=NCZv_Ktz4sFmmfpBwGKVVbbKrI36VQHwPapvLlWUUxE,120 +deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc,, +deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc,, +deepspeed/profiling/flops_profiler/profiler.py,sha256=NtUZADA8U0-9PioVuNKZnOIFPnQ_-dSlxWhdguJwaXw,51818 +deepspeed/runtime/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 +deepspeed/runtime/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc,, +deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc,, +deepspeed/runtime/__pycache__/compiler.cpython-310.pyc,, +deepspeed/runtime/__pycache__/config.cpython-310.pyc,, +deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc,, +deepspeed/runtime/__pycache__/constants.cpython-310.pyc,, +deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc,, +deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc,, +deepspeed/runtime/__pycache__/engine.cpython-310.pyc,, +deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc,, +deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc,, +deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc,, +deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc,, +deepspeed/runtime/__pycache__/quantize.cpython-310.pyc,, +deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc,, +deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc,, +deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc,, +deepspeed/runtime/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc,, +deepspeed/runtime/activation_checkpointing/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc,, +deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc,, +deepspeed/runtime/activation_checkpointing/checkpointing.py,sha256=AfqMBlOkP0DYiNJyqej_9zbzQdaP3aomJ0scwxn6CZs,44313 +deepspeed/runtime/activation_checkpointing/config.py,sha256=YQkhHYT6Qr2GJl_KusLU0pUJ-rbLLKKj0Zo7LkeZkAg,3988 +deepspeed/runtime/base_optimizer.py,sha256=qqzmVfz_Xdah_QogN3Nz6hiT8NUJacyGjX1uh7TTzQ0,3475 +deepspeed/runtime/bf16_optimizer.py,sha256=hvKo9HanSTB9kClFJEiBQDBNFVa5mS89pl0z2M_ZkxE,25572 +deepspeed/runtime/checkpoint_engine/__init__.py,sha256=LaqB9nJrtf9qqyhJ0Eyi_mDvYXOkGVK8-XIuxWWBVuw,420 +deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/checkpoint_engine/checkpoint_engine.py,sha256=82YbbHCTWv898UU075AemsYniypgmHLGUVLq6C-HwPU,1257 +deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py,sha256=1dnj6UgtfkaOqEq3NpwwwOE_3pnv9mmRuFSKTSoC1NY,5808 +deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py,sha256=Cr-9Px8TGIqVtrErYArLFZ84ueqKIGIXDh4KGO4njTg,1761 +deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py,sha256=5FDE7UUwM_xbB6StNjtynSyEfLpljyx0FwmWi4u9sME,5126 +deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py,sha256=UDWWw1hnn9oYWS0lRa_JP9Mp80Yic0iJ2g6Y7t5SZtg,1764 +deepspeed/runtime/checkpoint_engine/utils.py,sha256=ABeLEQGo9amYK0h_jkTytjENf2ybNkyBKphmyVdQyEI,1820 +deepspeed/runtime/comm/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc,, +deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc,, +deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc,, +deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc,, +deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc,, +deepspeed/runtime/comm/coalesced_collectives.py,sha256=2tdtNA68uONSuqZull-4T_p2eao2wChCLsRUIxM0BR8,11545 +deepspeed/runtime/comm/compressed.py,sha256=yflJJlYN2Wb47h6PUOnUJfmRIS8TJdjMi3LwVAbhvGA,5661 +deepspeed/runtime/comm/hccl.py,sha256=O1XqCH-DdBbJMlTxRkR2l5OUpMQYWXLhkgyZlPwbdCk,5176 +deepspeed/runtime/comm/mpi.py,sha256=m3-XVeEL6F5ln5krGaExKcBj1Qlxg7y_XwvvVo8APcU,10076 +deepspeed/runtime/comm/nccl.py,sha256=kAvrk95McM3accR11f2Grg1spe1kwP8UmC000zemsJw,7584 +deepspeed/runtime/compiler.py,sha256=tCN_Upy5tkLFRtT6Qa0u8-EbK0RLGP6Oaej0YyM8qJ4,712 +deepspeed/runtime/compression/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc,, +deepspeed/runtime/compression/cupy.py,sha256=gCQ7N_EE7TRgZmRaGRHuhGAXrtEkvfbAnJTBMGO-pXI,691 +deepspeed/runtime/config.py,sha256=tuJDGKzsXBASgi1tFUImqeGxsZ2ip3nDie_rYgLB-Zo,40536 +deepspeed/runtime/config_utils.py,sha256=Mjlms4PInPnKZct1xW-ceuulZCvhsspm1SwnuqTs_SE,8460 +deepspeed/runtime/constants.py,sha256=eizwrjKZpfT0qFgAcYBZNF4MGPexhMYgM20o7QZgkFc,15571 +deepspeed/runtime/data_pipeline/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/config.py,sha256=9GGQNeHTR-bog6aBpcxCm65jc5QO2Y3KWNMZoPfK4gs,7721 +deepspeed/runtime/data_pipeline/constants.py,sha256=x4rApIpDlX1zPViP01wMiKqGNPSiqzhzNZ4VlQ7zFp4,5702 +deepspeed/runtime/data_pipeline/curriculum_scheduler.py,sha256=g6Heo7I5L8mpyKJXSH8t1SwO2Lxk9jMfOIMbYDksV00,10025 +deepspeed/runtime/data_pipeline/data_routing/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_routing/basic_layer.py,sha256=QvJuYJRbFTZSkKovN5Z7PVQP5o_fNC-2QgYqrWGy8ns,5638 +deepspeed/runtime/data_pipeline/data_routing/helper.py,sha256=mFWBiepSdAAejDPSaIV6sXeKpRlfXoWFGF0CUmiDACY,1282 +deepspeed/runtime/data_pipeline/data_routing/scheduler.py,sha256=8ddlbZJ3RO0btPKmurh3288Lk1CIoMxRWHcrxlVMZkY,4638 +deepspeed/runtime/data_pipeline/data_routing/utils.py,sha256=ZGrHeImPXbVhLwOMQGWowXEK5YpIGXfC8i3RHVF4NAI,955 +deepspeed/runtime/data_pipeline/data_sampling/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc,, +deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py,sha256=eN9-LeQE66qratDmZ4vyZL1jPs2Fc-M4UM0HXG4Yg-Q,48993 +deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py,sha256=PJNw94Zc0L31DRk4DH9vkXXGvXPMqtisnbF79gFY7UY,19894 +deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py,sha256=QQdA40qmTehZqhBDX2oGeJI8m1r8c6Jc5Hx-ygPa-j8,21449 +deepspeed/runtime/data_pipeline/data_sampling/utils.py,sha256=5QtKBV_54fKL9I7ydqF40uBaLkUpcX5LFeYVz2KY_r4,1690 +deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py,sha256=o-fRgO_JiJ32nQh0rrh-Zz_uGfqsfqHzPsHTkl_KpLE,24311 +deepspeed/runtime/dataloader.py,sha256=rNiiaHzYOU3hp8OAiC2g9tFWPwACJTMhJ_EVrxvP5OU,6977 +deepspeed/runtime/domino/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 +deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc,, +deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc,, +deepspeed/runtime/domino/async_linear.py,sha256=4kDO5n0bLu71BuksABfK4bbi5Opp95JfeYJrOqRJuWQ,4576 +deepspeed/runtime/domino/transformer.py,sha256=xU4nTT-xy52S33rM2tuWlAjG6GxRpNEPQ2lf4MVT0Lg,25278 +deepspeed/runtime/eigenvalue.py,sha256=5SEkMHwnGPqRFyU7ymuoO6cOkRcDDcJnjsnp2wXa2JM,5872 +deepspeed/runtime/engine.py,sha256=tu30PwPILQwNeaqWJWzS83WtaH9SS7rG3jG1WfkKIg4,193515 +deepspeed/runtime/fp16/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 +deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc,, +deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc,, +deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc,, +deepspeed/runtime/fp16/fused_optimizer.py,sha256=w1ke8oaojnHEz-w8kVr_dy_3d72x1lnFt_z4UY18he4,22322 +deepspeed/runtime/fp16/loss_scaler.py,sha256=7E3yzMJEFzbKcIR2aeUx9-9GrB6_2sGv5wOX8NZ6ezY,11652 +deepspeed/runtime/fp16/onebit/__init__.py,sha256=8hnDOpk71O_EekeWgWIJ4CnbtYt67dre4flu5hoES-4,186 +deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc,, +deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc,, +deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc,, +deepspeed/runtime/fp16/onebit/adam.py,sha256=jRCUSTzBhvo1xPD5XvRxmHni8R2JXMBGL1_P7udBmpc,15694 +deepspeed/runtime/fp16/onebit/lamb.py,sha256=jMZBvuwG8GxoF4tD22ECq9jZyl8_XBcAKnDLyS0GDHA,23556 +deepspeed/runtime/fp16/onebit/zoadam.py,sha256=z3Pc2u9LyJXm5hCfQjAzj_wp72RJP8In6LYv5KyrChA,19554 +deepspeed/runtime/fp16/unfused_optimizer.py,sha256=qB2H34wcs1q05TqhyrA5wPYmC-RJduqjxswGA1rs3Ng,18029 +deepspeed/runtime/hybrid_engine.py,sha256=rfP2MpmnB-DqCiSRjwKyklcl78RtGNwFOQa-PiutlWU,20965 +deepspeed/runtime/lr_schedules.py,sha256=eDap8FWaxsp6M-VSxi07NWLKWCwEwyqbOLIQb6frQxU,39074 +deepspeed/runtime/model_checkpointing/__init__.py,sha256=9EfmW2osMT_aBTR0xOE5NadoNHh7mETjTJxf5lnz-4k,173 +deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc,, +deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc,, +deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc,, +deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc,, +deepspeed/runtime/model_checkpointing/config.py,sha256=jFcd7SQFwXj7JqWgpJHXinDN5j9X2-73U4aWd4BBFaU,2992 +deepspeed/runtime/model_checkpointing/constants.py,sha256=JJ2jAMNPTnWuFwr6JJelgayVBryCcIQhpXlxXBga0zg,2258 +deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py,sha256=hsW-v_h4VfZM8gbn57EBK-Ov4noINYpnKAku8GD9x6Y,10732 +deepspeed/runtime/model_checkpointing/utils.py,sha256=C78G1D098fb3ffbstnGkhcO8CEw2pUbGq9NLtxMSsl8,3275 +deepspeed/runtime/model_checkpointing/writer_factory.py,sha256=SaFlZX0cRvXrIVKTg_qSmOOd7NK3mYvpQx4NYdpPQz4,4930 +deepspeed/runtime/pipe/__init__.py,sha256=4Xc534VEOCSKYVc-ImMFLx5eTmPB_URVgWqweehFsVI,195 +deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc,, +deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc,, +deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc,, +deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc,, +deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc,, +deepspeed/runtime/pipe/engine.py,sha256=Rg89Sfun9GELR8Ob4OrClFyEmi8fMn6cOrOn6-WQ8eQ,63690 +deepspeed/runtime/pipe/module.py,sha256=FfLY5D8DZVGy0kEtcTs7Pqmc4U3V7Lq6TdPwZyRMYy0,30551 +deepspeed/runtime/pipe/p2p.py,sha256=t4K_Q2W2YI4K5bPOF1iqC5z7Ytyj4K1vHKA3gcSioBU,5344 +deepspeed/runtime/pipe/schedule.py,sha256=ZDzAeTPZTaYt77Wi2KSkycfuwneBGFsa7V6afQAzHo4,15546 +deepspeed/runtime/pipe/topology.py,sha256=Y_bQMiTNioGRb4twzWfsxlSkZc0iepbTbs9yvx7H-Cs,17698 +deepspeed/runtime/precision_config.py,sha256=HGZxwS8lygvliSHrtYLBl3B2eJmW37j7crRYzRt4UMQ,3612 +deepspeed/runtime/progressive_layer_drop.py,sha256=5zb3-BrMbRxYZL5lk1FtvaOknMi31xT1refDkgSMQbQ,1353 +deepspeed/runtime/quantize.py,sha256=OmimrTJV_RoVG8TLVdjcQR91kSatQCNdFkPf3U39fsk,7699 +deepspeed/runtime/sequence_parallel/__init__.py,sha256=2BSwDtJJ1v4uCmvYlXq0W17dbHPGC4v0-Xv0IfmtQb0,99 +deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc,, +deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc,, +deepspeed/runtime/sequence_parallel/parallel_state_sp.py,sha256=SfNF80rWVfeByWao4GcT3wCBOibzTNqZlFi4GO8l2Fs,3677 +deepspeed/runtime/sequence_parallel/ulysses_sp.py,sha256=0n4mc4u70lhFHNlSZ5cJze3T533qnKdXHluPryErMH8,55816 +deepspeed/runtime/sparse_tensor.py,sha256=QBoplFkxSYLqE9Jfj8mN2qNLMMpVpEgpP3nuIsobpRg,2466 +deepspeed/runtime/state_dict_factory.py,sha256=Zp2mJbIySbYgp2cl5pFvLkpROOOqX9eqGnG1b3d0kD0,18177 +deepspeed/runtime/swap_tensor/__init__.py,sha256=OiMgE9mkwbOsv3gC91ENJRAHc7zlVfMXXhn4pxN7zcc,134 +deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/swap_tensor/aio_config.py,sha256=CaII3llLuEecOBceu7MuaqUtyFyzxARYE3Fspyj7DQ0,1643 +deepspeed/runtime/swap_tensor/async_swapper.py,sha256=RlwlqBbkHa3nviVT4nrYzZ4s05nPch_dNMHTZcFOdxw,6349 +deepspeed/runtime/swap_tensor/constants.py,sha256=lmGtuinGdUVZ8lEgvKKj0dVLUuiLRQTfPaEhxxWPg7E,700 +deepspeed/runtime/swap_tensor/optimizer_utils.py,sha256=MTQIkQ9zT0C6gpyS8zeB9et40z1BXFnh1e1z5FJBaZA,21403 +deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py,sha256=0-t1zOLHC_PMm5BnllC0BwjToFpYvBYYHby_ALmIOYw,10629 +deepspeed/runtime/swap_tensor/partitioned_param_swapper.py,sha256=U3Oopxi88nOvXx7wkDAt8eGnk1fuvIoV5ClhDce6QHc,18791 +deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py,sha256=AiFBvtcln-TrI3pSnaQzLTvaVTNBfiObA9_muL35R_M,11346 +deepspeed/runtime/swap_tensor/utils.py,sha256=Yrgbja_CefG1juXQVS-_KfnUu4vmP5Cw-r0p_fJ8DgU,7820 +deepspeed/runtime/tensor_parallel/__init__.py,sha256=BwFdU2Aa2rmV01DRps6I9hRTGsOiBzsTGH7oaEoeUDE,198 +deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc,, +deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc,, +deepspeed/runtime/tensor_parallel/config.py,sha256=rcNLwzsdTG1wInLFF1JI916VjJFYxhuoFuZZNWiZfmE,2860 +deepspeed/runtime/tensor_parallel/tp_manager.py,sha256=nPEscixH1Of-jaTsyrXzGo3qSHK0gaWDpfnVbZzbDhg,2539 +deepspeed/runtime/torch_autocast.py,sha256=ERSLVErVR1athv4vKhGPebAT8oBWzS1oQexgfCPt7wA,3638 +deepspeed/runtime/utils.py,sha256=4WyJnDEc7k87RxQco63P1PIxikNABzB62aYN7E0v34A,46087 +deepspeed/runtime/weight_quantizer.py,sha256=-aS6hTZnMcc3ckoYaCI6MiHybwy2NZo5AqEZ8rT7G_o,7027 +deepspeed/runtime/zero/__init__.py,sha256=3hNEJ-32gAoqG24c00D2lrC4WY0qFQPr0AbNSqZIazU,611 +deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc,, +deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc,, +deepspeed/runtime/zero/config.py,sha256=tc0rFE0llwvM70WIyiLTc4iVWH8WZZhE0PftZp98AF8,14345 +deepspeed/runtime/zero/contiguous_memory_allocator.py,sha256=r1J-KRJqO_Aa4Rqlsb1cgtHPa4mpnxm_zrd19FzZVQU,10923 +deepspeed/runtime/zero/linear.py,sha256=vOhCaYR-E_Xg154-9VUsaF1rYJJUq6jytC_qxeSW9Zg,7998 +deepspeed/runtime/zero/mics.py,sha256=mKc4itpTLX3oQqciDiHHHmgF7cd_IRjOjgMv_5chk7E,22339 +deepspeed/runtime/zero/mics_utils.py,sha256=GWmBVESUd699i4YjCKScOwAwlwwFGs0si1dmPAeMz60,7500 +deepspeed/runtime/zero/offload_config.py,sha256=NN8h1kMNtH61GrSbnLTED2JnjzwGYL8v3IozOrSDIPI,3359 +deepspeed/runtime/zero/offload_states.py,sha256=WThDC2EhyaOArrehBNr-_VUNPc4hY5ZeTIXKsGh9Qhg,2956 +deepspeed/runtime/zero/parameter_offload.py,sha256=ObXU8MT1znWw8gg5vlIREvTzXqiGPIta4QBzVdAwFc8,26426 +deepspeed/runtime/zero/partition_parameters.py,sha256=Ob77bimPLPGXRsODOiIdnqg3jdNibK9MyVhNhd2omOk,110399 +deepspeed/runtime/zero/partitioned_param_coordinator.py,sha256=oVpLaxrhKbwsi3kvjlbJL79NSOhfmwQc4JJY5Ns9PUM,29471 +deepspeed/runtime/zero/partitioned_param_profiler.py,sha256=xTXgTnu1ew215TluuO2HvP78hZNbqPEhzJBieeOChf8,1801 +deepspeed/runtime/zero/stage3.py,sha256=88nREaCuBp62fCz_5uTRBiumtlSdLVvPK0y608QAcPU,155200 +deepspeed/runtime/zero/stage_1_and_2.py,sha256=A1t4UUHgHKBQ3KyTybM1LpvZ2qVdhtscfGWf9Sb73hE,126808 +deepspeed/runtime/zero/test.py,sha256=j7zu8Fnxm5_Wmp-aaUiQicPD1EUL2JGpODeSARFSZVs,2727 +deepspeed/runtime/zero/tiling.py,sha256=Uz6IgvJPAQ4GICXj4E509p_BIDvnczWVtBT7EjZU9rg,11727 +deepspeed/runtime/zero/utils.py,sha256=mpmn-5rjhKzd9BsXWoAf6cB90Ed4PnpynoCRgE0tBAs,5868 +deepspeed/sequence/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 +deepspeed/sequence/__pycache__/__init__.cpython-310.pyc,, +deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc,, +deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc,, +deepspeed/sequence/__pycache__/layer.cpython-310.pyc,, +deepspeed/sequence/cross_entropy.py,sha256=-YW4WYSHzZf_42OVu4WJuAp53GQAEnjDfIvGugAKA5k,2140 +deepspeed/sequence/fpdt_layer.py,sha256=MqPtxPXHxa5Gp6RFBIrQujHyJ_Qu8cJWH0g4gOEIGU0,56246 +deepspeed/sequence/layer.py,sha256=o-qB87cWrjJjf6152Xh0JWiCOzO-Kg6HCrAtEytHeW4,20516 +deepspeed/utils/__init__.py,sha256=n_hu2LhHCYt0EIvkNxEVJrSqyL_AhGWqQAAN6TkYPdg,1407 +deepspeed/utils/__pycache__/__init__.cpython-310.pyc,, +deepspeed/utils/__pycache__/bwc.cpython-310.pyc,, +deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc,, +deepspeed/utils/__pycache__/config.cpython-310.pyc,, +deepspeed/utils/__pycache__/debug.cpython-310.pyc,, +deepspeed/utils/__pycache__/exceptions.cpython-310.pyc,, +deepspeed/utils/__pycache__/groups.cpython-310.pyc,, +deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc,, +deepspeed/utils/__pycache__/logging.cpython-310.pyc,, +deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc,, +deepspeed/utils/__pycache__/numa.cpython-310.pyc,, +deepspeed/utils/__pycache__/nvtx.cpython-310.pyc,, +deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc,, +deepspeed/utils/__pycache__/timer.cpython-310.pyc,, +deepspeed/utils/__pycache__/torch.cpython-310.pyc,, +deepspeed/utils/__pycache__/types.cpython-310.pyc,, +deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc,, +deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc,, +deepspeed/utils/bwc.py,sha256=EbdKjvtp_6XwUc3XhCyPxYfIqe9v5I23ynXGG39aAcY,3970 +deepspeed/utils/comms_logging.py,sha256=Bx8Gz6g1BPgMZ7FjtGLlqRyhKGuPPgyKMAGVdW4K_pg,7996 +deepspeed/utils/config.py,sha256=O4l-6Da_8RWBeyJ4tZpKqJPCLNtYK7eMVaqsVCbQkRU,1344 +deepspeed/utils/debug.py,sha256=ilbE4Z7cz3Ww0DR7TzP6vPQbJjyZFaq-ayO6FagSTAo,4936 +deepspeed/utils/exceptions.py,sha256=h4J_9uk3HmKG8LdiSO8DKYzOmIYrD1MGwtZD1nA-Q3g,144 +deepspeed/utils/groups.py,sha256=CNBrWWq3J9m04nnpZjaViZYNnZEtEp89LViE1LVnDog,30800 +deepspeed/utils/init_on_device.py,sha256=Q4RFeRMi7PYGBHfTKh3E63X8IZquheDWYKbT56aOuuQ,3004 +deepspeed/utils/logging.py,sha256=Khqk1omHmqqcYJgI88YW-JxY0zFnNYZ7QAfLAsixX40,5016 +deepspeed/utils/mixed_precision_linkage.py,sha256=OAzA20w_15nLUurjltJRTkGC__-aNx68KWrf5cbYfRg,2484 +deepspeed/utils/numa.py,sha256=ggGkhc8gXvuzkSqpl7JzmQ4K3H9Y44pnaq9XUAgBq-Q,7174 +deepspeed/utils/nvtx.py,sha256=XeqDmA2fdAUTH1FtFQwEr0K5Y7-gTWe9ZqLSp5qttec,746 +deepspeed/utils/tensor_fragment.py,sha256=6uEsYS1-eq8iz9iViCx8LzA231jIpSPIjSFC5QUdQ_o,17039 +deepspeed/utils/timer.py,sha256=V1iXDzsYWmuE0l1hoMKditAJMr1y8sShO5MZOfyF9r0,10742 +deepspeed/utils/torch.py,sha256=xxVsKHjO8HqeDJeffSed1TiP36IAlMn8lsakLuj-MGo,888 +deepspeed/utils/types.py,sha256=IcEwFod7RqLTRCGf7lr2sbKhFRf1an1KvtQ98nswwNw,434 +deepspeed/utils/z3_leaf_module.py,sha256=9eVQubxo8VxE-hBkrI3qxrddUKFoOBTUL4SCKCQiBe8,4758 +deepspeed/utils/zero_to_fp32.py,sha256=Rkl1ZczytKix9vGMg0EELzdJYFqUM1yB9p3xvSaK9k8,33272 diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/WHEEL b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..e7fa31b6f3f78deb1022c1f7927f07d4d16da822 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..cfa0d2ef62101b531baf7ac75575e64b2c0cb6a2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/entry_points.txt @@ -0,0 +1,2 @@ +[pytest_randomly.random_seeder] +deepspeed = deepspeed.runtime.utils:set_random_seed diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/licenses/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..0f2c819def28c8c4109156014fde317020d4ddb9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/top_level.txt @@ -0,0 +1 @@ +deepspeed diff --git a/venv/lib/python3.10/site-packages/deepspeed/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c7e01a0ac969325ec381d69a6a9174c3811cda52 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/__init__.py @@ -0,0 +1,398 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import types +import json +from typing import Optional, Union +import torch +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler +from packaging import version as pkg_version + +# Skip Triton import for AMD due to pytorch-triton-rocm module breaking device API in DeepSpeed +if not (hasattr(torch.version, 'hip') and torch.version.hip is not None): + try: + import triton # noqa: F401 # type: ignore + HAS_TRITON = True + except ImportError: + HAS_TRITON = False +else: + HAS_TRITON = False + +from . import ops +from . import module_inject + +from .accelerator import get_accelerator +from .constants import TORCH_DISTRIBUTED_DEFAULT_PORT +from .runtime.engine import DeepSpeedEngine, DeepSpeedOptimizerCallable, DeepSpeedSchedulerCallable +from .runtime.engine import ADAM_OPTIMIZER, LAMB_OPTIMIZER +from .runtime.hybrid_engine import DeepSpeedHybridEngine +from .runtime.pipe.engine import PipelineEngine +from .inference.engine import InferenceEngine +from .inference.config import DeepSpeedInferenceConfig +from .runtime.lr_schedules import add_tuning_arguments +from .runtime.config import DeepSpeedConfig, DeepSpeedConfigError +from .runtime.activation_checkpointing import checkpointing +from .ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig +from .module_inject import replace_transformer_layer, revert_transformer_layer, set_autotp_mode + +from .utils import log_dist, OnDevice, logger +from .comm.comm import init_distributed + +from .runtime import zero, domino +from .runtime.compiler import is_compile_supported + +from .pipe import PipelineModule + +from .git_version_info import version, git_hash, git_branch + + +def _parse_version(version_str): + '''Parse a version string and extract the major, minor, and patch versions.''' + ver = pkg_version.parse(version_str) + return ver.major, ver.minor, ver.micro + + +# Export version information +__version__ = version +__version_major__, __version_minor__, __version_patch__ = _parse_version(__version__) +__git_hash__ = git_hash +__git_branch__ = git_branch + +# Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init +dist = None + + +def initialize(args=None, + model: torch.nn.Module = None, + optimizer: Optional[Union[Optimizer, DeepSpeedOptimizerCallable]] = None, + model_parameters: Optional[torch.nn.Module] = None, + training_data: Optional[torch.utils.data.Dataset] = None, + lr_scheduler: Optional[Union[_LRScheduler, DeepSpeedSchedulerCallable]] = None, + distributed_port: int = TORCH_DISTRIBUTED_DEFAULT_PORT, + mpu=None, + dist_init_required: Optional[bool] = None, + collate_fn=None, + config=None, + mesh_param=None, + config_params=None): + """Initialize the DeepSpeed Engine. + + Arguments: + args: an object containing local_rank and deepspeed_config fields. + This is optional if `config` is passed. + + model: Required: nn.module class before apply any wrappers + + optimizer: Optional: a user defined Optimizer or Callable that returns an Optimizer object. + This overrides any optimizer definition in the DeepSpeed json config. + + model_parameters: Optional: An iterable of torch.Tensors or dicts. + Specifies what Tensors should be optimized. + + training_data: Optional: Dataset of type torch.utils.data.Dataset + + lr_scheduler: Optional: Learning Rate Scheduler Object or a Callable that takes an Optimizer and returns a Scheduler object. + The scheduler object should define a get_lr(), step(), state_dict(), and load_state_dict() methods + + distributed_port: Optional: Master node (rank 0)'s free port that needs to be used for communication during distributed training + + mpu: Optional: A model parallelism unit object that implements + get_{model,data}_parallel_{rank,group,world_size}() + + dist_init_required: Optional: None will auto-initialize torch distributed if needed, + otherwise the user can force it to be initialized or not via boolean. + + collate_fn: Optional: Merges a list of samples to form a + mini-batch of Tensor(s). Used when using batched loading from a + map-style dataset. + + config: Optional: Instead of requiring args.deepspeed_config you can pass your deepspeed config + as an argument instead, as a path or a dictionary. + + config_params: Optional: Same as `config`, kept for backwards compatibility. + + Returns: + A tuple of ``engine``, ``optimizer``, ``training_dataloader``, ``lr_scheduler`` + + * ``engine``: DeepSpeed runtime engine which wraps the client model for distributed training. + + * ``optimizer``: Wrapped optimizer if a user defined ``optimizer`` is supplied, or if + optimizer is specified in json config else ``None``. + + * ``training_dataloader``: DeepSpeed dataloader if ``training_data`` was supplied, + otherwise ``None``. + + * ``lr_scheduler``: Wrapped lr scheduler if user ``lr_scheduler`` is passed, or + if ``lr_scheduler`` specified in JSON configuration. Otherwise ``None``. + """ + log_dist("DeepSpeed info: version={}, git-hash={}, git-branch={}".format(__version__, __git_hash__, + __git_branch__), + ranks=[0]) + + # Disable zero.Init context if it's currently enabled + zero.partition_parameters.shutdown_init_context() + + assert model is not None, "deepspeed.initialize requires a model" + + global dist + from deepspeed import comm as dist + dist_backend = get_accelerator().communication_backend_name() + dist.init_distributed(dist_backend=dist_backend, + distributed_port=distributed_port, + dist_init_required=dist_init_required) + + ##TODO: combine reuse mpu as mesh device and vice versa + # Set config using config_params for backwards compat + if config is None and config_params is not None: + config = config_params + + mesh_device = None + if mesh_param: + logger.info(f"mesh_param to Initialize mesh device: {mesh_param}") + mesh_device = dist.initialize_mesh_device(mesh_param, ("data_parallel", "sequence_parallel")) + #if config file has sequence parallelize and data parallelize, then use them to initialize mesh device + elif config is not None: + if "sequence_parallel_size" in config and "data_parallel_size" in config: + logger.info(f"config to Initialize mesh device: {config}") + mesh_device = dist.initialize_mesh_device((config["data_parallel_size"], config["sequence_parallel_size"]), \ + ("data_parallel", "sequence_parallel")) + + # Check for deepscale_config for backwards compat + if hasattr(args, "deepscale_config") and args.deepscale_config is not None: + logger.warning("************ --deepscale_config is deprecated, please use --deepspeed_config ************") + if hasattr(args, "deepspeed_config"): + assert (args.deepspeed_config + is None), "Not sure how to proceed, we were given both a deepscale_config and deepspeed_config" + args.deepspeed_config = args.deepscale_config + args.deepscale_config = None + + # Check that we have only one config passed + if hasattr(args, "deepspeed_config") and args.deepspeed_config is not None: + assert config is None, "Not sure how to proceed, we were given deepspeed configs in the deepspeed arguments and deepspeed.initialize() function call" + config = args.deepspeed_config + assert config is not None, "DeepSpeed requires --deepspeed_config to specify configuration file" + if not isinstance(model, PipelineModule): + config_class = DeepSpeedConfig(config, mpu, mesh_device=mesh_device) + if config_class.hybrid_engine.enabled: + engine = DeepSpeedHybridEngine(args=args, + model=model, + optimizer=optimizer, + model_parameters=model_parameters, + training_data=training_data, + lr_scheduler=lr_scheduler, + mpu=mpu, + dist_init_required=dist_init_required, + collate_fn=collate_fn, + config=config, + config_class=config_class) + else: + engine = DeepSpeedEngine(args=args, + model=model, + optimizer=optimizer, + model_parameters=model_parameters, + training_data=training_data, + lr_scheduler=lr_scheduler, + mpu=mpu, + dist_init_required=dist_init_required, + collate_fn=collate_fn, + config=config, + mesh_device=mesh_device, + config_class=config_class) + else: + assert mpu is None, "mpu must be None with pipeline parallelism" + mpu = model.mpu() + config_class = DeepSpeedConfig(config, mpu) + engine = PipelineEngine(args=args, + model=model, + optimizer=optimizer, + model_parameters=model_parameters, + training_data=training_data, + lr_scheduler=lr_scheduler, + mpu=mpu, + dist_init_required=dist_init_required, + collate_fn=collate_fn, + config=config, + config_class=config_class) + + # Restore zero.Init context if necessary + zero.partition_parameters.restore_init_context() + + return_items = [ + engine, + engine.optimizer, + engine.training_dataloader, + engine.lr_scheduler, + ] + return tuple(return_items) + + +def _add_core_arguments(parser): + r"""Helper (internal) function to update an argument parser with an argument group of the core DeepSpeed arguments. + The core set of DeepSpeed arguments include the following: + 1) --deepspeed: boolean flag to enable DeepSpeed + 2) --deepspeed_config : path of a json configuration file to configure DeepSpeed runtime. + + This is a helper function to the public add_config_arguments() + + Arguments: + parser: argument parser + Return: + parser: Updated Parser + """ + group = parser.add_argument_group('DeepSpeed', 'DeepSpeed configurations') + + group.add_argument('--deepspeed', + default=False, + action='store_true', + help='Enable DeepSpeed (helper flag for user code, no impact on DeepSpeed backend)') + + group.add_argument('--deepspeed_config', default=None, type=str, help='DeepSpeed json configuration file.') + + group.add_argument('--deepscale', + default=False, + action='store_true', + help='Deprecated enable DeepSpeed (helper flag for user code, no impact on DeepSpeed backend)') + + group.add_argument('--deepscale_config', + default=None, + type=str, + help='Deprecated DeepSpeed json configuration file.') + + return parser + + +def add_config_arguments(parser): + r"""Update the argument parser to enabling parsing of DeepSpeed command line arguments. + The set of DeepSpeed arguments include the following: + 1) --deepspeed: boolean flag to enable DeepSpeed + 2) --deepspeed_config : path of a json configuration file to configure DeepSpeed runtime. + + Arguments: + parser: argument parser + Return: + parser: Updated Parser + """ + parser = _add_core_arguments(parser) + + return parser + + +def default_inference_config(): + """ + Return a default DeepSpeed inference configuration dictionary. + """ + return DeepSpeedInferenceConfig().dict() + + +def init_inference(model, config=None, **kwargs): + """Initialize the DeepSpeed InferenceEngine. + + Description: all four cases are valid and supported in DS init_inference() API. + + # Case 1: user provides no config and no kwargs. Default config will be used. + + .. code-block:: python + + generator.model = deepspeed.init_inference(generator.model) + string = generator("DeepSpeed is") + print(string) + + # Case 2: user provides a config and no kwargs. User supplied config will be used. + + .. code-block:: python + + generator.model = deepspeed.init_inference(generator.model, config=config) + string = generator("DeepSpeed is") + print(string) + + # Case 3: user provides no config and uses keyword arguments (kwargs) only. + + .. code-block:: python + + generator.model = deepspeed.init_inference(generator.model, + tensor_parallel={"tp_size": world_size}, + dtype=torch.half, + replace_with_kernel_inject=True) + string = generator("DeepSpeed is") + print(string) + + # Case 4: user provides config and keyword arguments (kwargs). Both config and kwargs are merged and kwargs take precedence. + + .. code-block:: python + + generator.model = deepspeed.init_inference(generator.model, config={"dtype": torch.half}, replace_with_kernel_inject=True) + string = generator("DeepSpeed is") + print(string) + + Arguments: + model: Required: original nn.module object without any wrappers + + config: Optional: instead of arguments, you can pass in a DS inference config dict or path to JSON file + + Returns: + A deepspeed.InferenceEngine wrapped model. + """ + log_dist("DeepSpeed info: version={}, git-hash={}, git-branch={}".format(__version__, __git_hash__, + __git_branch__), + ranks=[0]) + + # Load config_dict from config first + if config is None: + config = {} + if isinstance(config, str): + with open(config, "r") as f: + config_dict = json.load(f) + elif isinstance(config, dict): + config_dict = config + else: + raise ValueError(f"'config' argument expected string or dictionary, got {type(config)}") + + # Update with values from kwargs, ensuring no conflicting overlap between config and kwargs + overlap_keys = set(config_dict.keys()).intersection(kwargs.keys()) + # If there is overlap, error out if values are different + for key in overlap_keys: + if config_dict[key] != kwargs[key]: + raise ValueError(f"Conflicting argument '{key}' in 'config':{config_dict[key]} and kwargs:{kwargs[key]}") + config_dict.update(kwargs) + + ds_inference_config = DeepSpeedInferenceConfig(**config_dict) + + engine = InferenceEngine(model, config=ds_inference_config) + + return engine + + +def tp_model_init(model, tp_size, dtype, config=None, **kwargs): + """ + Initialize the model for tensor parallelism. + + Args: + model (torch.nn.Module): The model to be initialized. + tp_size (int): The tensor parallelism size. + dtype (torch.dtype): The data type to be used for the model. + + Returns: + torch.nn.Module: The initialized model with tensor parallelism. + """ + # avoid re-entry + if hasattr(model, 'ds_autotp_parsed'): + logger.warning("ds_autotp_parsed' attribute already exists in the model, re-entry is not allowed.") + return + + set_autotp_mode(training=True) + + from deepspeed.runtime.tensor_parallel import TpTrainingManager + # The expected usage here is for it to be invoked by transformers package. + + #TODO: We should provide a custom TP mapping solution without using autoTP + #as modifying the autoTP logic may be more difficult for users compared to configuring it + + model = TpTrainingManager(model=model, tp_size=tp_size, dtype=dtype).module + + setattr(model, 'ds_autotp_parsed', True) + + return model diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6981f53b588133cc939af983dc429aa1adaf3c0f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..140c17b6a46dfe9ef070030dbea92368bcd620c2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/env_report.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/env_report.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..635373e3cc9c1fdb7db868634a31e61bef97f562 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/env_report.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acf9a7a17f9b6267d68e26a26f194ea13453a1fd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f8e1e7da2799ef0cd75e51024adef5c0fc65c06 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/constants.py b/venv/lib/python3.10/site-packages/deepspeed/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..8378e8f2264fd4e99b7309dc21a9844afa19c54e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/constants.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from datetime import timedelta + +############################################# +# Torch distributed constants +############################################# +TORCH_DISTRIBUTED_DEFAULT_PORT = 29500 + +# Default process group wide timeout, if applicable. +# This only applies to the gloo and nccl backends +# (only if NCCL_BLOCKING_WAIT or NCCL_ASYNC_ERROR_HANDLING is set to 1). +# To make an attempt at backwards compatibility with THD, we use an +# extraordinarily high default timeout, given that THD did not have timeouts. +default_pg_timeout = timedelta(minutes=int(os.getenv("DEEPSPEED_TIMEOUT", default=30))) +INFERENCE_GENERIC_MODE = 'generic' +INFERENCE_SPECIALIZED_MODE = 'specialized' + +CROSS_RANK = "CROSS_RANK" +CROSS_SIZE = "CROSS_SIZE" +LOCAL_RANK = 'LOCAL_RANK' diff --git a/venv/lib/python3.10/site-packages/deepspeed/env_report.py b/venv/lib/python3.10/site-packages/deepspeed/env_report.py new file mode 100644 index 0000000000000000000000000000000000000000..37e33b1e873972e10194ef70a7fd4470af6a7e30 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/env_report.py @@ -0,0 +1,195 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +import deepspeed +import subprocess +import argparse +from .ops.op_builder.all_ops import ALL_OPS +from .git_version_info import installed_ops, torch_info, accelerator_name +from deepspeed.accelerator import get_accelerator + +GREEN = '\033[92m' +RED = '\033[91m' +YELLOW = '\033[93m' +END = '\033[0m' +SUCCESS = f"{GREEN} [SUCCESS] {END}" +OKAY = f"{GREEN}[OKAY]{END}" +WARNING = f"{YELLOW}[WARNING]{END}" +FAIL = f'{RED}[FAIL]{END}' +INFO = '[INFO]' + +color_len = len(GREEN) + len(END) +okay = f"{GREEN}[OKAY]{END}" +warning = f"{YELLOW}[WARNING]{END}" + + +def op_report(verbose=True): + max_dots = 23 + max_dots2 = 11 + h = ["op name", "installed", "compatible"] + print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) + print("DeepSpeed C++/CUDA extension op report") + print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) + + print("NOTE: Ops not installed will be just-in-time (JIT) compiled at\n" + " runtime if needed. Op compatibility means that your system\n" + " meet the required dependencies to JIT install the op.") + + print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) + print("JIT compiled ops requires ninja") + ninja_status = OKAY if ninja_installed() else FAIL + print('ninja', "." * (max_dots - 5), ninja_status) + print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) + print(h[0], "." * (max_dots - len(h[0])), h[1], "." * (max_dots2 - len(h[1])), h[2]) + print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) + installed = f"{GREEN}[YES]{END}" + no = f"{YELLOW}[NO]{END}" + for op_name, builder in ALL_OPS.items(): + dots = "." * (max_dots - len(op_name)) + is_compatible = OKAY if builder.is_compatible(verbose) else no + is_installed = installed if installed_ops.get(op_name, + False) and accelerator_name == get_accelerator()._name else no + dots2 = '.' * ((len(h[1]) + (max_dots2 - len(h[1]))) - (len(is_installed) - color_len)) + print(op_name, dots, is_installed, dots2, is_compatible) + print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) + + +def ninja_installed(): + try: + import ninja # noqa: F401 # type: ignore + except ImportError: + return False + return True + + +def nvcc_version(): + import torch.utils.cpp_extension + cuda_home = torch.utils.cpp_extension.CUDA_HOME + if cuda_home is None: + return f"{RED} [FAIL] cannot find CUDA_HOME via torch.utils.cpp_extension.CUDA_HOME={torch.utils.cpp_extension.CUDA_HOME} {END}" + try: + output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True) + except FileNotFoundError: + return f"{RED} [FAIL] nvcc missing {END}" + output_split = output.split() + release_idx = output_split.index("release") + release = output_split[release_idx + 1].replace(',', '').split(".") + return ".".join(release) + + +def installed_cann_path(): + if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): + return os.environ["ASCEND_HOME_PATH"] + return None + + +def installed_cann_version(): + import re + ascend_path = installed_cann_path() + if ascend_path is None: + return f"CANN_HOME does not exist, unable to compile NPU op(s)" + cann_version = "" + for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): + if cann_version: + break + install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] + if install_files: + filepath = os.path.join(dirpath, install_files[0]) + with open(filepath, "r") as f: + for line in f: + if line.find("version") != -1: + cann_version = line.strip().split("=")[-1] + break + return cann_version + + +def get_shm_size(): + try: + shm_stats = os.statvfs('/dev/shm') + except (OSError, FileNotFoundError, ValueError, AttributeError): + return "UNKNOWN", None + + shm_size = shm_stats.f_frsize * shm_stats.f_blocks + shm_hbytes = human_readable_size(shm_size) + warn = [] + if shm_size < 512 * 1024**2: + warn.append( + f" {YELLOW} [WARNING] /dev/shm size might be too small, if running in docker increase to at least --shm-size='1gb' {END}" + ) + if get_accelerator().communication_backend_name() == "nccl": + warn.append( + f" {YELLOW} [WARNING] see more details about NCCL requirements: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#sharing-data {END}" + ) + return shm_hbytes, warn + + +def human_readable_size(size): + units = ['B', 'KB', 'MB', 'GB', 'TB'] + i = 0 + while size >= 1024 and i < len(units) - 1: + size /= 1024 + i += 1 + return f'{size:.2f} {units[i]}' + + +def debug_report(): + max_dots = 33 + + report = [("torch install path", torch.__path__), ("torch version", torch.__version__), + ("deepspeed install path", deepspeed.__path__), + ("deepspeed info", f"{deepspeed.__version__}, {deepspeed.__git_hash__}, {deepspeed.__git_branch__}")] + if get_accelerator().device_name() == 'cuda': + hip_version = getattr(torch.version, "hip", None) + report.extend([("torch cuda version", torch.version.cuda), ("torch hip version", hip_version), + ("nvcc version", (None if hip_version else nvcc_version())), + ("deepspeed wheel compiled w.", f"torch {torch_info['version']}, " + + (f"hip {torch_info['hip_version']}" if hip_version else f"cuda {torch_info['cuda_version']}")) + ]) + elif get_accelerator().device_name() == 'npu': + import torch_npu + report.extend([("deepspeed wheel compiled w.", f"torch {torch_info['version']}"), + ("torch_npu install path", torch_npu.__path__), ("torch_npu version", torch_npu.__version__), + ("ascend_cann version", installed_cann_version())]) + else: + report.extend([("deepspeed wheel compiled w.", f"torch {torch_info['version']} ")]) + + report.append(("shared memory (/dev/shm) size", get_shm_size())) + + print("DeepSpeed general environment info:") + for name, value in report: + warns = [] + if isinstance(value, tuple): + value, warns = value + print(name, "." * (max_dots - len(name)), value) + if warns: + for warn in warns: + print(warn) + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--hide_operator_status', + action='store_true', + help='Suppress display of installation and compatibility statuses of DeepSpeed operators. ') + parser.add_argument('--hide_errors_and_warnings', action='store_true', help='Suppress warning and error messages.') + args = parser.parse_args() + return args + + +def main(hide_operator_status=False, hide_errors_and_warnings=False): + if not hide_operator_status: + op_report(verbose=not hide_errors_and_warnings) + debug_report() + + +def cli_main(): + args = parse_arguments() + main(hide_operator_status=args.hide_operator_status, hide_errors_and_warnings=args.hide_errors_and_warnings) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/git_version_info.py b/venv/lib/python3.10/site-packages/deepspeed/git_version_info.py new file mode 100644 index 0000000000000000000000000000000000000000..70c536d2f78eee24c21db985c43fc9a3c017d4ee --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/git_version_info.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + # This is populated by setup.py + from .git_version_info_installed import * # noqa: F401 # type: ignore +except ModuleNotFoundError: + import os + if os.path.isfile('version.txt'): + # Will be missing from checkouts that haven't been installed (e.g., readthedocs) + version = open('version.txt', 'r').read().strip() + else: + version = "0.0.0" + git_hash = '[none]' + git_branch = '[none]' + + from .ops.op_builder.all_ops import ALL_OPS + installed_ops = dict.fromkeys(ALL_OPS.keys(), False) + accelerator_name = "" + torch_info = {'version': "0.0", "cuda_version": "0.0", "hip_version": "0.0"} + +# compatible_ops list is recreated for each launch +from .ops.op_builder.all_ops import ALL_OPS + +compatible_ops = dict.fromkeys(ALL_OPS.keys(), False) +for op_name, builder in ALL_OPS.items(): + op_compatible = builder.is_compatible() + compatible_ops[op_name] = op_compatible + compatible_ops["deepspeed_not_implemented"] = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/git_version_info_installed.py b/venv/lib/python3.10/site-packages/deepspeed/git_version_info_installed.py new file mode 100644 index 0000000000000000000000000000000000000000..c2692d4544024549a3b18539ed5266f8aebf76f2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/git_version_info_installed.py @@ -0,0 +1,6 @@ +version='0.17.2' +git_hash='unknown' +git_branch='unknown' +installed_ops={'deepspeed_not_implemented': False, 'async_io': False, 'deepspeed_ccl_comm': False, 'deepspeed_shm_comm': False, 'cpu_adam': False, 'fused_adam': False} +accelerator_name='cpu' +torch_info={'version': '0.0', 'bf16_support': False, 'cuda_version': '0.0', 'nccl_version': '0.0', 'hip_version': '0.0'} diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ad95c58f76090188b5d9c73e1e3df4eb2b2c678f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .transformers.ds_transformer import DeepSpeedTransformerInference +from .transformers.clip_encoder import DSClipEncoder diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d649bdff2998d339af87783ae03549774593d51 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad528ca9fe48e0816ae43eadbcb8a2043966348e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2625a7dba14f2a13c08d03c643eb2fe2e3c1f608 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6503be7fd03b02dd22333203a9286b2c897bc344 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bdced73b9855965a4c27fcf0cd01304d712c6016 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4149d2933d447221b0c866366f92c147876c0f68 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b56d2e8c38b202987ca339142aa50791839b3b4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e9894384691904286574b1175f95a7b4ec1a061 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..848a5b48dcf1336a4c80c20a9e21388cdff9e514 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.accelerator import get_accelerator +from ..features.cuda_graph import CUDAGraph + + +class DSClipEncoder(CUDAGraph, torch.nn.Module): + + def __init__(self, enc, enable_cuda_graph=False): + super().__init__(enable_cuda_graph=enable_cuda_graph) + enc.text_model._build_causal_attention_mask = self._build_causal_attention_mask + self.enc = enc + self.device = self.enc.device + self.dtype = self.enc.dtype + self.cuda_graph_created = [False, False] + self.static_inputs = [None, None] + self.static_kwargs = [None, None] + self.static_output = [None, None] + self._cuda_graphs = [None, None] + self.iter = 0 + self.config = self.enc.config + + def _build_causal_attention_mask(self, bsz, seq_len, dtype): + mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype, device=get_accelerator().current_device_name()) + mask.fill_(torch.tensor(torch.finfo(dtype).min)) + mask.triu_(1) + mask = mask.unsqueeze(1) + return mask + + def _graph_replay(self, *inputs, **kwargs): + for i in range(len(inputs)): + if torch.is_tensor(inputs[i]): + self.static_inputs[self.iter][i].copy_(inputs[i]) + for k in kwargs: + if torch.is_tensor(kwargs[k]): + self.static_kwargs[self.iter][k].copy_(kwargs[k]) + get_accelerator().replay_graph(self._cuda_graphs[self.iter]) + return self.static_output[self.iter] + + def forward(self, *inputs, **kwargs): + if self.enable_cuda_graph: + if self.cuda_graph_created[self.iter]: + outputs = self._graph_replay(*inputs, **kwargs) + else: + self._create_cuda_graph(*inputs, **kwargs) + outputs = self._graph_replay(*inputs, **kwargs) + self.iter = (self.iter + 1) % 2 + return outputs + else: + return self.enc(*inputs, **kwargs) + + def _create_cuda_graph(self, *inputs, **kwargs): + # warmup to create the workspace and cublas handle + cuda_stream = torch.cuda.Stream() + cuda_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(cuda_stream): + for i in range(3): + ret = self._forward(*inputs, **kwargs) + torch.cuda.current_stream().wait_stream(cuda_stream) + + # create cuda_graph and assign static_inputs and static_outputs + self._cuda_graphs[self.iter] = get_accelerator().create_graph() + self.static_inputs[self.iter] = inputs + self.static_kwargs[self.iter] = kwargs + + with get_accelerator().capture_to_graph(self._cuda_graphs[self.iter]): + self.static_output[self.iter] = self._forward(*self.static_inputs[self.iter], + **self.static_kwargs[self.iter]) + + self.cuda_graph_created[self.iter] = True + + def _forward(self, *inputs, **kwargs): + return self.enc(*inputs, **kwargs) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_base.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_base.py new file mode 100644 index 0000000000000000000000000000000000000000..b4b113904997fd124aecdfe51529aab6c9e1b25c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_base.py @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch.nn as nn + + +class DeepSpeedTransformerBase(nn.module): + + def __init__(self): + pass + + # this would be the new clean base class that will replace DeepSpeedTransformerInference. + # we currently don't know how this will look like but keeping it here as a placeholder. diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bert.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..13075553ec8b6d818744cc5404ef4db31283ac8e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bert.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedBERTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed BERT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py new file mode 100644 index 0000000000000000000000000000000000000000..c48c7ed58ae5e2746c513a1acea509bbfd129784 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedBloomInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed Bloom Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..3c3baed1f6186a89ed2b852f27ad56a5ab00d9f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedGPTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed GPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py new file mode 100644 index 0000000000000000000000000000000000000000..325bfb4f7e181d8b718748fd5ebf2f4f1ca3af64 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py @@ -0,0 +1,58 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedLlama2Inference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed OPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) + + def forward(self, *args, **kwargs): + + input = args[0] + input_mask = None + get_present = True + + self.allocate_workspace(input.size()) + + # We set the prev key/value to None when there is a prompt + if input.shape[1] > 1: + self.layer_past = None + layer_past = self.layer_past + + input_type = input.dtype + + if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ + and input.dtype == torch.float: + target_dtype = torch.half if self.dtype == torch.int8 else self.dtype + input = input.to(target_dtype) + + with torch.no_grad(): + attention_output, key, value, context_outputtn_ctx, inp_norm = \ + self.attention(input, + input_mask, + None, + layer_past, + get_present, + None, None, None, + self.norm_w, + self.norm_b, + None) + self.layer_past = (key, value) + output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) + + output = output.to(input_type) + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..055ed6d27d7db4cb39cf868d5c189ef517f4a08a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedMegatronGPTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed Megatron GPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_opt.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_opt.py new file mode 100644 index 0000000000000000000000000000000000000000..7bc5524d71c7155fad0b40eafeca490705f3e324 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_opt.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedOPTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed OPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..7e3c81b714c06b84fff1d18daf28ecccbc1db4f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py @@ -0,0 +1,191 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn as nn +from deepspeed import comm as dist +from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp +from deepspeed.utils.logging import log_dist + +from deepspeed.ops.transformer.inference.ds_mlp import DeepSpeedMLP +from deepspeed.ops.transformer.inference.ds_attention import DeepSpeedSelfAttention, BloomSelfAttention +from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp +from deepspeed.accelerator import get_accelerator +import deepspeed +if deepspeed.HAS_TRITON and get_accelerator().is_triton_supported(): + from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP + from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention + + +class DeepSpeedTransformerInference(nn.Module): + """Initialize the DeepSpeed Transformer Layer. + Arguments: + layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, + layer_id will be 0,1,2...23 when each layer object is instantiated + config: An object of DeepSpeedInferenceConfig + mp_group: Model parallelism group initialized on the modeling side. + quantize_scales: This argument groups all the layers' scales used for quantization + quantize_groups: Number of groups used for quantizing the model + merge_count: Shows the number of model-parallel checkpoints merged before running inference. + We use this argument to control the quantization scale for the model parameters if a bigger + quantize-grouping than 1 is used. + mlp_extra_grouping: This flag is used to show a 2x higher number of groups used for the MLP part + of a Transformer layer. We use this feature for quantization to reduce the convergence impact + for specific downstream tasks. + """ + layer_id = 0 + workspace = None + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super(DeepSpeedTransformerInference, self).__init__() + + self.config = config + self.config.layer_id = DeepSpeedTransformerInference.layer_id + DeepSpeedTransformerInference.layer_id += 1 + + data_type = torch.half if self.config.dtype == torch.int8 else self.config.dtype + + if DeepSpeedTransformerInference.layer_id == 1: + log_dist(f"DeepSpeed-Inference config: {self.config.__dict__}", [0]) + if deepspeed.HAS_TRITON and self.config.use_triton: + log_dist(f"Injecting Triton kernels ...", [0]) + + if self.config.bigscience_bloom: + self.attention = BloomSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count) + assert not self.config.use_triton + else: + if deepspeed.HAS_TRITON and self.config.use_triton: + self.attention = TritonSelfAttention(self.config) + else: + self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, + merge_count) + + if deepspeed.HAS_TRITON and self.config.use_triton: + self.mlp = TritonMLP(self.config) + else: + self.mlp = DeepSpeedMLP(self.config, mp_group, quantize_scales, quantize_groups, merge_count, + mlp_extra_grouping) + + device = get_accelerator().current_device_name() # if config.bigscience_bloom else 'cpu' + if self.config.set_empty_params: + self.norm_w = None + self.norm_b = None + else: + self.norm_w = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), + requires_grad=False) + self.norm_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), + requires_grad=False) + self.layer_past = None + self.layer_norm = LayerNormOp() + if DeepSpeedTransformerInference.workspace is None: + DeepSpeedTransformerInference.workspace = WorkspaceOp(self.config) + self._should_allocate_workspace = True + + def allocate_workspace(self, size): + # Allocate memory only on first layer forward + if self.config.layer_id == 0 and self._should_allocate_workspace: + DeepSpeedTransformerInference.workspace.allocate_workspace( + self.config.hidden_size, self.config.heads, size[1], size[0], DeepSpeedTransformerInference.layer_id, + self.config.mp_size, self.config.bigscience_bloom, + dist.get_rank() if dist.is_initialized() else 0, self.config.max_out_tokens, + self.config.min_out_tokens) + self._should_allocate_workspace = False + + @classmethod + def reset_cache(cls): + if cls.workspace is not None: + cls.workspace.reset_cache() + + def forward( + self, + input=None, + input_mask=None, + attention_mask=None, + attn_mask=None, + head_mask=None, + layer_past=None, + get_key_value=False, + get_present=False, + encoder_output=None, + enc_dec_attn_mask=None, + x=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + use_cache=False, + alibi=None, + output_attentions=False, + # TODO(arashb): 'layer_head_mask' and 'past_key_value' are only added to satisfy the OPT models API. + # This needs to be redesigned later! + layer_head_mask=None, + past_key_value=None, + **kwargs): + + if x is not None: + input = x + if "hidden_states" in kwargs: + input = kwargs["hidden_states"] + + input_mask = (input_mask if attn_mask is None else attn_mask) if attention_mask is None else attention_mask + + self.allocate_workspace(input.size()) + + get_present = (get_present or get_key_value or use_cache) + input_mask = input_mask if attention_mask is None else attention_mask + + # We set the prev key/value to None when there is a prompt + if input.shape[1] > 1: + self.layer_past = None + layer_past = layer_past if layer_past is not None else self.layer_past + head_mask = layer_head_mask if layer_head_mask is not None else head_mask + + attn_mask = None + if isinstance(input, tuple): + attn_mask = input[1] + input = input[0] + input_type = input.dtype + + if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ + and input.dtype == torch.float: + target_dtype = torch.half if self.config.dtype == torch.int8 else self.config.dtype + input = input.to(target_dtype) + + with torch.no_grad(): + attention_output, key, value, context_outputtn_ctx, inp_norm = \ + self.attention(input, + input_mask, + head_mask, + layer_past, + get_present, + encoder_hidden_states, + encoder_attention_mask, + output_attentions, + self.norm_w, + self.norm_b, + alibi, + **kwargs) + + presents = (key, value) + self.layer_past = presents if layer_past is None else None + output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) + + if not self.config.pre_layer_norm: + output = self.layer_norm(output, self.norm_w, self.norm_b, self.config.epsilon) + + output = output.to(input_type) + if get_present: + output = (output, presents) + + if self.config.return_single_tuple: + return (output, ) + elif self.config.return_tuple: + return output if type(output) is tuple else (output, attn_mask) + else: + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fc2f979a04b7adb8aaa638d8e2db3d6b2946d98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .replace_module import replace_transformer_layer, revert_transformer_layer, ReplaceWithTensorSlicing, GroupQuantizer, generic_injection +from .module_quantize import quantize_transformer_layer +from .replace_policy import HFBertLayerPolicy +from .layers import LinearAllreduce, LinearLayer, EmbeddingLayer, Normalize, set_autotp_mode +from .policy import DSPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5843b8eae3585d5acb207ab3664d3ce83f10ac14 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2aab19375809def7fc0acd0be7b1063a29f8c59e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1df6bc5312408dca6dae73ff84e08b6e4752823 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26f4cf99c2ae7ccbc2e532007d756b99284e5584 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9139fd19fba43cfa8b41246bfac4b532280aa583 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb8150805609840f2c76bd8e02aeb9950ad2d2e3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e122146f8b2c7f1a7555083340e2077a5c8a27e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c732953190afa13a42574e0676cfe11d66235eff Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbeea9444375c2bfd4a2f963183fdf80e458e982 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f8953f31d1a59d4115ffcf9cdd7633ed576a693 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51334220208926bc087fe27a70c34297d827f7f1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bce51fe3800c11cd8a3d7cb447b0dade3fb13246 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e89180cb01ac69e5f0d646f36f59219734d2c6d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp.py new file mode 100644 index 0000000000000000000000000000000000000000..82cd9042071e7de85cffdcbdba8dad6aa94568bc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp.py @@ -0,0 +1,509 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Automatic Tensor Parallelism +import re + +from torch import nn +from .replace_policy import replace_policies +from typing import Optional +import torch +from deepspeed import comm as dist +from .layers import * +from deepspeed.accelerator import get_accelerator +from .fusedqkv_utils import require_tp_fused_qkvw +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list +from deepspeed.utils import groups +from deepspeed.module_inject.layers import is_autotp_training_mode + + +def move(tensor, device, copy=True): + if tensor.is_meta: + return torch.empty_like(tensor, device=device) + else: + # Using new tensors help in freeing memory (after split for example) was done before by calling clone(). + # Using copy=True instead of clone() will help in case of cpu --> cpu. + # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced. + return tensor.to(device, copy=copy) + + +class ReplaceWithTensorSlicing: + + def __init__(self, mp_group=None, mp_size=1, out_dim=1, in_dim=0): + if mp_group is not None: + self.gpu_index = dist.get_rank(group=mp_group) + else: + self.gpu_index = 0 + self.out_dim = out_dim + self.in_dim = in_dim + self.mp_size = mp_size + + def merge_assert(self, dim1, dim2): + assert dim1 > dim2, \ + 'Merging tensors is not allowed here! Please use deepspeed load_checkpoint\ + for merging your checkpoints before replacing the transformer layer with\ + inference-kernels' + + def strided_copy(self, + dst: Optional[torch.Tensor], + src: Optional[torch.Tensor], + num_splits: int, + int8: bool = False, + allocate_tensor: bool = False): + if src is None: + return src + src_shape = src.shape + dst_shape = dst.shape + + outer_dim = 0 if int8 else -1 + + if allocate_tensor: + dst = torch.empty_like(dst) + + src_split = torch.split(src.data, src.shape[outer_dim] // num_splits, dim=outer_dim) + if (len(src_shape) == 2 and len(dst_shape) == 2): + if src_shape[outer_dim] == dst_shape[self.out_dim]: + try: + dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) + except: + print(dst.shape, src.shape) + exit() + dst = torch.nn.parameter.Parameter(dst, requires_grad=False) + if hasattr(src, 'scale'): + dst.scale = src.scale + return dst + self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) + qkv_size = dst_shape[self.out_dim] // num_splits + qkv_split = [torch.split(src_s, qkv_size, dim=outer_dim) for src_s in src_split] + weight_split = [ + torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=outer_dim) for i in range(len(qkv_split[0])) + ] + dst = dst.reshape(-1).data.copy_(weight_split[self.gpu_index].contiguous().reshape(-1)).reshape( + weight_split[self.gpu_index].shape) + else: + if src_shape[0] == dst_shape[0]: + return torch.nn.parameter.Parameter(src) + qkv_size = dst_shape[0] // num_splits + qkv_split = [torch.split(src_s, qkv_size, dim=0) for src_s in src_split] + bias_split = [torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=0) for i in range(len(qkv_split[0]))] + dst.data.copy_(bias_split[self.gpu_index].contiguous()) + + dst = torch.nn.parameter.Parameter(dst, requires_grad=False) + if hasattr(src, 'scale'): + dst.scale = src.scale + return dst + + def copy(self, dst, src, int8=False, allocate_tensor=False): + if src is None: + return src + assert not dst.data.is_meta # the torch.Tensor.copy_ method used below will silently fail on meta tensors + if allocate_tensor: + dst = torch.empty_like(dst) + outer_dim = 0 if int8 else 1 + inner_dim = 1 if int8 else 0 + src_shape = src.shape + dst_shape = dst.shape + if (len(src_shape) == 2 and len(dst_shape) == 2): + + if src_shape[inner_dim] == dst_shape[self.in_dim] and src_shape[outer_dim] == dst_shape[self.out_dim]: + dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) + else: + if src_shape[inner_dim] != dst_shape[self.in_dim]: + self.merge_assert(src_shape[inner_dim], dst_shape[self.in_dim]) + dst.data.copy_(src[:, self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim]] if inner_dim == 1 else \ + src[self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim], :]) + else: + self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) + dst.data.copy_(src[:, self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim]] if outer_dim == 1 else \ + src[self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim], :]) + else: + if src_shape[0] == dst_shape[0]: + dst = src if src.dtype == dst.dtype else dst.data.copy_(src) + else: + dst.data.copy_(src[self.gpu_index * dst_shape[-1]:(self.gpu_index + 1) * dst_shape[-1]]) + dst = torch.nn.parameter.Parameter(dst, requires_grad=False) + if hasattr(src, 'scale'): + dst.scale = src.scale + return dst + + +class Loading(): + + def is_load_module(module): + load_layers = [nn.Linear, nn.Embedding, nn.LayerNorm] + load_layer_names = [ + "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear", + "MistralRMSNorm", "T5LayerNorm", "MixtralRMSNorm", "Phi3RotaryEmbedding", "Phi3SuScaledRotaryEmbedding", + "Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding", "Qwen2RMSNorm", + "Qwen3RMSNorm", "Qwen3MoeRMSNorm", "DeepseekV2RMSNorm", "DeepseekV3RMSNorm", + "DeepseekV2YarnRotaryEmbedding", "DeepseekV3YarnRotaryEmbedding", "MoEGate" + ] + return module.__class__ in load_layers or module._get_name() in load_layer_names + + def load_buffer(module, state_dict, prefix): + for name in module._buffers.keys(): + if module._buffers[name].data.is_meta: + module._buffers[name] = torch.nn.parameter.Parameter( + data=torch.empty_like(module._buffers[name].data, device="cpu"), + requires_grad=module._buffers[name].data.requires_grad) + if prefix + name in state_dict.keys(): + module._buffers[name].data.copy_(state_dict[prefix + name]) + + def load(module, state_dict, prefix, mp_group=None): + mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group) + if hasattr(module, 'weight'): + if module.weight.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data, device="cpu"), + requires_grad=module.weight.data.requires_grad) + if 'query_key_value' in prefix: + module.weight = mp_replace.strided_copy(module.weight.data, + state_dict[prefix + 'weight'], + num_splits=3) + else: + module.weight = mp_replace.copy(module.weight.data, state_dict[prefix + 'weight']) + else: + if hasattr(module, 'norm') and hasattr(module.norm, 'weight'): + if module.norm.weight.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.norm.weight = torch.nn.parameter.Parameter( + data=torch.empty_like(module.norm.weight.data, device="cpu"), + requires_grad=module.norm.weight.data.requires_grad) + module.norm.weight = mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight']) + + if prefix + 'bias' in state_dict.keys(): + if hasattr(module, 'bias'): + if module.bias.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), + requires_grad=module.bias.data.requires_grad) + module.bias = mp_replace.copy(module.bias, state_dict[prefix + 'bias']) + else: + if hasattr(module, 'norm') and hasattr(module.norm, 'bias'): + if module.norm.bias.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.norm.bias = torch.nn.parameter.Parameter( + data=torch.empty_like(module.norm.bias.data, device="cpu"), + requires_grad=module.norm.bias.data.requires_grad) + module.norm.bias = mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias']) + + +class AutoTP(): + + def __init__(self, + module, + all_reduce_linears, + prefix, + state_dict, + linear_layer_setting, + orig_layer_impl, + keep_module_on_host=False): + self.module = module + self.all_reduce_linears = all_reduce_linears + self.prefix = prefix + self.state_dict = state_dict + + self.mp_size = None + self.mp_group = None + self.linear_layer_setting = linear_layer_setting + self.orig_layer_impl = orig_layer_impl + self.linear_policies = None + self.conv_linear_layer = False + TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host) + + def in_module_list(module, module_list): + for item in module_list: + if type(item).__name__ == type(module).__name__: + return True + return False + + def get_module_list(model): + mlist = [] + for child in model.children(): + if isinstance(child, nn.ModuleList): + for module in child.children(): + if not mlist: + mlist = [module] + elif not AutoTP.in_module_list(module, mlist): + mlist = mlist + [module] + else: + mlist = mlist + AutoTP.get_module_list(child) + return mlist + + def supported(model): + unsupported = ['deberta', 'flaubert', 'fsmt', 'gpt2', 'led', 'longformer', 'xlm', 'xlnet'] + model = str(model) + key = re.search(r": (.*?)Model", model) + if key is None: + key = re.search(r": (.*?)Stack", model) + if key is None: + key = re.match(r"(.*?)Model", model) + assert key is not None, "Not able to determine model policy automatically. Please provide policy." + if key.group(1).lower() in unsupported: + return False + return True + + def get_layers(parent, module): + layer_list = [] + for key, submodule in module._modules.items(): + if isinstance(submodule, nn.Linear): + layer_list = layer_list + [parent + "." + key] + elif isinstance(submodule, nn.LayerNorm) or key == 'LayerNorm' or key == 'layer_norm': + layer_list = layer_list + ["ln"] + else: + layer_list = layer_list + AutoTP.get_layers(key, submodule) + return layer_list + + def update_policy_list(policy_list, new_module, new_gems): + if len(policy_list): + for i, policy in enumerate(policy_list): + # if module already exists in policy, combine gems and remove duplicates + if policy[0] == type(new_module): + new_gems = set(new_gems + policy[1]) + policy_list[i] = tuple([type(new_module), new_gems]) + return policy_list + policy_list.append(tuple([type(new_module), new_gems])) + return policy_list + + def kernel_supported(module_list): + policy = [] + for plcy in replace_policies: + # instantiate a throw-away policy in order to populate the _orig_layer_class + _ = plcy(None) + if isinstance(plcy._orig_layer_class, list): + for orig_layer_class in plcy._orig_layer_class: + policy.append(orig_layer_class) + elif plcy._orig_layer_class is not None: + policy.append(plcy._orig_layer_class) + for child in module_list: + if child.__class__ in policy: + return True + return False + + def tp_parser(model): + policy_list = [] + module_list = [] + layer_list = [] + gem_list = [] + + module_list = AutoTP.get_module_list(model) + assert AutoTP.supported(model), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \ + if AutoTP.kernel_supported(module_list) else "AutoTP not supported for model. Please provide policy." + norm_layer_name_list = ['LayerNorm', 'layer_norm', 'ln_1', 'ln_2'] + #ln_1 , ln_2 for Qwen + for module in module_list: + for key, submodule in module._modules.items(): + if isinstance(submodule, nn.Linear): + layer_list = layer_list + ["." + key] + elif isinstance(submodule, nn.LayerNorm) or key in norm_layer_name_list: + layer_list = layer_list + ["ln"] + else: + layer_list = layer_list + AutoTP.get_layers(key, submodule) + for i, layer in enumerate(layer_list): + if layer == 'ln': + if layer_list[i - 1] != 'ln': + gem_list = gem_list + [layer_list[i - 1]] + elif 'out_proj' in layer: + gem_list = gem_list + [layer] + elif 'o_proj' in layer: + gem_list = gem_list + [layer] + elif 'down_proj' in layer: + gem_list = gem_list + [layer] + elif 'attention.dense' in layer and 'GPTNeoX' in str(model): + gem_list = gem_list + [layer] + elif 'self_attention.dense' in layer and 'falcon' in str( + type(module)): # this is a hack to get the right linear layer for this model! + gem_list = gem_list + [layer] + # Mixtral-7x8b used w2*act(w1*w3) linear. need to replace w2 to linearallreduce. + elif 'w2' in layer and 'Mixtral' in str(type(module)): + gem_list = gem_list + [layer] + elif 'self_attn.dense' in layer and 'Phi' in str(type(module)): + gem_list = gem_list + [layer] + elif 'self_attention.dense' in layer and 'ChatGLM' in str(model): + gem_list = gem_list + [layer] + elif 'dense_4h_to_h' in layer and 'ChatGLM' in str(model): + gem_list = gem_list + [layer] + + layer_list = [] + if gem_list != []: + gem_list = list(set(gem_list)) + policy_list = AutoTP.update_policy_list(policy_list, module, gem_list) + gem_list = [] + assert len(policy_list), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \ + if AutoTP.kernel_supported(module_list) else "Not able to determine model policy automatically. Please provide policy." + return policy_list + + def set_tensor_parallel_config(self, mp_size, mp_group): + + if is_autotp_training_mode(): + self.mp_group = groups.get_tensor_model_parallel_group() + self.mp_size = groups.get_tensor_model_parallel_world_size() + return + + self.mp_size = mp_size + self.mp_group = mp_group + + def _replace(self, child, name, conv_linear_layer): + # This function should clearly define the routing rules for specific layers + # and avoid any complex shard-related logic. + if getattr(child, "replaced", False) == True: + return + + weight_shape = child.weight.shape + mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) + # For TP layer skip, e.g., MoE gate, deepseek low rank layer skip + if "mlp.gate" == name or "q_a_proj" in name or "kv_a_proj_with_mqa" in name or name == "block_sparse_moe.gate" or ( + ('mlp.shared_expert_gate' == name or 'mlp.gate' == name) and 'qwen2_moe' in str(type(self.module))): + return child + # For Yuan model + if 'Yuan' in str(self.module): + if 'v_proj' in name: + return Yuan_LinearLayer(child, self.mp_group) + + elif 'o_proj' in name: + return Yuan_LinearAllreduce(child, self.mp_group) + + # For MLP including chunk layer. + if 'gate_up_proj' in name or ('dense_h_to_4h' in name and 'GLM' in str(self.module)): + return GateUpPack_LinearLayer(child, self.mp_group) + # For Arctic model, bypass to all_reduce replacement for w2 weights + arctic_w2_all_reduce_linear = False + if 'Arctic' in str(self.module) and 'w2' in name: + arctic_w2_all_reduce_linear = True + # For MoE MLP model, e.g., deepseek and jamba + down_proj = False + if 'down_proj' in name: + down_proj = True + if name in self.all_reduce_linears or arctic_w2_all_reduce_linear or down_proj: + + setattr(child, "replaced", True) + if self.conv_linear_layer: + return Conv_LinearALlreduce(child, self.mp_group, name=name) + elif name == "lm_head" or name == 'embed_out': + return LmHeadLinearAllreduce(child, self.mp_group) + + return LinearAllreduce(child, self.mp_group, name=name) + else: + + setattr(child, "replaced", True) + if self.conv_linear_layer: + conv_LinearLayer(child, self.mp_group) + elif require_tp_fused_qkvw(name, self.mp_size): + #Check and handle fused qkv for TP + return fused_LinearLayer(child, self.mp_group, fused_module=self.module) + + return LinearLayer(child, self.mp_group, name=name) + + def _slice_embedding(self, child, name, conv_linear_layer): + if getattr(child, "replaced", False) == True: + return + mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) + + if hasattr(child.weight, 'ds_tensor'): + data = child.weight.ds_tensor.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size), dim=1) + else: + data = child.weight.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size, name), dim=1) + data = data[mp_replace.gpu_index].to(get_accelerator().current_device_name()) + data = torch.nn.parameter.Parameter(data, requires_grad=False) + + new_embedding = nn.Embedding(child.weight.shape[0], get_shard_size(child.weight.shape[1], self.mp_size, name)) + new_embedding.weight.data.copy_(data) + setattr(child, "replaced", True) + return new_embedding + + def update_mp_params(self, child): + if getattr(child, "replaced", False) == True: + return + param_list = [ + "n_heads", "inner_dim", "num_heads", "num_kv", "num_attention_heads", "num_attn_heads", "all_head_size", + "embed_dim", "hidden_size", "num_key_value_heads", "num_kv_heads", "kv_n_heads", "d_model", + "num_attention_heads_per_partition", "num_multi_query_groups_per_partition", "hidden_size_per_partition" + ] + for param in param_list: + if "Yuan" in str(child) and 'embed_dim' in param_list: + param_list.remove('embed_dim') + if hasattr(child, param): + param_val = getattr(child, param) + setattr(child, param, get_shard_size(param_val, self.mp_size)) + setattr(child, "replaced", True) + + def update_linear_policies(self): + self.conv_linear_layer = False + if self.linear_layer_setting is not None: + self.linear_policies = {self.linear_layer_setting[0]: self._replace} + if len(self.linear_layer_setting) == 2: + self.linear_policies.update({self.linear_layer_setting[1]: self._slice_embedding}) + else: + import transformers + if self.orig_layer_impl is transformers.models.gpt2.modeling_gpt2.GPT2Block: + try: + self.conv_linear_layer = True + self.linear_policies = {transformers.pytorch_utils.Conv1D: self._replace} + except ImportError: + self.linear_policies = {nn.Linear: self._replace} + else: + self.linear_policies = {nn.Linear: self._replace, nn.Embedding: self._slice_embedding} + + def _replace_module(self, r_module, prev_name='', prev_class_name=''): + for name, child in r_module.named_children(): + if prev_class_name == "": + class_name = prev_name + elif prev_name == "": + class_name = prev_class_name + else: + class_name = prev_class_name + '.' + prev_name + checking_key = self.prefix + '.' + class_name + '.' + name + '.' if class_name != "" else self.prefix + '.' + name + '.' + if Loading.is_load_module(child) and self.state_dict is not None: + if any(checking_key in item for item in self.state_dict): + Loading.load(child, self.state_dict, checking_key, self.mp_group) + else: + continue + if len(child._buffers) != 0 and self.state_dict is not None: + Loading.load_buffer(child, self.state_dict, checking_key) + if child.__class__ in self.linear_policies: + setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name, + self.conv_linear_layer)) + elif any(isinstance(child, lp) for lp in self.linear_policies): + # Added for falcon model support + # Note: isinstance will account for class inheritance, child.__class__ does not + key = None + for lp in self.linear_policies: + if isinstance(child, lp): + key = lp + break + assert key is not None + setattr(r_module, name, self.linear_policies[key](child, prev_name + '.' + name, + self.conv_linear_layer)) + else: + self.update_mp_params(child) + self._replace_module(child, name, class_name) + return r_module + + def get_model_num_kv_heads(self, config): + num_kv_heads = None + # multi_query_group_num is for chatglm2 & chatglm3 + kv_head_names = [ + 'multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads', + 'attention_heads' + ] + for name in kv_head_names: + if hasattr(config, name): + num_kv_heads = getattr(config, name) + if num_kv_heads is not None: + break + return num_kv_heads + + def _replace_last_linear_module(self, r_module): + if hasattr(r_module, "lm_head"): + name = "lm_head" + child = r_module.lm_head + elif hasattr(r_module, "embed_out"): + name = "embed_out" + child = r_module.embed_out + else: + return r_module + if child.__class__ in self.linear_policies: + setattr(r_module, name, self.linear_policies[child.__class__](child, name, self.conv_linear_layer)) + return r_module diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp_model_utils.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp_model_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a71b1a54d6f6031c18899b1a5294dd5dd963e92d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp_model_utils.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed import comm as dist +import torch +from typing import Optional +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list + + +def build_bloom_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor: + """ + Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it + relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value + `softmax(l+a) = softmax(l)`. Based on + https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 + TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly. + + Args: + Returns tensor shaped (batch_size * num_heads, 1, max_seq_len) + attention_mask (`torch.Tensor`): + Token-wise attention mask, this should be of shape (batch_size, max_seq_len). + num_heads (`int`, *required*): + number of heads + dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`): + dtype of the output tensor + """ + import math + batch_size, seq_length = attention_mask.shape + closest_power_of_2 = 2**math.floor(math.log2(num_heads)) + base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))), + device=attention_mask.device, + dtype=torch.float32) + powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.pow(base, powers) + + if closest_power_of_2 != num_heads: + extra_base = torch.tensor(2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))), + device=attention_mask.device, + dtype=torch.float32) + num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2) + extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) + + # Note: alibi will added to the attention bias that will be applied to the query, key product of attention + # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length) + # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length) + # => the query_length dimension will then be broadcasted correctly + # This is more or less identical to T5's relative position bias: + # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527 + arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :] + alibi = slopes[..., None] * arange_tensor + if dist.is_initialized(): + num_heads_per_rank = get_shard_size(num_heads, dist.get_world_size()) + offset = sum(get_shard_size_list(num_heads, dist.get_world_size())[0:dist.get_rank()]) + alibi = alibi.view(batch_size, num_heads, 1, seq_length) + alibi = alibi[:, offset:num_heads_per_rank + offset, :, :] + return alibi.reshape(batch_size * num_heads_per_rank, 1, seq_length).to(dtype) + else: + return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype) + + +def get_alibi_mask(self, tensor, seq_length_with_past): + mask = self.get_alibi_mask_orig(tensor, seq_length_with_past) + if not self.training and dist.is_initialized(): + num_heads_per_rank = get_shard_size(self.n_head, dist.get_world_size()) + offset = sum(get_shard_size_list(self.n_head, dist.get_world_size())[0:dist.get_rank()]) + mask = mask[offset:num_heads_per_rank + offset, :seq_length_with_past, :seq_length_with_past] + + return mask + + +def build_mpt_atten_bias_tensor(self, + device, + dtype, + attention_mask: Optional[torch.ByteTensor] = None, + prefix_mask: Optional[torch.ByteTensor] = None, + sequence_id: Optional[torch.LongTensor] = None): + (attn_bias, attention_mask) = self._attn_bias_orig(device, + dtype, + attention_mask=attention_mask, + prefix_mask=prefix_mask, + sequence_id=sequence_id) + if dist.is_initialized(): + num_heads_per_rank = get_shard_size(self.config.n_heads, dist.get_world_size()) + offset = sum(get_shard_size_list(self.config.n_heads, dist.get_world_size())[0:dist.get_rank()]) + attn_bias = attn_bias[:, offset:num_heads_per_rank + offset, :, :] + return attn_bias, attention_mask + + +def build_mpt_alibi_tensor(self, num_heads, sequence_length, alibi_bias_max=8, device=None) -> torch.Tensor: + r""" + Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it + relies on a translation invariance of softmax for quick implementation. This implementation has been copied from + the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi: + https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292 + """ + alibi = self.build_mpt_alibi_tensor_orig(num_heads, sequence_length, alibi_bias_max, device) + if dist.is_initialized(): + num_heads_per_rank = int(num_heads / dist.get_world_size()) + offset = dist.get_rank() * num_heads_per_rank + alibi = alibi[offset:num_heads_per_rank + offset, :, :] + return alibi diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..993d14071659303ea6702b0d5182c6d04c931905 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .bert import DS_BERTContainer, HFBertLayerPolicy +from .bloom import DS_BloomContainer, BLOOMLayerPolicy, supported_models +from .distil_bert import DS_DistilBERTContainer, HFDistilBertLayerPolicy +from .gpt2 import DS_GPT2Container, HFGPT2LayerPolicy +from .gptj import DS_GPTJContainer, HFGPTJLayerPolicy +from .gptneo import DS_GPTNEOContainer, HFGPTNEOLayerPolicy +from .gptneox import DS_GPTNEOXContainer, GPTNEOXLayerPolicy +from .llama import DS_LLAMAContainer, LLAMALayerPolicy +from .llama2 import LLAMA2LayerPolicy, DS_LLAMA2Container +from .internlm import DS_InternLMContainer, InternLMLayerPolicy +from .megatron_gpt import DS_MegatronGPTContainer, MegatronLayerPolicy +from .megatron_gpt_moe import DS_MegatronGPTMoEContainer, MegatronMoELayerPolicy +from .opt import DS_OPTContainer, HFOPTLayerPolicy +from .clip import DS_CLIPContainer, HFCLIPLayerPolicy +from .unet import UNetPolicy +from .vae import VAEPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..923e1c264af0d375b68bd6c4a180e8ae65076187 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5fbfdf8125d1e5bfb3901f1b84da00154c26961 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..022a65b51ef6f457403191d975be4da143ca0453 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..142129d93547d351e8cf00390cb5fa3aa631f7d4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71258bd8924ec289ee2e65d6591f38cd18dcb529 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..022541e8b36ad8481f5544b59704bdca0e55c974 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9658edcf096a448a74e6d4f8675f6fb6bb9272b9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9508feb18fece4614eca8ec96abef7d19963fbe Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6fc1d9c3570254acfe29ed62b3a01f5b6a1d077 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20b61ddb249f7381aaa32e8db9530554fc788748 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..038b291a68bf1ce696f6f4e24c4b20fbe7c89bf6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a81336c56abd511a04f5d463869e3df7d33bf975 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cc1f680066364d364cc77a9eb9355225475371a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffe78d415cd1603a1d1f061561fb6417f0c4f880 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a2f8b05af386551ed85536b7603e9816d69fbd6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d11434d939c72588cd4c6167bc66075eaa35edcd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8857af09d7f331995f83877647b1ddc8f68852b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e66279c3d92bc21329c700c01997c5420a67dfba Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa76a9e11111cd96f9e15cf94f5d0dc666457495 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..83e109167ffec669b6b1d77c2996010a70678f27 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py @@ -0,0 +1,322 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Create a container object to save model-specific tensors using the policy file above. +from abc import ABC + +import torch + +import deepspeed +from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig +from deepspeed.accelerator import get_accelerator + +# If the intermediate size attribute is set DEFAULT_INTERMEDIATE_SIZE +# it is assumed the intermediate size is 4x the embedding dimension +DEFAULT_INTERMEDIATE_SIZE = -1 + + +class BaseConvolutionContainer(ABC): + # not implemented + def __init__(self): + pass + + +class BaseTransformerContainer(ABC): + + def __init__(self, policy, config, model_config, layer_id, child): + self.policy = policy + self.config = config + self.model_config = model_config + self.layer_id = layer_id + self.child = child + + self.megatron_v2 = self.policy.is_megatron_v2 + self.scale_attention = self.policy.scale_attention + self.ckpt_load_enabled = False + + # configuration for models. todo: can this be moved to a pydantic model config? + self.hidden_size = None + self.intermediate_size = None + self.num_attention_heads = None + self.mp_size = self.config.tensor_parallel.tp_size + self.pre_layer_norm = self.model_config.do_layer_norm_before if \ + hasattr(self.model_config, 'do_layer_norm_before') else self.policy.pre_attn_norm + self.dtype = self.config.dtype + self.attn_linear_layer = self.policy.linear_layer + self.mlp_linear_layer = self.policy.linear_layer + self.return_tuple = self.config.return_tuple + self.triangular_masking = True + self.local_attention = ((self.model_config.attention_layers[self.layer_id] == "local") if hasattr( + self.model_config, 'attention_layers') else False) + self.window_size = getattr(self.model_config, "window_size", 1) + self.mlp_act_func_type = self.policy.mlp_act_func_type + self.norm_type = self.policy.norm_type + self.training_mp_size = self.config.training_mp_size + self.bigscience_bloom = False + self.max_out_tokens = self.config.max_out_tokens + self.min_out_tokens = self.config.min_out_tokens + self.scale_attn_by_inverse_layer_idx = getattr(self.config, "scale_attn_by_inverse_layer_idx", False) + self.use_mup = self.policy.use_mup + self.return_single_tuple = False + self.rotary_dim = self.get_rotary_dim() + self.mlp_after_attn = (self.rotary_dim is None or self.rotary_dim < 0) + + # Attention tensors + self.qkvw = None + self.qkvb = None + self.dense_w = None + self.dense_b = None + # MLP tensors + self._h4h_w = None + self._h4h_b = None + self._4hh_w = None + self._4hh_b = None + # LayerNorm tensors + self.attn_nw = None + self.attn_nb = None + self.input_nw = None + self.input_nb = None + + self.mp_group = None + self.use_triton = False + + # Triton + self.use_triton = config.use_triton and deepspeed.HAS_TRITON + + def create_ds_model_config(self): + self.set_hidden_heads(*self.policy.get_hidden_heads()) + assert self.num_attention_heads % self.mp_size == 0,\ + "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ + "This is because the attention computation is partitioned evenly among the parallel GPUs." + + self.ds_model_config = DeepSpeedInferenceConfig( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + heads=self.num_attention_heads, + layer_norm_eps=self.layernorm_epsilon, + dtype=self.dtype, + pre_layer_norm=self.pre_layer_norm, + norm_type=self.norm_type, + mp_size=self.mp_size, + return_tuple=self.return_tuple, + triangular_masking=self.triangular_masking, + local_attention=self.local_attention, + window_size=self.window_size, + rotary_dim=self.rotary_dim, + mlp_after_attn=self.mlp_after_attn, + mlp_act_func_type=self.mlp_act_func_type, + training_mp_size=self.training_mp_size, + bigscience_bloom=self.bigscience_bloom, + max_out_tokens=self.max_out_tokens, + min_out_tokens=self.min_out_tokens, + scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, + use_mup=self.use_mup, + return_single_tuple=self.return_single_tuple, + set_empty_params=self.config.set_empty_params, + transposed_mode=self.config.transposed_mode, + use_triton=self.use_triton, + triton_autotune=self.config.triton_autotune) + + if self.use_triton and deepspeed.HAS_TRITON: + from .bert import DS_BERTContainer + if not isinstance(self, DS_BERTContainer): + raise NotImplementedError("Triton kernels are only for BERT-like models yet") + + if not self.config.triton_autotune: + from deepspeed.ops.transformer.inference.triton.matmul_ext import fp16_matmul + fp16_matmul.skip_autotune() + + return self.ds_model_config + + def check_meta_tensor_support(self): + if hasattr(self.qkvw, 'is_meta'): + if self.qkvw.is_meta: + assert self.ckpt_load_enabled, "Meta tensors are not supported for this model currently." + else: + raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") + + def initialize_tensors(self, enable_training=False): + # Set the tensors from policy (user module) to container (DS module) + self.set_attention(*self.policy.attention(enable_training=enable_training)) + self.set_mlp(*self.policy.mlp(enable_training=enable_training)) + self.set_layernorm(*self.policy.layernorm()) + #self.check_meta_tensor_support() + + def convert_to_required_dtype(self): + # Note: converting tensors to fp16 requires that we do it in-place using self.__dict__ and not make a list/dict copy + if self.dtype in [torch.half, torch.bfloat16]: + for k, v in self.__dict__.items(): + # The list comprehension is used for MoE tensor lists + if isinstance(v, list) and all((isinstance(tensor, torch.Tensor) \ + or isinstance(tensor, torch.nn.Parameter)) for tensor in v): + self.__dict__[k] = [moe_tensor.to(self.dtype) for moe_tensor in v] + + if isinstance(v, torch.Tensor) or isinstance(v, torch.nn.Parameter): + self.__dict__[k] = v.to(self.dtype) + + def get_rotary_dim(self): + if hasattr(self.model_config, 'rotary_dim'): + return self.model_config.rotary_dim + if hasattr(self.child, 'attention') and hasattr(self.child.attention, 'rotary_ndims'): + return self.child.attention.rotary_ndims + return -1 + + def set_moe(self, moe=False): + self.moe = moe + + def set_tensor_parallel_config(self, mp_size, mp_group): + self.mp_size = mp_size + self.mp_group = mp_group + + def set_quantization_config(self, quantizer): + self.quantizer = quantizer + + def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon, intermediate_size): + """ + Args: + hidden_size: embedding dimension of the model + num_attention_heads: number of attention heads in the model + epsilon: epsilon value for layer norm (same value used for all norms) + intermediate_size: Size of MLP projection. If `DEFAULT_INTERMEDIATE_SIZE` is passed + it is assumed to be `4 * hidden_size` + """ + self.hidden_size = hidden_size + if intermediate_size == DEFAULT_INTERMEDIATE_SIZE: + self.intermediate_size = 4 * hidden_size + else: + self.intermediate_size = intermediate_size + self.num_attention_heads = num_attention_heads + self.layernorm_epsilon = epsilon + + def set_attention(self, qkvw, qkvb, dense_w, dense_b): + self.qkvw = qkvw + self.qkvb = qkvb + self.dense_w = dense_w + self.dense_b = dense_b + + def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): + self._h4h_w = _h4h_w + self._h4h_b = _h4h_b + self._4hh_w = _4hh_w + self._4hh_b = _4hh_b + + def set_layernorm(self, attn_nw, attn_nb, input_nw, input_nb): + self.attn_nw = attn_nw + self.attn_nb = attn_nb + self.input_nw = input_nw + self.input_nb = input_nb + + def apply_weight_quantization(self): + # quantize attention weights + self.attention_quantization() + + # quantize mlp weights + self.mlp_quantization() + + def attention_quantization(self): + self.module.attention.attn_qkvw = self.quantizer.quantize(self.module.attention.attn_qkvw) + self.module.attention.attn_ow = self.quantizer.quantize(self.module.attention.attn_ow) + + def mlp_quantization(self): + self.module.mlp.inter_w = self.quantizer.quantize(self.module.mlp.inter_w) + self.module.mlp.output_w = self.quantizer.quantize(self.module.mlp.output_w) + + def apply_tensor_parallelism(self, mp_replace): + # setup the new Attention module + self.attention_qkv_mp(mp_replace) + self.attention_o_mp(mp_replace) + + # setup the new MLP module + self.mlp_inter_mp(mp_replace) + self.mlp_output_mp(mp_replace) + + # Apply weight quantization + # TODO(cmikeh2): Re-enable this once verified + #self.apply_weight_quantization() + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + self.module.attention.attn_qkvw = mp_replace.strided_copy(self.module.attention.attn_qkvw, + self.qkvw, + num_splits=3, + int8=reversed_dim) + self.module.attention.attn_qkvb = mp_replace.strided_copy(self.module.attention.attn_qkvb, + self.qkvb, + num_splits=3, + int8=reversed_dim) + + def attention_o_mp(self, mp_replace, reversed_dim=False): + self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, self.dense_w, int8=reversed_dim) + self.module.attention.attn_ob = mp_replace.copy(self.module.attention.attn_ob, + self.dense_b, + int8=reversed_dim, + allocate_tensor=reversed_dim) + + def mlp_inter_mp(self, mp_replace, reversed_dim=False): + self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim) + self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim) + + def mlp_output_mp(self, mp_replace, reversed_dim=False): + self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim) + self.module.mlp.output_b = mp_replace.copy(self.module.mlp.output_b, + self._4hh_b, + int8=reversed_dim, + allocate_tensor=reversed_dim) + + def copy_data_to_new_module(self): + params = {'attn_nw': self.attn_nw, 'attn_nb': self.attn_nb} + for key in params: + if params[key] is None: + setattr(self.module.mlp, key, None) + else: + setattr(self.module.mlp, key, + torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) + + params = {'norm_w': self.input_nw, 'norm_b': self.input_nb} + for key in params: + if params[key] is None: + setattr(self.module, key, None) + else: + setattr(self.module, key, + torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) + + def transpose(self): + self.transpose_attention() + self.transpose_mlp() + + def transpose_attention(self): + if self.attn_linear_layer: + self.qkvw = self.transpose_impl(self.qkvw.data) + self.dense_w = self.transpose_impl(self.dense_w.data) + + def transpose_mlp(self): + if self.mlp_linear_layer: + self._h4h_w = self.transpose_impl(self._h4h_w.data) + self._4hh_w = self.transpose_impl(self._4hh_w.data) + + def transpose_impl(self, data): + data = data.contiguous() + data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) + data = data.reshape(data.shape[-1], data.shape[-2]) + data.to(get_accelerator().current_device_name()) + return data + + def get_all_params(self): + params = [ + self.attn_nw, + self.attn_nb, + self.input_nw, + self.input_nb, + ] + + params.extend(self.get_attn_params()) + params.extend(self.get_mlp_params()) + + return params + + def get_attn_params(self): + return [self.qkvw, self.qkvb, self.dense_w, self.dense_b] + + def get_mlp_params(self): + return [self._h4h_w, self._h4h_b, self._4hh_w, self._4hh_b] diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base_moe.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..4be1b849ba70da04b6b08ea011c27bbbcf96b8bf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base_moe.py @@ -0,0 +1,130 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Create a container object to save model-specific tensors using the policy file above. +from .base import * +from deepspeed import comm as dist +import deepspeed.ops.transformer as transformer_inference +from deepspeed.accelerator import get_accelerator + + +class BaseTransformerMoEContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + # Call the init function of the parent class to initialize the tensors and configs from parent class + super().__init__(**kwargs) + + self.num_experts = self.policy.get_num_experts() + self.ep_world_size = dist.get_world_size() + self.local_ep_size = 1 if self.num_experts < self.ep_world_size else self.num_experts // self.ep_world_size + + self.layer_norm_eps = self.config.layer_norm_eps if hasattr(self.config, 'layer_norm_eps') else 1e-12, + + # MoE models will have a list of mlp related tensors + self._h4h_w = [] + self._h4h_b = [] + self._4hh_w = [] + self._4hh_b = [] + + # Residual MoE needs extra parameters + self._res_h4h_w = None + self._res_h4h_b = None + self._res_4hh_w = None + self._res_4hh_b = None + self._res_coef = None + + def create_ds_model_config(self): + self.set_hidden_heads(*self.policy.get_hidden_heads()) + assert self.num_attention_heads % self.mp_size == 0,\ + "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ + "This is because the attention computation is partitioned evenly among the parallel GPUs." + + self.ds_model_config = transformer_inference.DeepSpeedMoEInferenceConfig( + hidden_size=self.hidden_size, + heads=self.num_attention_heads, + layer_norm_eps=self.layer_norm_eps, + fp16=self.fp16, + pre_layer_norm=self.pre_layer_norm, + mp_size=self.mp_size, + q_int8=self.quantize, + moe_experts=self.local_ep_size, + global_experts=self.num_experts, + mlp_type=self.config.moe.type, + scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, + ) + + return self.ds_model_config + + def initialize_tensors(self): + # Set the tensors from policy (user module) to container (DS module) + self.set_attention(*self.policy.attention()) + self.set_mlp(self.config.moe.type) + self.set_layernorm(*self.policy.layernorm()) + + def set_mlp(self, config_moe_type): + if config_moe_type == 'standard': + self._h4h_w, self._h4h_b, \ + self._4hh_w, self._4hh_b = self.policy.mlp() + else: + self._h4h_w, self._h4h_b, self._4hh_w, \ + self._4hh_b, self._res_h4h_w, self._res_h4h_b, \ + self._res_4hh_w, self._res_4hh_b, \ + self._res_coef = self.policy.mlp(config_moe_type) + + def transpose(self): + self.transpose_attention() + self.transpose_mlp() + + if self.config.moe.type == 'residual': + self.transpose_residual() + + def transpose_mlp(self): + self._h4h_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._h4h_w] + self._4hh_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._4hh_w] + + def transpose_residual(self): + self._res_h4h_w.data = self.transpose_impl(self._res_h4h_w.data) + self._res_4hh_w.data = self.transpose_impl(self._res_4hh_w.data) + self._res_coef.data = self.transpose_impl(self._res_coef.data) + + def apply_tensor_parallelism(self, mp_replace): + # setup the new Attention module + self.attention_qkv_mp(mp_replace) + self.attention_o_mp(mp_replace) + + # quantize attention weights + self.attention_quantization() + + # setup the new MLP module + self.mlp_mp() + + def mlp_mp(self): + gpu_index = dist.get_rank() + for ep_index in range(self.local_ep_size): + # mlp inter + self.module.mlp[ep_index].inter_w.data = self._h4h_w[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + self.module.mlp[ep_index].inter_b.data = self._h4h_b[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + + # mlp output + self.module.mlp[ep_index].output_w.data = self._4hh_w[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + self.module.mlp[ep_index].output_b.data = self._4hh_b[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + + def copy_data_to_new_module(self): + self.module.attn_nw.data = self.attn_nw.to(get_accelerator().current_device_name()) + self.module.attn_nb.data = self.attn_nb.to(get_accelerator().current_device_name()) + + self.module.norm_w.data.copy_(self.input_nw.to(get_accelerator().current_device_name())) + self.module.norm_b.data.copy_(self.input_nb.to(get_accelerator().current_device_name())) + + if self.config.moe.type == 'residual': + self.module.res_mlp.inter_w.data = self._res_h4h_w.to(get_accelerator().current_device_name()) + self.module.res_mlp.inter_b.data = self._res_h4h_b.to(get_accelerator().current_device_name()) + self.module.res_mlp.output_w.data = self._res_4hh_w.to(get_accelerator().current_device_name()) + self.module.res_mlp.output_b.data = self._res_4hh_b.to(get_accelerator().current_device_name()) + self.module.res_coef.data = self._res_coef.to(get_accelerator().current_device_name()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bert.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bert.py new file mode 100644 index 0000000000000000000000000000000000000000..20ae575f45144733a82b609eed21284746cb95d0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bert.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy + + +class DS_BERTContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + self.return_tuple = True + self.triangular_masking = False + self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFBertLayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=False): + super().__init__(inference, pre_attn_norm=False) + self.client_module = client_module + self.cuda_graph_supported = True + + if HFBertLayerPolicy._orig_layer_class is None: + try: + import transformers + HFBertLayerPolicy._orig_layer_class = [ + transformers.models.bert.modeling_bert.BertLayer, + transformers.models.roberta.modeling_roberta.RobertaLayer + ] + except: + HFBertLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + if self.pre_attn_norm: + attention_layernorm = self.client_module.PostAttentionLayerNorm + else: + attention_layernorm = self.client_module.attention.output.LayerNorm + return self.client_module.attention.self.query.weight.shape[1], \ + self.client_module.attention.self.num_attention_heads, \ + attention_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.attention.self.query.weight + qb = self.client_module.attention.self.query.bias + kw = self.client_module.attention.self.key.weight + kb = self.client_module.attention.self.key.bias + vw = self.client_module.attention.self.value.weight + vb = self.client_module.attention.self.value.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.attention.output.dense.weight, \ + self.client_module.attention.output.dense.bias, \ + + def mlp(self, enable_training=False): + if self.pre_attn_norm: + intermediate_ff = self.client_module.intermediate.dense_act + else: + intermediate_ff = self.client_module.intermediate.dense + + return intermediate_ff.weight, intermediate_ff.bias, \ + self.client_module.output.dense.weight, \ + self.client_module.output.dense.bias + + def layernorm(self): + if self.pre_attn_norm: + attention_layernorm = self.client_module.PostAttentionLayerNorm + transformer_layernorm = self.client_module.PreAttentionLayerNorm + else: + attention_layernorm = self.client_module.attention.output.LayerNorm + transformer_layernorm = self.client_module.output.LayerNorm + return attention_layernorm.weight, \ + attention_layernorm.bias, \ + transformer_layernorm.weight, \ + transformer_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bloom.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bloom.py new file mode 100644 index 0000000000000000000000000000000000000000..7a9b9ca2065bf969efba949588bc9ac4c95ec8dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bloom.py @@ -0,0 +1,143 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.hybrid_engine import HybridEngineContainer +from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy + +from ..policy import maybe_get_lora + +supported_models = {None} + + +class DS_BloomContainer(MetaTensorContainer, HybridEngineContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + # Check transformers version, error if > 4.43.4 (breaks at 4.44.0) + from importlib.metadata import version + v_transformers = version('transformers') + vers = v_transformers.split('.') + major = int(vers[0]) + minor = int(vers[1]) + if major > 4 or (major == 4 and minor > 43): + import sys + sys.exit( + f"Transformers version {v_transformers} exceeds version 4.43.4! After transformers version 4.43.4, BLOOM inference with DeepSpeed is no longer supported." + ) + + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + self.bigscience_bloom = True + self.triangular_masking = False + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + self.module = DeepSpeedBloomInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + self.module.config.invert_mask = False + return self.module + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + self.module.attention.attn_qkvw = mp_replace.copy(self.module.attention.attn_qkvw, self.qkvw) + self.module.attention.attn_qkvb = mp_replace.copy(self.module.attention.attn_qkvb, self.qkvb) + + def get_lora_matched_pair(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)] + return ret + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, self.policy. + client_module.self_attention.query_key_value, self.policy.client_module.self_attention.dense + ] + ] + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attention.query_key_value.weight', \ + 'self_attention.query_key_value.bias', \ + 'self_attention.dense.weight', \ + 'self_attention.dense.bias', \ + 'mlp.dense_h_to_4h.weight', \ + 'mlp.dense_h_to_4h.bias', \ + 'mlp.dense_4h_to_h.weight', \ + 'mlp.dense_4h_to_h.bias', \ + 'post_attention_layernorm.weight', \ + 'post_attention_layernorm.bias', \ + 'input_layernorm.weight', \ + 'input_layernorm.bias' + ) + for i in range(0, 2): + maybe_copy(module.attention, + sd, + weight_quantizer, + mp_replace, + transformer_param_names[i], + prefix + param_names[i], + qkv=True, + megatron_v2=self.policy.is_megatron_v2, + split_qkv=self.policy.split_qkv) + for i in range(2, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(4, 10): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(10, 12): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i]) + + +class BLOOMLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True, use_load_prefix=True, split_qkv=False): + super().__init__(inference, linear_layer=True, use_load_prefix=use_load_prefix, split_qkv=split_qkv) + self.client_module = client_module + try: + import transformers + BLOOMLayerPolicy._orig_layer_class = transformers.models.bloom.modeling_bloom.BloomBlock + global supported_models + supported_models.update({transformers.models.bloom.modeling_bloom.BloomModel}) + except Exception as e: + print(f"WARNING! Setting BLOOMLayerPolicy._orig_layer_class to None due to Exception: {e}") + BLOOMLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.self_attention.hidden_size, \ + self.client_module.self_attention.num_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + return self.client_module.self_attention.query_key_value.weight, \ + self.client_module.self_attention.query_key_value.bias, \ + self.client_module.self_attention.dense.weight, \ + self.client_module.self_attention.dense.bias, + + def mlp(self, enable_training=False): + return self.client_module.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.dense_4h_to_h.bias + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + self.client_module.post_attention_layernorm.bias, \ + self.client_module.input_layernorm.weight, \ + self.client_module.input_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/clip.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..afe4a76086d80ceec7dce15c9cfb2a9b8718b0c4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/clip.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy + + +class DS_CLIPContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFCLIPLayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=False): + super().__init__(inference, pre_attn_norm=True, scale_attention=True) + self.client_module = client_module + self.cuda_graph_supported = True + + if HFCLIPLayerPolicy._orig_layer_class is None: + try: + import transformers + HFCLIPLayerPolicy._orig_layer_class = transformers.models.clip.modeling_clip.CLIPEncoderLayer + except: + HFCLIPLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.self_attn.q_proj.weight.shape[1], \ + self.client_module.self_attn.num_heads, \ + self.client_module.layer_norm1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + qb = self.client_module.self_attn.q_proj.bias + kw = self.client_module.self_attn.k_proj.weight + kb = self.client_module.self_attn.k_proj.bias + vw = self.client_module.self_attn.v_proj.weight + vb = self.client_module.self_attn.v_proj.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.self_attn.out_proj.weight, \ + self.client_module.self_attn.out_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.fc1.weight, \ + self.client_module.mlp.fc1.bias, \ + self.client_module.mlp.fc2.weight, \ + self.client_module.mlp.fc2.bias + + def layernorm(self): + return self.client_module.layer_norm2.weight, \ + self.client_module.layer_norm2.bias, \ + self.client_module.layer_norm1.weight, \ + self.client_module.layer_norm1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/distil_bert.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/distil_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..ecd0562438b5ac634cf8b4536fd3413d0f9ed9d8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/distil_bert.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy + + +class DS_DistilBERTContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + self.triangular_masking = False + self.return_single_tuple = True + self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFDistilBertLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=False, preln=False): + super().__init__(inference) + self.client_module = client_module + self.preln = preln + self.cuda_graph_supported = True + if HFDistilBertLayerPolicy._orig_layer_class is None: + try: + import transformers + HFDistilBertLayerPolicy._orig_layer_class = [ + transformers.models.distilbert.modeling_distilbert.TransformerBlock, + ] + except: + HFDistilBertLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attention.q_lin.weight.shape[1], \ + self.client_module.attention.n_heads, \ + self.client_module.sa_layer_norm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.attention.q_lin.weight + qb = self.client_module.attention.q_lin.bias + kw = self.client_module.attention.k_lin.weight + kb = self.client_module.attention.k_lin.bias + vw = self.client_module.attention.v_lin.weight + vb = self.client_module.attention.v_lin.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.attention.out_lin.weight, \ + self.client_module.attention.out_lin.bias + + def mlp(self, enable_training=False): + intermediate_ff = self.client_module.ffn.lin1 + + return intermediate_ff.weight, intermediate_ff.bias, \ + self.client_module.ffn.lin2.weight, \ + self.client_module.ffn.lin2.bias + + def layernorm(self): + attention_layernorm = self.client_module.sa_layer_norm + transformer_layernorm = self.client_module.output_layer_norm + return attention_layernorm.weight, \ + attention_layernorm.bias, \ + transformer_layernorm.weight, \ + transformer_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fc2eb2a65531e61bcea078c0035bd2ee6ee861a6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .gated_mlp import HybridGatedMLPContainer +from .megatron import MegatronContainer +from .meta_tensor import MetaTensorContainer +from .split_qkv import HybridSplitQKVContainer diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2e0fefb8cb1b38c9a720ca7a0b7144c35920640 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e77e5a265dc7994eb429128f49256b1b3324c28a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6682782ba7778a60b6d7654a56370884d979e186 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..020a773cc8a8fcdfbd3808cabe9858109a70c93e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..927c08f789979738e6f26c52415f7f9a01c99f6f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50e60628f15732c3f9c079b1a887a89919da0ab8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb0ffb5ad10ff593f4bdaf7dbd58d0c318694835 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..24f0826db14ed08f373a2824f845ac6f9d7d9508 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import abstractmethod + +from .hybrid_engine import HybridEngineContainer + + +class HybridGatedMLPContainer(HybridEngineContainer): + """ + The HybridGatedMLPContainer supports models for which the first MLP layer + is represented with two separate weights, one for the activation function + and one for the gating function. + """ + + def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): + super().set_mlp(_h4h_w, _h4h_b, _4hh_w, _4hh_b) + self.set_mlp_gate() + + @abstractmethod + def set_mlp_gate(self): + """ + In `set_mlp_gate`, it is necessary to populate the following variables (where appropriate) + for the given model: + self.inter_up_w: inter up weight + self.inter_up_b: inter up bias + self.inter_gate_w: inter gate weight + self.inter_gate_b: inter gate bias + If the parameter does not exist in the original model, set the attribute to None. + """ + raise NotImplementedError("A set_mlp_gate() function must be defined in the model container \ + in order to set the unfused inter up and gate tensors.") + + def mlp_inter_mp(self, mp_replace, reversed_dim=False): + # Only need to alter behavior if we can't do the normal destructive copy + if self.module.mlp.inter_w is None: + params = [ + (self.module.mlp.inter_up_w, self.inter_up_w), + (self.module.mlp.inter_up_b, self.inter_up_b), + (self.module.mlp.inter_gate_w, self.inter_gate_w), + (self.module.mlp.inter_gate_b, self.inter_gate_b), + ] + for dst, src in params: + dst = mp_replace.copy(dst[:self.inter_up_w.shape[0] // mp_replace.mp_size], + src, + int8=reversed_dim, + allocate_tensor=reversed_dim) if src is not None else None + else: + self.module.mlp.inter_w = mp_replace.strided_copy(self.module.mlp.inter_w, + self._h4h_w, + num_splits=2, + int8=reversed_dim) + self.module.mlp.inter_b = mp_replace.strided_copy(self.module.mlp.inter_b, + self._h4h_b, + num_splits=2, + int8=reversed_dim) + + def release_mlp(self): + super().release_mlp() + gated_mlp_params = [ + (self.module.mlp.inter_up_w, self.inter_up_w), + (self.module.mlp.inter_up_b, self.inter_up_b), + (self.module.mlp.inter_gate_w, self.inter_gate_w), + (self.module.mlp.inter_gate_b, self.inter_gate_b), + ] + + self._release_params(gated_mlp_params) + + def reset_mlp(self): + self._h4h_w.data[:self.inter_up_w.shape[0]] = self.inter_up_w.data + self._h4h_w.data[self.inter_up_w.shape[0]:] = self.inter_gate_w.data + + if self.inter_up_b is not None: + self._h4h_b.data[:self.inter_up_b.shape[0]] = self.inter_up_b.data + self._h4h_b.data[self.inter_up_b.shape[0]:] = self.inter_gate_b.data + + inter_data = [self.inter_up_w.data, self.inter_gate_w.data] + if self.inter_up_b is not None: + inter_data.extend([self.inter_up_b.data, self.inter_gate_b.data]) + + self.inter_up_w.data = self._h4h_w.data[:self.inter_up_w.shape[0]] + self.inter_gate_w.data = self._h4h_w.data[self.inter_up_w.shape[0]:] + + if self.inter_up_b is not None: + self.inter_up_b.data = self._h4h_b.data[:self.inter_up_b.shape[0]] + self.inter_gate_b.data = self._h4h_b.data[self.inter_up_b.shape[0]:] + + for data in inter_data: + del data + + def set_mlp_params_wo_copy(self, Z3_enabled=False): + self.module.mlp.output_w = self._4hh_w + self.module.mlp.output_b = self._4hh_b + + if not Z3_enabled: + # In initialize_tensors, we create a fused inter projection with the appropriate shape + # and copy the up projection and gate projection into it + self.module.mlp.inter_w = self._h4h_w + self.module.mlp.inter_b = self._h4h_b + + self.inter_up_w.data = self._h4h_w[:self.inter_up_w.shape[0], :] + self.inter_gate_w.data = self._h4h_w[self.inter_up_w.shape[0]:, :] + + if self.inter_up_b is not None: + self.inter_up_b.data = self._h4h_b[:self.inter_up_w.shape[0]] if self._h4h_b is not None else None + self.inter_gate_b.data = self._h4h_b[self.inter_up_w.shape[0]:] if self._h4h_b is not None else None + else: + self.module.mlp.inter_up_w = self.inter_up_w + self.module.mlp.inter_up_b = self.inter_up_b + self.module.mlp.inter_gate_w = self.inter_gate_w + self.module.mlp.inter_gate_b = self.inter_gate_b + + def get_mlp_params(self): + params = super().get_mlp_params() + params.extend([self.inter_up_w, self.inter_up_b, self.inter_gate_w, self.inter_gate_b]) + return params diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..3f702abcf319a6db997b6e2607dcba26f3440841 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py @@ -0,0 +1,212 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from typing import List, Tuple + +import torch + + +class HybridEngineContainer(ABC): + """ + This container identifies which methods need to be overridden in addition to + the base container to enable use in the RLHF pipeline. These methods are not + necessary for inference alone. + + NOTE: If you are using this feature with a container that + also inherits from `MetaTensorContainer`, ensure that `MetaTensorContainer` + is inherited before `HybridEngineContainer` in the class definition. + """ + + def initialize_tensors(self, enable_training=False): + """ + Same purposes as the base container, but also grabs the hooks for any LoRA + parameters. If it's necessary to override specific sub-components of the model, + it's best to augment the specific `set_[component]` itself rather than modifying + the `initialize_tensors` method. See the `HybridSplitQKVContainer` for an example. + """ + super().initialize_tensors(enable_training=enable_training) + self.set_lora_params() + + def transform_for_training(self): + """ + If the views on certain parameters are largely incompatible, it may be necessary to do + more substantial transformations to the parameters. This method should be overridden to + transform the inference format to what is necessary for training. + """ + pass + + def transform_for_inference(self): + """ + If the views on certain parameters are largely incompatible, it may be necessary to do + more substantial transformations to the parameters. This method should be overridden to + transform the training format to what is necessary for inference. + """ + pass + + @abstractmethod + def set_lora_params(self): + """ + If available, set the LoRA parameters for the module. An implementation + for this would iterate over all parameters of the model and use the `maybe_get_lora` helper + method to check if the parameter does in fact have any LoRA params. + """ + raise NotImplementedError("A set_lora_params() function must be defined for the relevant parameters.") + + @abstractmethod + def get_lora_matched_pair(self): + """Get the pair of lora params and its matched model parameters.""" + raise NotImplementedError("get_lora_matched_pair() must be defined for the relevant parameters.") + + def fuse_lora(self): + """Fuse the LoRA parameters for the inference mode.""" + for maybe_lora_param, param in self.get_lora_matched_pair(): + if len(maybe_lora_param) == 3: + lora_right_weight, \ + lora_left_weight, \ + lora_scaling = maybe_lora_param + param.data += lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) + + def unfuse_lora(self): + """Unfuse the LoRA parameters for the training mode.""" + for maybe_lora_param, param in self.get_lora_matched_pair(): + if len(maybe_lora_param) == 3: + lora_right_weight, \ + lora_left_weight, \ + lora_scaling = maybe_lora_param + param.data -= lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) + + def apply_tensor_parallelism(self, mp_replace, reversed_dim=False): + """ + Add support for reversed dim in tensor parallelism. If necessary, override + the called methods to handle partitioned weights (i.e. if qkv is split, override + the `attention_qkv_mp` method). If the model component is not split, it should + be safe to use the default implementation. + """ + # Setup the new Attention module + self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim) + self.attention_o_mp(mp_replace, reversed_dim=reversed_dim) + + # Setup the new MLP module + self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim) + self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim) + + # Apply weight quantization + # TODO(cmikeh2): Re-enable this once verified + #self.apply_weight_quantization() + + def _release_params(self, param_pairs: List[Tuple[torch.Tensor, torch.Tensor]]): + """ + Helper for `release_[component]` methods. Accepts a list of tuples where the first + element is the module param that needs to be deleted, and the second is the reassignment + from the container. + """ + for module_param, container_param in param_pairs: + if module_param is not None: + del module_param + module_param = container_param + + def release_memory(self): + """ + Delete module parameters if they exist and point them back to the container. The primary + purpose of this is for TP-inference with ZeRO-3. In this scenario, we need to delete the + parameters we've created for inference to free their memory. + """ + general_params = [ + (self.module.attention.attn_ow, self.dense_w), + (self.module.attention.attn_ob, self.dense_b), + (self.module.mlp.attn_nw, self.attn_nw), + (self.module.mlp.attn_nb, self.attn_nb), + (self.module.norm_w, self.input_nw), + (self.module.norm_b, self.input_nb), + ] + + self._release_params(general_params) + + self.release_qkv() + self.release_mlp() + + def release_qkv(self): + """ + Release for QKV parameters (as well as any aliases). + """ + qkv_params = [ + (self.module.attention.attn_qkvw, self.qkvw), + (self.module.attention.attn_qkvb, self.qkvb), + ] + + self._release_params(qkv_params) + + def release_mlp(self): + """ + Release for MLP parameters (as well as any aliases). + """ + mlp_params = [ + (self.module.mlp.inter_w, self._h4h_w), + (self.module.mlp.inter_b, self._h4h_b), + (self.module.mlp.output_w, self._4hh_w), + (self.module.mlp.output_b, self._4hh_b), + ] + + self._release_params(mlp_params) + + def reset_params(self): + """ + The purpose of reset params is to get the weights from the FP16 training + copy of the model and copy to them to contiguous inference view. This only needs + to be performed when the container parameters cannot be used directly for inference. + """ + self.reset_qkv() + self.reset_mlp() + + def reset_qkv(self): + """ + Perform any necessary resets of the model parameters for the QKV components. + """ + pass + + def reset_mlp(self): + """ + Perform any necessary resets of the model parameters for the MLP components. + """ + pass + + def get_lora_params(self): + """ + Return a list of all parameters that would have LoRA for the module. + """ + if not hasattr(self, "lora_params"): + self.set_lora_params() + return self.lora_params + + def set_params_wo_copy(self, Z3_enabled=False): + """ + Rather than copying into, set the parameters directly. This is necessary to provide + an inexpensive (low-memory-overhead) view onto the FP16 forward weights. + """ + self.module.mlp.attn_nw = self.attn_nw + self.module.mlp.attn_nb = self.attn_nb + self.module.norm_w = self.input_nw + self.module.norm_b = self.input_nb + self.set_attn_params_wo_copy(Z3_enabled=Z3_enabled) + self.set_mlp_params_wo_copy(Z3_enabled=Z3_enabled) + + def set_attn_params_wo_copy(self, **kwargs): + """ + Narrower sub-method for finer grained overriding. + """ + self.module.attention.attn_ow = self.dense_w + self.module.attention.attn_ob = self.dense_b + self.module.attention.attn_qkvw = self.qkvw + self.module.attention.attn_qkvb = self.qkvb + + def set_mlp_params_wo_copy(self, **kwargs): + """ + Narrower sub-method for finer grained overriding. + """ + self.module.mlp.inter_w = self._h4h_w + self.module.mlp.inter_b = self._h4h_b + self.module.mlp.output_w = self._4hh_w + self.module.mlp.output_b = self._4hh_b diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py new file mode 100644 index 0000000000000000000000000000000000000000..d40f2a6b57e86087dc13ff2b74a389c67826fa76 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from .hybrid_engine import HybridEngineContainer +from .megatron import MegatronContainer + + +class HybridMegatronContainer(MegatronContainer, HybridEngineContainer): + + def _align_qkv(self, x: torch.Tensor): + """ + Internal helper for accepting the head-contiguous weight matrix and chunking + the query, key, and value components. + """ + attention_head_size = x.shape[0] // self.num_attention_heads + new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] + x_1 = x.view(*new_x_shape) + div_dim = len(x_1.size()) - 2 if len(x.shape) == 2 else -1 + (q, k, v) = torch.split(x_1, (x_1.shape[div_dim] // 3), dim=div_dim) + if len(q.shape) > 2: + x.data.copy_( + torch.cat((q.reshape(-1, q.shape[-1]), k.reshape(-1, q.shape[-1]), v.reshape(-1, q.shape[-1])), + dim=0).reshape(x.shape)) + else: + x.data.copy_(torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)) + + def transform_for_inference(self) -> None: + """ + Overrides the HybridEngineContainer implementation. + + The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V + are sequential in memory. This is different from the default layout in which all of the Qs + are sequential, followed by all of the Ks, and then all of the Vs. Here, we take the default + layout and transform it to the inference layout. + """ + if hasattr(self.qkvw, 'ds_id'): + from deepspeed.runtime.zero import GatheredParameters + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + param_list = [self.qkvw, self.qkvb] + non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + with GatheredParameters(non_active_params): + self._align_qkv(self.qkvw) + self._align_qkv(self.qkvb) + else: + self._align_qkv(self.qkvw) + self._align_qkv(self.qkvb) + + def _partition_qkv(self, x: torch.Tensor): + """ + Internal helper for taking contiguous QKV and partitioning it for contiguous + heads. + """ + q_k_v = torch.split(x, (x.shape[0] // 3), dim=0) + attention_head_size = q_k_v[0].shape[0] // self.num_attention_heads + new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] + q, k, v = [data.view(*new_x_shape) for data in q_k_v] + if len(q.shape) > 2: + x.data.copy_(torch.cat((q, k, v), dim=-2).reshape(-1, q.shape[-1])) + else: + x.data.copy_(torch.cat((q, k, v), dim=-1).reshape(-1)) + + def transform_for_training(self): + """ + Overrides the HybridEngineContainer implementation. + + The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V + are sequential in memory. This is different from the default layout in which all of the Qs + are sequential, followed by all of the Ks, and then all of the Vs. This function takes the inference format and reverts it back to the default format. + """ + # If parameter is distributed, handle gathering it + if hasattr(self.qkvw, 'ds_id'): + from deepspeed.runtime.zero import GatheredParameters + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + param_list = [self.qkvw, self.qkvb] + non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + with GatheredParameters(non_active_params): + self._partition_qkv(self.qkvw) + self._partition_qkv(self.qkvb) + else: + self._partition_qkv(self.qkvw) + self._partition_qkv(self.qkvb) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/megatron.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/megatron.py new file mode 100644 index 0000000000000000000000000000000000000000..4daccf7d7c8d4db06bd65243bd83a5f1fd50df2e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/megatron.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from abc import ABC + + +class MegatronContainer(ABC): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.megatron_v2 = self.policy.is_megatron_v2 + + def _align_qkv_transposed(self, x): + attention_head_size = x.shape[-1] // self.num_attention_heads + new_x_shape = x.size()[:-1] + (self.num_attention_heads, attention_head_size) + x_1 = x.view(*new_x_shape) + (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=(x_1.dim() - 1)) + if len(q.shape) > 2: + return torch.cat((q.reshape(q.shape[0], -1), k.reshape(q.shape[0], -1), v.reshape(q.shape[0], -1)), + dim=-1).reshape(x.shape) + else: + return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) + + def transpose(self): + super().transpose() + if self.megatron_v2: + self.qkvw = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvw).contiguous()) + self.qkvb = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvb).contiguous()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..57b136663be370863dad81cc9d6a4a29425376b3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py @@ -0,0 +1,70 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from packaging import version as pkg_version +import torch + + +class MetaTensorContainer(ABC): + """ + NOTE: If you are using this feature with a container that + also inherits from `HybridEngineContainer`, ensure that `MetaTensorContainer` + is inherited before `HybridEngineContainer` in the class definition. + """ + + def __init__(self, **kwargs): + if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): + raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") + super().__init__(**kwargs) + self.is_meta = False + self.ckpt_load_enabled = True + + def initialize_tensors(self, enable_training=False): + super().initialize_tensors(enable_training=enable_training) + self.is_meta = self.qkvw.is_meta + + def apply_tensor_parallelism(self, mp_replace, **kwargs): + if self.is_meta: + if self.qkvb is None: + self.module.attention.attn_qkvb = None + if self.dense_b is None: + self.module.attention.attn_ob = None + else: + super().apply_tensor_parallelism(mp_replace, **kwargs) + + def copy_data_to_new_module(self): + if self.is_meta: + if self.attn_nw is None: + self.module.mlp.attn_nw = self.attn_nw + self.module.mlp.attn_nb = self.attn_nb + else: + super().copy_data_to_new_module() + + def transpose(self): + if not self.is_meta: + super().transpose() + + @abstractmethod + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + """ + Load all the transformer parameter from the checkpoint file (sd). + In addition to the parameter names, we require two + more parameters to help read the data correctly + from the checkpoint and split the qkv heads in the + right order: + 1. `use_load_prefix` (Default: False): this specifies + whether we need to use the name of first abstraction + layer of the model for searching the parameter's name + in a checkpoint file. For more information of how this + is used please see + https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/load_checkpoint.py + 2. `split_qkv` (Default: True): we use this flag when splitting + the qkv parameter into heads. If it is False, it means the heads + of q, k, and v are stored together and needs to split in the + DeepSpeed-Inference API. + """ + raise NotImplementedError("A load_params() function must be defined in the model container \ + when inheriting the MetaTensorContainer feature") diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/split_qkv.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/split_qkv.py new file mode 100644 index 0000000000000000000000000000000000000000..f4c14d4e425a7e9096bae3bf5788ab7b9ca0dcc2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/split_qkv.py @@ -0,0 +1,159 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import abstractmethod +import torch + +from .hybrid_engine import HybridEngineContainer + + +class HybridSplitQKVContainer(HybridEngineContainer): + + def set_attention(self, qkvw, qkvb, dense_w, dense_b): + super().set_attention(qkvw, qkvb, dense_w, dense_b) + self.set_q_k_v() + + @abstractmethod + def set_q_k_v(self): + """ + In `set_q_k_v`, it is necessary to populate the following variables (where appropriate) + for the given model: + self.qw: q weight + self.qb: q bias + self.kw: k weight + self.kb: k bias + self.vw: v weight + self.vb: v bias + """ + raise NotImplementedError("A set_q_k_v() function must be defined in the model container \ + in order to set the unfused q, k, and v tensors.") + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + # Only need to alter + if self.module.attention.attn_qkvw is None: + params = [ + (self.module.attention.attn_qw, self.qw), + (self.module.attention.attn_qb, self.qb), + (self.module.attention.attn_kw, self.kw), + (self.module.attention.attn_kb, self.kb), + (self.module.attention.attn_vw, self.vw), + (self.module.attention.attn_vb, self.vb), + ] + for dst, src in params: + dst = mp_replace.copy( + dst[:self.qw.shape[0] // mp_replace.mp_size], src, int8=reversed_dim, + allocate_tensor=reversed_dim) if src is not None else None + else: + super().attention_qkv_mp(mp_replace) + + def release_qkv(self): + super().release_qkv() + split_qkv_params = [ + (self.module.attention.attn_qw, self.qw), + (self.module.attention.attn_qb, self.qb), + (self.module.attention.attn_kw, self.kw), + (self.module.attention.attn_kb, self.kb), + (self.module.attention.attn_vw, self.vw), + (self.module.attention.attn_vb, self.vb), + ] + + self._release_params(split_qkv_params) + + def reset_qkv(self): + self.qkvw.data[:self.qw.shape[0]] = self.qw.data + self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data + self.qkvw.data[2 * self.qw.shape[0]:] = self.vw.data + + qkv_data = [self.qw.data, self.kw.data, self.vw.data] + + self.qw.data = self.qkvw.data[:self.qw.shape[0]] + self.kw.data = self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.qkvw.data[2 * self.qw.shape[0]:] + + if self.qkvb is not None: + self.qkvb.data[:self.qw.shape[0]] = self.qb.data + self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data + self.qkvb.data[2 * self.qw.shape[0]:] = self.vb.data + + qkv_data.extend([self.qb.data, self.kb.data, self.vb.data]) + + self.qb.data = self.qkvb.data[:self.qw.shape[0]] + self.kb.data = self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vb.data = self.qkvb.data[2 * self.qw.shape[0]:] + + for data in qkv_data: + del data + + def reset_qkv_experimental(self): + """ + WIP - experimental and likely to be changed/improved. + Unused by keeping for now. + """ + if self.module.attention.attn_qkvw is None: + self.module.attention.attn_qkvw = torch.empty(self.qw.shape[0] * 3, + self.qw.shape[0], + dtype=self.qw.dtype, + device=self.qw.device) + self.module.attention.attn_qkvb = torch.empty(self.qw.shape[0] * 3, + dtype=self.qw.dtype, + device=self.qw.device) + self.module.attention.attn_qkvw.data[:self.qw.shape[0]] = self.qw.data + self.module.attention.attn_qkvb.data[:self.qw.shape[0]] = self.qb.data + self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data + self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data + self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] = self.vw.data + self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] = self.vb.data + + qkv_data = [self.qw.data, \ + self.qb.data, \ + self.kw.data, \ + self.kb.data, \ + self.vw.data, \ + self.vb.data] + + self.qw.data = self.module.attention.attn_qkvw.data[:self.qw.shape[0]] + self.qb.data = self.module.attention.attn_qkvb.data[:self.qw.shape[0]] + self.kw.data = self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.kb.data = self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] + self.vb.data = self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] + + for data in qkv_data: + del data + + def set_attn_params_wo_copy(self, Z3_enabled=False): + self.module.attention.attn_ow = self.dense_w + self.module.attention.attn_ob = self.dense_b + if not Z3_enabled: + # In initialize_tensors, we create a fused qkvw with the appropriate shape + # and copy the qw, qb, kw, kb, vw, vb into it + self.module.attention.attn_qkvw = self.qkvw + self.module.attention.attn_qkvb = self.qkvb + + # We reset the data for qw (which is the original model parameter) to point + # to the fused weight matrix we have created here + self.qw.data = self.qkvw[:self.qw.shape[0], :] + self.kw.data = self.qkvw[self.qw.shape[0]:2 * self.qw.shape[0], :] + self.vw.data = self.qkvw[self.qw.shape[0] * 2:, :] + + # Assume if one of the biases is not None, then all of them are not None + if self.qb is not None: + self.qb.data = self.qkvb[:self.qw.shape[0]] + self.kb.data = self.qkvb[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vb.data = self.qkvb[self.qw.shape[0] * 2:] + else: + # In ZeRO-3 this will be managed by ZeRO and handled separately in the + # forward of ds_attention + self.module.attention.attn_qw = self.qw + self.module.attention.attn_qb = self.qb + self.module.attention.attn_kw = self.kw + self.module.attention.attn_kb = self.kb + self.module.attention.attn_vw = self.vw + self.module.attention.attn_vb = self.vb + + def get_attn_params(self): + params = super().get_attn_params() + params.extend([self.qw, self.qb, self.kw, self.kb, self.vw, self.vb]) + return params diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gpt2.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gpt2.py new file mode 100644 index 0000000000000000000000000000000000000000..7a19aac34b447a7d5116ccceeb7687faa61ea174 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gpt2.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +from ..policy import TransformerPolicy + + +class DS_GPT2Container(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFGPT2LayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True): + # HuggingFace GPT2 uses convolutional layer instead of linear layer + super().__init__(inference, linear_layer=False) + self.client_module = client_module + try: + import transformers + HFGPT2LayerPolicy._orig_layer_class = transformers.models.gpt2.modeling_gpt2.GPT2Block + except: + HFGPT2LayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attn.embed_dim, \ + self.client_module.attn.num_heads, \ + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + return self.client_module.attn.c_attn.weight, \ + self.client_module.attn.c_attn.bias, \ + self.client_module.attn.c_proj.weight, \ + self.client_module.attn.c_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.c_fc.weight, \ + self.client_module.mlp.c_fc.bias, \ + self.client_module.mlp.c_proj.weight, \ + self.client_module.mlp.c_proj.bias + + def layernorm(self): + return self.client_module.ln_2.weight, \ + self.client_module.ln_2.bias, \ + self.client_module.ln_1.weight, \ + self.client_module.ln_1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptj.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptj.py new file mode 100644 index 0000000000000000000000000000000000000000..17c0a5027a4c94dcc206b69ad6f3fa6d1d746e17 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptj.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.split_qkv import HybridSplitQKVContainer +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from ..policy import maybe_copy_qkv + +from ..policy import maybe_get_lora + + +class DS_GPTJContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.fc_in, self.policy.client_module.mlp.fc_out, + self.policy.client_module.attn.q_proj, self.policy.client_module.attn.k_proj, + self.policy.client_module.attn.v_proj, self.policy.client_module.attn.out_proj + ] + ] + + def get_lora_matched_pair(self): + fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), + (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attn.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.attn.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.attn.v_proj.weight + self.vb = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attn.q_proj.weight', \ + 'attn.k_proj.weight', \ + 'attn.v_proj.weight', \ + 'attn.out_proj.weight', \ + 'mlp.fc_in.weight', \ + 'mlp.fc_in.bias', \ + 'mlp.fc_out.weight', \ + 'mlp.fc_out.bias', \ + 'ln_1.weight', \ + 'ln_1.bias' + ) + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + for i in range(4, 8): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(8, 10): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i + 2], + prefix + param_names[i]) + + +class HFGPTJLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True): + super().__init__(inference, scale_attention=True) + self.client_module = client_module + try: + import transformers + HFGPTJLayerPolicy._orig_layer_class = transformers.models.gptj.modeling_gptj.GPTJBlock + except: + HFGPTJLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attn.embed_dim, \ + self.client_module.attn.num_attention_heads, \ + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.attn.q_proj.weight + kw = self.client_module.attn.k_proj.weight + vw = self.client_module.attn.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.attn.out_proj.weight, \ + None, + + def mlp(self, enable_training=False): + return self.client_module.mlp.fc_in.weight, \ + self.client_module.mlp.fc_in.bias, \ + self.client_module.mlp.fc_out.weight, \ + self.client_module.mlp.fc_out.bias + + def layernorm(self): + return None, \ + None, \ + self.client_module.ln_1.weight, \ + self.client_module.ln_1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneo.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneo.py new file mode 100644 index 0000000000000000000000000000000000000000..fca673b375e18a9b74741f86a9692bfd386f6a51 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneo.py @@ -0,0 +1,145 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.split_qkv import HybridSplitQKVContainer +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from ..policy import maybe_copy_qkv + +from ..policy import maybe_get_lora + + +class DS_GPTNEOContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.c_fc, self.policy.client_module.mlp.c_proj, + self.policy.client_module.attn.attention.q_proj, self.policy.client_module.attn.attention.k_proj, + self.policy.client_module.attn.attention.v_proj, self.policy.client_module.attn.attention.out_proj + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attn.attention.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.attn.attention.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.attn.attention.v_proj.weight + self.vb = None + + def get_lora_matched_pair(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), + (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attn.attention.q_proj.weight', \ + 'attn.attention.k_proj.weight', \ + 'attn.attention.v_proj.weight', \ + 'attn.attention.out_proj.weight', \ + 'attn.attention.out_proj.bias', \ + 'mlp.c_fc.weight', \ + 'mlp.c_fc.bias', \ + 'mlp.c_proj.weight', \ + 'mlp.c_proj.bias', \ + 'ln_2.weight', \ + 'ln_2.bias', \ + 'ln_1.weight', \ + 'ln_1.bias' + ) + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 5): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + for i in range(5, 11): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + for i in range(11, 13): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + + +class HFGPTNEOLayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__(inference, scale_attention=False) + self.client_module = client_module + try: + import transformers + HFGPTNEOLayerPolicy._orig_layer_class = transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock + except: + HFGPTNEOLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attn.attention.embed_dim, \ + self.client_module.attn.attention.num_heads, \ + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def get_q_k_v(self): + return self.client_module.attn.attention.q_proj.weight, \ + None, \ + self.client_module.attn.attention.k_proj.weight, \ + None, \ + self.client_module.attn.attention.v_proj.weight, \ + None + + def attention(self, enable_training=False): + qw = self.client_module.attn.attention.q_proj.weight + kw = self.client_module.attn.attention.k_proj.weight + vw = self.client_module.attn.attention.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.attn.attention.out_proj.weight, \ + self.client_module.attn.attention.out_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.c_fc.weight, \ + self.client_module.mlp.c_fc.bias, \ + self.client_module.mlp.c_proj.weight, \ + self.client_module.mlp.c_proj.bias + + def layernorm(self): + return self.client_module.ln_2.weight, \ + self.client_module.ln_2.bias, \ + self.client_module.ln_1.weight, \ + self.client_module.ln_1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneox.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneox.py new file mode 100644 index 0000000000000000000000000000000000000000..3978b166b329b341232da61811ecaf5683c5578d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneox.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.hybrid_megatron import HybridMegatronContainer +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from packaging import version as pkg_version + +from ..policy import maybe_get_lora + + +class DS_GPTNEOXContainer(MetaTensorContainer, HybridMegatronContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + + if self.megatron_v2: + self.module.config.rotate_half = True + self.module.config.rotate_every_two = False + + return self.module + + def get_lora_matched_pair(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)] + return ret + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + if GPTNEOXLayerPolicy.version == 0: + attention = self.policy.client_module.attention + else: + attention = self.policy.client_module.self_attention + + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, + attention.query_key_value, attention.dense + ] + ] + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attention.query_key_value.weight', \ + 'attention.query_key_value.bias', \ + 'attention.dense.weight', \ + 'attention.dense.bias', \ + 'mlp.dense_h_to_4h.weight', \ + 'mlp.dense_h_to_4h.bias', \ + 'mlp.dense_4h_to_h.weight', \ + 'mlp.dense_4h_to_h.bias', \ + 'post_attention_layernorm.weight', \ + 'post_attention_layernorm.bias', \ + 'input_layernorm.weight', \ + 'input_layernorm.bias' + ) + for i in range(0, 2): + maybe_copy(module.attention, + sd, + weight_quantizer, + mp_replace, + transformer_param_names[i], + prefix + param_names[i], + qkv=True, + megatron_v2=self.policy.is_megatron_v2, + split_qkv=self.policy.split_qkv, + heads=self.policy.client_module.attention.num_attention_heads) + for i in range(2, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(4, 10): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(10, 12): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i]) + + +class GPTNEOXLayerPolicy(TransformerPolicy): + _orig_layer_class = None + version = 0 + + def __init__(self, client_module, inference=True, megatron_v2=True, split_qkv=False): + super().__init__(inference, megatron_v2=megatron_v2, split_qkv=split_qkv) + self.client_module = client_module + if GPTNEOXLayerPolicy._orig_layer_class is None: + if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): + GPTNEOXLayerPolicy._orig_layer_class = None + else: + try: + from transformers import GPTNeoXLayer + GPTNEOXLayerPolicy._orig_layer_class = GPTNeoXLayer + except ImportError: + GPTNEOXLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + if GPTNEOXLayerPolicy.version == 0: + attention = self.client_module.attention + else: + attention = self.client_module.self_attention + + return self.client_module.attention.hidden_size, \ + self.client_module.attention.num_attention_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + if GPTNEOXLayerPolicy.version == 0: + attention = self.client_module.attention + else: + attention = self.client_module.self_attention + + return attention.query_key_value.weight, \ + attention.query_key_value.bias, \ + attention.dense.weight, \ + attention.dense.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.dense_4h_to_h.bias + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + self.client_module.post_attention_layernorm.bias, \ + self.client_module.input_layernorm.weight, \ + self.client_module.input_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/internlm.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/internlm.py new file mode 100644 index 0000000000000000000000000000000000000000..31255d4b3ca5a0199ee572219b6d75fd70d478c3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/internlm.py @@ -0,0 +1,181 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import importlib + +import torch +from torch.nn.parameter import Parameter + +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +from deepspeed.utils.types import ActivationFuncType, NormType + +from ..policy import (TransformerPolicy, maybe_copy, maybe_copy_geglu, maybe_copy_qkv, maybe_get_lora, + transformer_param_names) +from .base import * +from .features import HybridGatedMLPContainer, HybridSplitQKVContainer + + +class DS_InternLMContainer(HybridGatedMLPContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = True + _config.rotate_every_two = False + _config.rotary_dim = self.hidden_size // self.num_attention_heads + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, + self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, + self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, + self.policy.client_module.self_attn.o_proj.weight + ] + ] + + def get_lora_matched_pair(self): + up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), + (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = self.policy.client_module.self_attn.q_proj.bias + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = self.policy.client_module.self_attn.k_proj.bias + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = self.policy.client_module.self_attn.v_proj.bias + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.mlp.up_proj.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.o_proj.weight', \ + 'mlp.up_proj.weight', \ + 'mlp.gate_proj.weight', \ + 'mlp.down_proj.weight', \ + 'input_layernorm.weight', \ + 'post_attention_layernorm.weight' + 'self_attn.q_proj.bias', \ + 'self_attn.k_proj.bias', \ + 'self_attn.v_proj.bias', \ + 'self_attn.o_proj.bias', \ + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvb', [prefix + param_names[9], prefix + param_names[10], prefix + param_names[11]], + split_qkv=self.policy.split_qkv) + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[2], + prefix + param_names[3]) + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[3], + prefix + param_names[12]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + +class InternLMLayerPolicy(TransformerPolicy): + _orig_layer_class = [] + _orig_layer_class_inited = False + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + + self._init_orig_layer_class_once() + + def _init_orig_layer_class_once(self): + if InternLMLayerPolicy._orig_layer_class_inited: + return + + for sub_pkg in ['', '.internlm-7b', '.internlm-chat-7b']: + try: + from transformers.utils import TRANSFORMERS_DYNAMIC_MODULE_NAME + module = importlib.import_module(f"{TRANSFORMERS_DYNAMIC_MODULE_NAME}{sub_pkg}.modeling_internlm") + if module.InternLMDecoderLayer not in InternLMLayerPolicy._orig_layer_class: + InternLMLayerPolicy._orig_layer_class.append(module.InternLMDecoderLayer) + except ImportError: + continue + + InternLMLayerPolicy._orig_layer_class_inited = True + + def get_hidden_heads(self): + return self.client_module.self_attn.q_proj.weight.shape[1], \ + self.client_module.self_attn.num_heads, \ + self.client_module.input_layernorm.variance_epsilon, \ + self.client_module.mlp.gate_proj.weight.shape[0] + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + kw = self.client_module.self_attn.k_proj.weight + vw = self.client_module.self_attn.v_proj.weight + qb = self.client_module.self_attn.q_proj.bias + kb = self.client_module.self_attn.k_proj.bias + vb = self.client_module.self_attn.v_proj.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.self_attn.o_proj.weight, \ + self.client_module.self_attn.o_proj.bias + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.mlp.up_proj.weight + mlp1_gate = self.client_module.mlp.gate_proj.weight + mlp2 = self.client_module.mlp.down_proj.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + None, \ + self.client_module.input_layernorm.weight, \ + None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..7af333dc1ee4771f52bcc837cf01c847486bdfd7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py @@ -0,0 +1,164 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer +from deepspeed.utils.types import ActivationFuncType, NormType +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter + +from ..policy import ( + TransformerPolicy, + transformer_param_names, + maybe_copy, + maybe_copy_qkv, + maybe_copy_geglu, + maybe_get_lora, +) + + +class DS_LLAMAContainer(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, + BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = True + _config.rotate_every_two = False + _config.rotary_dim = self.hidden_size // self.num_attention_heads + _config.rope_theta = self.policy.client_module.self_attn.rope_theta + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, + self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, + self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, + self.policy.client_module.self_attn.o_proj.weight + ] + ] + + def get_lora_matched_pair(self): + up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), + (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = None + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.mlp.up_proj.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.o_proj.weight', \ + 'mlp.up_proj.weight', \ + 'mlp.gate_proj.weight', \ + 'mlp.down_proj.weight', \ + 'post_attention_layernorm.weight', \ + 'input_layernorm.weight', + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + # This line is necessary for proper output when kernels + meta tensors are used in Llama models + # TODO: Investigate root-cause and fix meta tensor loading + module.mlp.output_b = None + + +class LLAMALayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + try: + import transformers + LLAMALayerPolicy._orig_layer_class = transformers.models.llama.modeling_llama.LlamaDecoderLayer # type: ignore + except: + LLAMALayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + hidden_heads = ( + self.client_module.self_attn.q_proj.in_features, + self.client_module.self_attn.num_heads, + self.client_module.input_layernorm.variance_epsilon, + self.client_module.mlp.gate_proj.out_features, + ) + return hidden_heads + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + kw = self.client_module.self_attn.k_proj.weight + vw = self.client_module.self_attn.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.self_attn.o_proj.weight, \ + None + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.mlp.up_proj.weight + mlp1_gate = self.client_module.mlp.gate_proj.weight + mlp2 = self.client_module.mlp.down_proj.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + None, \ + self.client_module.input_layernorm.weight, \ + None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama2.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama2.py new file mode 100644 index 0000000000000000000000000000000000000000..b531890ab859dfbff51a401819846299d8e77c08 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama2.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer +from deepspeed.utils.types import ActivationFuncType, NormType +from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference +import torch +from torch.nn.parameter import Parameter + +from ..policy import ( + TransformerPolicy, + transformer_param_names, + maybe_copy, + maybe_copy_qkv, + maybe_copy_geglu, + maybe_get_lora, +) + + +class DS_LLAMA2Container(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, + BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = False + _config.rotate_every_two = True + _config.rotary_dim = self.hidden_size // self.num_attention_heads + _config.num_kv = self.policy.client_module.attention.n_kv_heads + self.module = DeepSpeedLlama2Inference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.feed_forward.w3.weight, self.policy.client_module.feed_forward.w1.weight, + self.policy.client_module.feed_forward.w2.weight, self.policy.client_module.attention.wq.weight, + self.policy.client_module.attention.wk.weight, self.policy.client_module.attention.wv.weight, + self.policy.client_module.attention.wo.weight + ] + ] + + def get_lora_matched_pair(self): + up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), + (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attention.wq.weight + self.qb = None + self.kw = self.policy.client_module.attention.wk.weight + self.kb = None + self.vw = self.policy.client_module.attention.wv.weight + self.vb = None + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.feed_forward.w2.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.feed_forward.w1.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attention.wq.weight', \ + 'attention.wk.weight', \ + 'attention.wv.weight', \ + 'attention.wo.weight', \ + 'feed_forward.w3.weight', \ + 'feed_forward.w1.weight', \ + 'feed_forward.w2.weight', \ + 'ffn_norm.weight', \ + 'attention_norm.weight' + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + +class LLAMA2LayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + try: + import llama + LLAMA2LayerPolicy._orig_layer_class = llama.model.TransformerBlock # type: ignore + except: + LLAMA2LayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attention.wq.weight.shape[1], \ + self.client_module.n_heads, \ + self.client_module.ffn_norm.eps, \ + (self.client_module.feed_forward.w1.weight.shape[0] * \ + deepspeed.comm.get_world_size() if deepspeed.comm.is_initialized() else 1) # this is a hack to inject when model is already partitioned! + + def attention(self, enable_training=False): + qw = self.client_module.attention.wq.weight + kw = self.client_module.attention.wk.weight + vw = self.client_module.attention.wv.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.attention.wo.weight, \ + None + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.feed_forward.w3.weight + mlp1_gate = self.client_module.feed_forward.w1.weight + mlp2 = self.client_module.feed_forward.w2.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.ffn_norm.weight, \ + None, \ + self.client_module.attention_norm.weight, \ + None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..2851dd246d99c94235cc6c0abce05f0c73375138 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt.py @@ -0,0 +1,117 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.megatron import MegatronContainer +from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference +import torch +from ..policy import TransformerPolicy +from packaging import version as pkg_version + + +class DS_MegatronGPTContainer(MegatronContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + + if self.megatron_v2: + self.module.config.rotate_half = True + self.module.config.rotate_every_two = False + + return self.module + + +# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp +# TODO: Generalize MoE overall goal, expand beyond Megatron +class MegatronLayerPolicy(TransformerPolicy): + _orig_layer_class = None + version = 0 + moe_type = 'standard' + megatron_v2 = True + use_mup = False + + def __init__(self, client_module, inference=True): + super().__init__(inference, megatron_v2=MegatronLayerPolicy.megatron_v2, use_mup=MegatronLayerPolicy.use_mup) + self.client_module = client_module + # we use megatron version to differentiate between the old and new + # megatron-lm source code + if MegatronLayerPolicy._orig_layer_class is None: + if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): + MegatronLayerPolicy._orig_layer_class = None + else: + try: + from megatron.model.transformer import ParallelTransformerLayer + MegatronLayerPolicy._orig_layer_class = ParallelTransformerLayer + MegatronLayerPolicy.version = 1 + except ImportError: + MegatronLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + if MegatronLayerPolicy.version == 0: + return self.client_module.attention.query_key_value.weight.shape[1], \ + self.client_module.attention.num_attention_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + else: + return self.client_module.self_attention.query_key_value.weight.shape[1], \ + self.client_module.self_attention.num_attention_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + if self.inference: + if MegatronLayerPolicy.version == 0: + attention = self.client_module.attention + else: + attention = self.client_module.self_attention + + return attention.query_key_value.weight, \ + attention.query_key_value.bias, \ + attention.dense.weight, \ + attention.dense.bias + + def mlp(self, moe_type='standard', enable_training=False): + from deepspeed.moe.utils import has_moe_layers + moe, _ = has_moe_layers(self.client_module) + + if moe: + moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \ + self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts + num_experts = len(moe_experts) + if moe_type == 'standard': + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)] + else: + + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \ + self.client_module.mlp.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.mlp.dense_4h_to_h.bias, \ + self.client_module.mlp.coefficient.weight + + else: + return self.client_module.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.dense_4h_to_h.bias + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + self.client_module.post_attention_layernorm.bias, \ + self.client_module.input_layernorm.weight, \ + self.client_module.input_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..c4063be05b6c5f232470685a292be51667c027a8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .base_moe import * +from .features.megatron import MegatronContainer +from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference +import torch +from .megatron_gpt import MegatronLayerPolicy +from packaging import version as pkg_version + + +class DS_MegatronGPTMoEContainer(MegatronContainer, BaseTransformerMoEContainer): + + def __init__(self, policy, config, model_config, layer_id): + super().__init__(policy, config, model_config, layer_id) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + + if self.megatron_v2: + self.module.config.rotate_half = True + self.module.config.rotate_every_two = False + + return self.module + + +# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp +# TODO: Generalize MoE overall goal, expand beyond Megatron +class MegatronMoELayerPolicy(MegatronLayerPolicy): + _orig_layer_class = None + version = 0 + moe_type = 'standard' + num_experts = 1 + + def __init__(self, client_module, inference=True): + super().__init__(inference) + self.client_module = client_module + # we use megatron version to differentiate between the old and new + # megatron-lm source code + if MegatronMoELayerPolicy._orig_layer_class is None: + if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): + MegatronMoELayerPolicy._orig_layer_class = None + else: + try: + from megatron.model.transformer import ParallelTransformerLayer + MegatronMoELayerPolicy._orig_layer_class = ParallelTransformerLayer + except ImportError: + MegatronMoELayerPolicy._orig_layer_class = None + + def get_num_experts(self): + return self.num_experts + + def mlp(self, moe_type='standard', enable_training=False): + # for now, all of this is tightly coupled to megatron-deepspeed moe implementation + # todo: think and refactor this to be more general + + #from deepspeed.moe.utils import has_moe_layers + #moe, _ = has_moe_layers(self.client_module) + + moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \ + self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts + num_experts = len(moe_experts) + self.num_experts = num_experts + + if moe_type == 'standard': + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)] + else: + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \ + self.client_module.mlp.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.mlp.dense_4h_to_h.bias, \ + self.client_module.mlp.coefficient.weight diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/opt.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/opt.py new file mode 100644 index 0000000000000000000000000000000000000000..eba619c79d6ebf5325c12f42c921b77dcb6e5ad7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/opt.py @@ -0,0 +1,160 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import MetaTensorContainer, HybridSplitQKVContainer +from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from ..policy import maybe_copy_qkv +from ..policy import maybe_get_lora +from deepspeed.utils.types import ActivationFuncType + + +class DS_OPTContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedOPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.fc1, + self.policy.client_module.fc2, + self.policy.client_module.self_attn.q_proj, + self.policy.client_module.self_attn.k_proj, + self.policy.client_module.self_attn.v_proj, + self.policy.client_module.self_attn.out_proj, + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = self.policy.client_module.self_attn.q_proj.bias + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = self.policy.client_module.self_attn.k_proj.bias + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = self.policy.client_module.self_attn.v_proj.bias + + def get_lora_matched_pair(self): + fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), + (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.q_proj.bias', \ + 'self_attn.k_proj.bias', \ + 'self_attn.v_proj.bias', \ + 'self_attn.out_proj.weight', \ + 'self_attn.out_proj.bias', \ + 'fc1.weight', \ + 'fc1.bias', \ + 'fc2.weight', \ + 'fc2.bias', \ + 'final_layer_norm.weight', \ + 'final_layer_norm.bias', \ + 'self_attn_layer_norm.weight', \ + 'self_attn_layer_norm.bias' + ) + + for i in range(0, 6, 3): + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + transformer_param_names[i // 3], + [prefix + param_names[i], prefix + param_names[i + 1], prefix + param_names[i + 2]], + split_qkv=self.policy.split_qkv) + for i in range(6, 8): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], + prefix + param_names[i]) + for i in range(8, 14): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], + prefix + param_names[i]) + for i in range(14, 16): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], + prefix + param_names[i]) + + +class HFOPTLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True, use_load_prefix=True): + super().__init__(inference, linear_layer=True, pre_attn_norm=True, use_load_prefix=use_load_prefix) + self.client_module = client_module + try: + import transformers + HFOPTLayerPolicy._orig_layer_class = transformers.models.opt.modeling_opt.OPTDecoderLayer + except: + HFOPTLayerPolicy._orig_layer_class = None + + if hasattr(TransformerPolicy, "hf_model_config") and hasattr(TransformerPolicy.hf_model_config, + "activation_function"): + if TransformerPolicy.hf_model_config.activation_function == "relu": + self.mlp_act_func_type = ActivationFuncType.ReLU + elif TransformerPolicy.hf_model_config.activation_function in ["gelu", "gelu_new"]: + self.mlp_act_func_type = ActivationFuncType.GELU + else: + raise ValueError("Unsupported activation function: {}".format( + TransformerPolicy.hf_model_config.activation_function)) + else: + self.mlp_act_func_type = ActivationFuncType.ReLU # default + + def get_hidden_heads(self): + return self.client_module.self_attn.embed_dim, \ + self.client_module.self_attn.num_heads, \ + self.client_module.self_attn_layer_norm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + qb = self.client_module.self_attn.q_proj.bias + + kw = self.client_module.self_attn.k_proj.weight + kb = self.client_module.self_attn.k_proj.bias + + vw = self.client_module.self_attn.v_proj.weight + vb = self.client_module.self_attn.v_proj.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + return qkvw, \ + qkvb, \ + self.client_module.self_attn.out_proj.weight, \ + self.client_module.self_attn.out_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.fc1.weight, \ + self.client_module.fc1.bias, \ + self.client_module.fc2.weight, \ + self.client_module.fc2.bias + + def layernorm(self): + return self.client_module.final_layer_norm.weight, \ + self.client_module.final_layer_norm.bias, \ + self.client_module.self_attn_layer_norm.weight, \ + self.client_module.self_attn_layer_norm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/unet.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..48179265553150824855140750da3fc6e41503ef --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/unet.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from torch.nn.parameter import Parameter + +from ..policy import DSPolicy +from ...model_implementations.diffusers.unet import DSUNet + + +class UNetPolicy(DSPolicy): + + def __init__(self): + super().__init__() + try: + import diffusers + self._orig_layer_class = diffusers.models.unet_2d_condition.UNet2DConditionModel + except AttributeError: + self._orig_layer_class = diffusers.models.unets.unet_2d_condition.UNet2DConditionModel + except ImportError: + self._orig_layer_class = None + + def match(self, module): + return isinstance(module, self._orig_layer_class) + + def match_replaced(self, module): + return isinstance(module, DSUNet) + + def apply(self, module, enable_cuda_graph=True): + # TODO(cmikeh2): Enable cuda graph should be an inference configuration + return DSUNet(module, enable_cuda_graph=enable_cuda_graph) + + def attention(self, client_module): + qw = client_module.to_q.weight + kw = client_module.to_k.weight + vw = client_module.to_v.weight + + if qw.shape[1] == kw.shape[1]: + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=False) + + return qkvw, \ + client_module.to_out[0].weight, \ + client_module.to_out[0].bias, \ + qw.shape[-1], \ + client_module.heads + else: + #return None + #kvw = Parameter(torch.cat((kw, vw), dim=0), requires_grad=False) + return qw, \ + kw, vw, \ + client_module.to_out[0].weight, \ + client_module.to_out[0].bias, \ + qw.shape[-1], \ + client_module.heads diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/vae.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..d26d0ef77ca9d10e27bad47859dba589190af7b8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/vae.py @@ -0,0 +1,42 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..policy import DSPolicy +from ...model_implementations.diffusers.vae import DSVAE + + +class VAEPolicy(DSPolicy): + + def __init__(self): + super().__init__() + try: + import diffusers + if hasattr(diffusers.models, "autoencoders"): + # Diffusers >= 0.25.0 + # Changes location to 'autoencoders' directory + self._orig_layer_class = diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL + elif hasattr(diffusers.models.vae, "AutoencoderKL"): + # Diffusers < 0.12.0 + self._orig_layer_class = diffusers.models.vae.AutoencoderKL + else: + # Diffusers >= 0.12.0 & < 0.25.0 + # Changes location of AutoencoderKL + self._orig_layer_class = diffusers.models.autoencoder_kl.AutoencoderKL + except ImportError: + self._orig_layer_class = None + + def match(self, module): + return isinstance(module, self._orig_layer_class) + + def match_replaced(self, module): + return isinstance(module, DSVAE) + + def apply(self, module, enable_cuda_graph=True): + # TODO(cmikeh2): Enable cuda graph should be an inference configuration + return DSVAE(module, enable_cuda_graph=enable_cuda_graph) + + # NOTE (lekurile): Should we have a diffusers policy class? + def attention(self, client_module): + pass diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/fusedqkv_utils.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/fusedqkv_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0609c6001dd2f716b3c1953500369dc9e67a4987 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/fusedqkv_utils.py @@ -0,0 +1,231 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +import torch +from deepspeed.utils.logging import warning_once +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, get_num_kv_heads, get_n_embd, get_num_attention_heads + + +def split_by_qkvlist_and_refuse(qkv_list, split_size, split_dim=0, cat_dim=0): + qkv_split_list = [torch.split(mat, split_size, dim=split_dim) for mat in qkv_list] + tp_fusedqkv_list = [ + torch.cat([qkv_s[i] for qkv_s in qkv_split_list], dim=cat_dim) for i in range(len(qkv_split_list[0])) + ] + return tp_fusedqkv_list + + +def require_tp_fused_qkvw(name, mp_size): + fused_qkvw_name_list = ['qkv_proj', 'query_key_value', 'attn.Wqkv', 'self_attn.W_pack', 'c_attn'] + + if mp_size == 1: + return False + for fused_name in fused_qkvw_name_list: + if fused_name in name: + return True + return False + + +def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index): + + module_str = str(module).strip() + if src is None: + return + fused_type_dict = { + 'CodeGenBlock': 'codegentype', + 'BloomBlock': 'bloomtype', + 'GLMBlock': 'glmtype', + "MPTBlock": 'glmtype', + "MptBlock": 'glmtype', + "BaichuanLayer": 'glmtype', + "QWenBlock": 'qwentype', + "FalconDecoderLayer": 'bloomtype', + "GPTBigCodeBlock": 'bigcodetype', + "DecoderLayer": 'glmtype', + "Phi3DecoderLayer": "phi3type" + } + + def _codegen_type_transpose(input, mp_size, codegen_mp_num=4): + # codegen_mp_num defined in https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py + assert get_num_kv_heads() % ( + mp_size * codegen_mp_num) == 0, "codgen autoTP requires num_kv_heads % (mp_size*codegen_mp_num) == 0" + #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias) + + shape = input.shape + dst_shape = get_shard_size(shape[0], mp_size) + num_mp_blocks = input.reshape(codegen_mp_num, shape[0] // codegen_mp_num, shape[1]) + + #num_mp_blocks : [codegen_mp_num, 3*hidden_dim/codegen_mp_num, :] + src_split = list(torch.split(num_mp_blocks, num_mp_blocks.shape[1] // 3, dim=1)) + src_split = [x.reshape(codegen_mp_num * mp_size, -1, shape[1]) for x in src_split] + + split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size(shape[0] // 3, mp_size), 0, 1) + tp_fuseqkv_weight = torch.cat(split_fusedqkv, dim=0).reshape(shape[0], -1) + + return tp_fuseqkv_weight[gpu_index * dst_shape:(gpu_index + 1) * dst_shape] + + def _glm_type_transpose(input, mp_size): + #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias) + + # For chatglm2 & chatglm3(kv_heads=2), need to special handle. + if get_num_kv_heads() == 2: + shape = input.shape + hidden_dim = get_n_embd() + kv_dim = (shape[0] - hidden_dim) // get_num_kv_heads() + q = input[:hidden_dim] + k = input[hidden_dim:hidden_dim + kv_dim] + v = input[hidden_dim + kv_dim:] + q_split = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0) + k_split = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0) + v_split = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0) + return torch.cat((q_split[gpu_index], k_split[gpu_index], v_split[gpu_index]), dim=0) + else: + shape = input.shape + src_split = torch.split(input, shape[0] // 3, dim=0) + + split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size)) + return split_fusedqkv[gpu_index] + + def _bloom_type_transpose(input, mp_size): + shape = input.shape + + split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size), dim=0) + return split_fusedqkv[gpu_index] + + def _qwen_type_transpose(input, mp_size, module): + if not hasattr(module, "_ds_fusedqkv_entered"): + # Adjust splitting absolute value variables + setattr(module, "_ds_fusedqkv_entered", True) + module.attn.split_size = get_shard_size(module.attn.split_size, mp_size) + return _glm_type_transpose(input, mp_size) + + def _bigcode_type_transpose(input, mp_size): + n_embd = get_n_embd() + q = input[:n_embd] + kv = input[n_embd:] + shape = q.shape + split_q = q.split(get_shard_size_list(shape[0], mp_size), dim=0) + return torch.cat((split_q[gpu_index], kv), dim=0) + + def _phi3_type_transpose(input, mp_size): + num_kv_heads = get_num_kv_heads() + num_heads = get_num_attention_heads() + hidden_size = input.shape[1] + head_dim = hidden_size // num_heads + q_pos = input.shape[0] - 2 * num_kv_heads * head_dim + q = input[:q_pos] + k = input[q_pos:q_pos + num_kv_heads * head_dim] + v = input[q_pos + num_kv_heads * head_dim:] + split_q = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0) + split_k = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0) + split_v = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0) + return torch.cat((split_q[gpu_index], split_k[gpu_index], split_v[gpu_index]), dim=0) + + def _transpose_fused_qkvw(src, mp_size, fused_qkv_type=None, module=None): + + # suppose num_heads=n, q(n)_w means the n-th q head linear weight, the weight format are as following + # bloomtype: [q(1)_w,k(1)_w,v(1)_w,q(2)_w,k(2)_w,v(2)_w,...,q(n)_w,k(n)_w,v(n)_w] + # glmtype: [q(1)_w, q(2)_w,...,q(n)_w,k(1)_w,k(2)_w,...,k(n)_w,v(1)_w,v(2)_w,...,v(n)_w] + # codegentype: [q(1)_w,q(2)_w,...,q(n/t)_w,k(1)_w,k(2)_w,...,k(n/t)_w,v(1)_2,v(2)_w,...v(n/t)_w,q(n/t+1)_w,...], where t is a const defined in model file. + + if fused_qkv_type == 'bloomtype': + return _bloom_type_transpose(src, mp_size) + elif fused_qkv_type == 'codegentype': + return _codegen_type_transpose(src, mp_size) + elif fused_qkv_type == 'glmtype': + return _glm_type_transpose(src, mp_size) + elif fused_qkv_type == 'qwentype': + return _qwen_type_transpose(src, mp_size, module) + elif fused_qkv_type == 'bigcodetype': + return _bigcode_type_transpose(src, mp_size) + elif fused_qkv_type == 'phi3type': + return _phi3_type_transpose(src, mp_size) + + raise ValueError("unknown fused_qkv_type") + + module_name_matches = [k for k in fused_type_dict.keys() if k in module_str] + if module_name_matches: + # There can be overlap with matches (e.g., "DecoderLayer" and "FalconDecoderLayer"). + # We take the longest matching module_name + module_name = max(module_name_matches, key=len) + fused_type = fused_type_dict[module_name] + return _transpose_fused_qkvw(src, mp_size, fused_type, module) + warning_once(f"Unrecognized fusedkqv weight type, default to using bloom type," + f"please check in prepare_tp_fused_qkvw() to avoid potential calculation errors") + return _bloom_type_transpose(src, mp_size) + + +# For share qk type: +# q = [q1,...,q_{n/4}, q_{n/2+1},...,q_{3n/4}, k1,...,k_{n/4}, k_{n/2+1},...,k_{3n/4}] +# k = [q_{n/4+1},...,q_{n/2}, q_{3n/4+1},...,qn, k_{n/4+1},...,k_{n/2}, k{3n/4+1},...,kn] +# Avoid modifying the modeling code. We adjust the value and oproj weight to fit this qk type. +def shard_value_with_share_qk( + weight, + bias, + rank, + world_size, + shard_value=True # True -> shard_value; False -> shard_oproj +): + if shard_value: + total_size = weight.shape[0] + weight_cat_dim = 0 + else: + total_size = weight.shape[1] + weight_cat_dim = 1 + num_heads = get_num_kv_heads() + head_dim = total_size // num_heads + assert (num_heads % world_size == 0) + if world_size > num_heads // 2: + RuntimeError(f"world_size {world_size} is larger than half of num_heads {num_heads}") + head_per_rank = num_heads // world_size + q_head_start = rank * head_per_rank + # mapping q_head to v_head + v_head_ids = [] + i = 0 + # mapping neighbor q_head to v_head + while i < head_per_rank: + v_head_ids.append(q_head_start // 2) + q_head_start += 2 + i = i + 2 + + # mapping neighbor k_head to v_head + v_head_ids.extend([i + num_heads // 2 for i in v_head_ids]) + sharded_weight = [] + sharded_bias = [] + for head_id in v_head_ids: + if shard_value: + sharded_weight.append(weight[head_id * head_dim:(head_id + 1) * head_dim]) + if bias is not None: + sharded_bias.append(bias.data[head_id * head_dim:(head_id + 1) * head_dim]) + else: + sharded_weight.append(weight[:, head_id * head_dim:(head_id + 1) * head_dim]) + sharded_weight = torch.cat(sharded_weight, dim=weight_cat_dim) + if bias is not None: + if shard_value: + sharded_bias = torch.cat(sharded_bias, dim=0) + else: + bias = bias / float(world_size) + return torch.nn.Parameter(sharded_weight), torch.nn.Parameter(sharded_bias) + else: + return torch.nn.Parameter(sharded_weight), None + + +# For phi3 with chunk mlp, adjust the weight order. +def shard_chunk_mlp( + weight, + bias, + rank, + world_size, +): + weight_gate, weight_states = weight.chunk(2, dim=0) + total_size = weight_gate.shape[0] + split_weight_gate = weight_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + split_weight_states = weight_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + shard_weight = torch.cat((split_weight_gate[rank], split_weight_states[rank]), dim=0) + if bias is not None: + bias_gate, bias_states = bias.chunk(2, dim=0) + split_bias_gate = bias_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + split_bias_states = bias_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + return shard_weight, torch.cat((split_bias_gate[rank], split_bias_states[rank]), dim=0) + + return shard_weight, None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/inject.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/inject.py new file mode 100644 index 0000000000000000000000000000000000000000..401da1bd6ef76b268f65e48b77c30a8ed65bfc64 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/inject.py @@ -0,0 +1,112 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy +import torch +from deepspeed.ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig + + +def module_inject(layer_obj, model, config, micro_batch_size, max_seq_length, seed, preln, fp16=True): + for name, child in model.named_children(): + if isinstance(child, layer_obj): + print('REPLACING BertLayer') + + cuda_config = DeepSpeedTransformerConfig(batch_size=micro_batch_size, + max_seq_length=max_seq_length, + hidden_size=config.hidden_size, + heads=config.num_attention_heads, + attn_dropout_ratio=config.attention_probs_dropout_prob, + hidden_dropout_ratio=config.hidden_dropout_prob, + num_hidden_layers=config.num_hidden_layers, + initializer_range=config.initializer_range, + seed=seed, + fp16=fp16, + pre_layer_norm=preln) + + new_module = DeepSpeedTransformerLayer(cuda_config) + + # copy relevant state from child -> new module + qw = child.attention.self.query.weight + qb = child.attention.self.query.bias + kw = child.attention.self.key.weight + kb = child.attention.self.key.bias + vw = child.attention.self.value.weight + vb = child.attention.self.value.bias + + qkvw = torch.cat((qw, kw, vw), 0) + qkvb = torch.cat((qb, kb, vb), 0) + + new_module.attn_qkvw.data = qkvw + new_module.attn_qkvb.data = qkvb + new_module.attn_ow.data = child.attention.output.dense.weight + new_module.attn_ob.data = child.attention.output.dense.bias + if preln: + attention_layerNorm = child.PostAttentionLayerNorm + else: + attention_layerNorm = child.attention.output.LayerNorm + new_module.attn_nw.data = attention_layerNorm.weight + new_module.attn_nb.data = attention_layerNorm.bias + if preln: + intermediate_FF = child.intermediate.dense_act + else: + intermediate_FF = child.intermediate.dense + new_module.inter_w.data = intermediate_FF.weight + new_module.inter_b.data = intermediate_FF.bias + new_module.output_w.data = child.output.dense.weight + new_module.output_b.data = child.output.dense.bias + if preln: + transformer_LayerNorm = child.PreAttentionLayerNorm + else: + transformer_LayerNorm = child.output.LayerNorm + new_module.norm_w.data = transformer_LayerNorm.weight + new_module.norm_b.data = transformer_LayerNorm.bias + + setattr(model, name, copy.deepcopy(new_module)) + + else: + module_inject(layer_obj, child, config, micro_batch_size, max_seq_length, seed, preln, fp16) + + return model + + +def test_hi(): + from turing.nvidia_modelingpreln import BertConfig as BertConfigPreLN + from turing.nvidia_modelingpreln import BertForQuestionAnswering as BertForQuestionAnsweringPreLN + from turing.nvidia_modelingpreln import BertLayer + bert_model_config = { + "vocab_size_or_config_json_file": 119547, + "hidden_size": 1024, + "num_hidden_layers": 1, + "num_attention_heads": 16, + "intermediate_size": 4096, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.1, + "attention_probs_dropout_prob": 0.1, + "hidden_dropout_prob": 0.1, + "attention_probs_dropout_prob": 0.1, + "max_position_embeddings": 512, + "type_vocab_size": 2, + "initializer_range": 0.02 + } + bert_config = BertConfigPreLN(**bert_model_config) + base_model = BertForQuestionAnsweringPreLN(bert_config, args=None) + + #base_model = LinearStack() + + test_model = copy.deepcopy(base_model) + test_model = module_inject(BertLayer, test_model, bert_config, 4, 384, 1234) + + print('BASE', base_model) + print('TEST', test_model) + + #base_model.eval() + #test_model.eval() + + #test_input = torch.rand(1, base_model.input_dim) + + #base_output = base_model(test_input) + #test_output = test_model(test_input) + # + #assert torch.allclose(base_output, test_output, atol=3e-8) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/layers.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..a9faac71361a4aec8c76e6976121ed794b554115 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/layers.py @@ -0,0 +1,821 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +from torch import nn +from torch.nn import functional as F +from torch.nn.parameter import Parameter +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list +from deepspeed.runtime.zero.utils import is_zero_param +from abc import ABC, abstractmethod +from typing import Iterable, Any, Optional, List, Tuple +from .fusedqkv_utils import shard_value_with_share_qk, shard_chunk_mlp, prepare_tp_fused_qkvw +from deepspeed.runtime.tensor_parallel import AUTOTP_MODE +from copy import deepcopy +from typing import Union + +__all__ = [ + "TensorParallel_Layer", "LinearAllreduce", "LinearLayer", "LmHeadLinearAllreduce", "Yuan_LinearAllreduce", + "Yuan_LinearLayer", "GateUpPack_LinearLayer", "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer" +] + +DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE +DS_IS_REPLACED_MODULE = 'ds_is_replaced_module' +DS_TENSOR_MODEL_PARALLEL = 'tensor_model_parallel' + + +def get_auto_tp_mode(): + global DEEPSPEED_AUTOTP_MODE + return DEEPSPEED_AUTOTP_MODE + + +def is_autotp_training_mode(): + global DEEPSPEED_AUTOTP_MODE + return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING + + +def set_autotp_mode(training=False): + """ + Set the DEEPSPEED_AUTOTP_MODE based on the training flag + """ + global DEEPSPEED_AUTOTP_MODE + if training: + DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.TRAINING + else: + DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE + + +def add_bias(input, bias): + if bias is None: + return input + if is_autotp_training_mode(): + # Training mode - avoid inplace to ensure correct autograd + input = input + bias + return input + else: + input += bias + return input + + +class RowParallel(torch.autograd.Function): + """ + A custom autograd function for performing row-wise parallelism. + """ + + @staticmethod + def symbolic(graph, input): + """Symbolic function for tracing.""" + return input + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, is_inference_mode: bool) -> torch.Tensor: + """ + Forward pass. + """ + ctx.group = group + if group == None: + return input + if is_inference_mode: + dist.inference_all_reduce(input, group=group) + else: + dist.all_reduce(input.contiguous(), group=group) + return input + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None]: + """ + Backward pass. + """ + return None, grad_output, None + + +class AsyncColumnParallel(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, weight, bias) -> torch.Tensor: + """ + Forward pass. + """ + ctx.use_bias = bias is not None + ctx.group = group + output = torch.matmul(input, weight.transpose(-1, -2)) + if bias is not None: + output = add_bias(output, bias) + + ctx.save_for_backward(input, weight) + + return output + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]: + + input, weight = ctx.saved_tensors + grad_input = grad_output.matmul(weight) + handle = dist.all_reduce(grad_input.contiguous(), group=ctx.group, async_op=True) + grad_weight = grad_output.view(-1, grad_output.shape[-1]).t().matmul(input.view(-1, input.shape[-1])) + grad_bias = grad_output.sum(0) if ctx.use_bias else None + handle.wait() + return None, grad_input, grad_weight, grad_bias + + +class ColumnParallel(torch.autograd.Function): + """ + Custom autograd function for column-wise parallelism. + """ + + @staticmethod + def symbolic(graph, input): + """Symbolic function for tracing.""" + return dist.all_reduce(input.contiguous(), dist.get_tensor_model_parallel_group()) + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor: + """ + Forward pass. + """ + ctx.group = group + return input + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]: + """ + Backward pass. + """ + if ctx.group == None: + return None, grad_output + + dist.all_reduce(grad_output.contiguous(), group=ctx.group) + return None, grad_output + + +class TensorParallel_Layer(nn.Module, ABC): + """ + A base class for model layers with tensor parallelism support. + This class is designed to be extended by specific layers that require distributed + operations and parameter gather/partitioning during inference or training. + + Attributes: + mode (str): The mode of operation[INFERENCE or TRAINING], default is "INFERENCE". + mp_group (Optional[dist.ProcessGroup]): The process group used for model parallelism. + tp_world_size (int): The world size of tensor parallelism, i.e., the number of parallel workers. + tp_index (int): The rank (ID) of the current worker in tensor parallelism. + support_training (bool): Flag indicating whether the layer supports training (default: False). + name (Optional[str]): The name of the layer, if provided. + """ + ##### Initialize Parameter List ##### + + # keep_module_on_host determines whether to keep the module on the host. + # Checkpoints are first loaded to the host (sometimes directly from disk to avoid filling host memory), + # so an additional copy is unnecessary. + keep_module_on_host: bool = False + + ##### Runtime Parameter List ##### + tp_overlap_comm: bool = False + """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ + + def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any): + """ + Initializes the TensorParallel_Layer with optional model parallelism group and layer name. + + Args: + mp_group (Optional[dist.ProcessGroup]): The process group for model parallelism. + If None, no model parallelism is set. + """ + super().__init__() + self.support_training: bool = False + self.mp_group = mp_group + if mp_group is not None: + self.tp_world_size: int = dist.get_world_size(self.mp_group) + self.tp_index: int = dist.get_rank(self.mp_group) + else: + self.tp_world_size: int = 1 + self.tp_index: int = 0 + + # backward compatibility + self.world_size = self.tp_world_size + self.rank = self.tp_index + + self.name = getattr(self, 'name', None) + if kwargs.get('name') is not None: + self.name = kwargs.get('name') # Set the layer name if provided. + + @classmethod + def set_keep_module_on_host(cls, value: bool): + """ + Set the static variable keep_module_on_host. + + Args: + value (bool): The new value for keep_module_on_host. + """ + cls.keep_module_on_host = value + + @abstractmethod + def forward(self, input): + """ + Forward pass method. Must be implemented by subclasses to define layer-specific operations. + """ + pass + + @abstractmethod + def gather_params(self, params_list): + """ + Gathers parameters across devices for distributed training. Must be implemented by subclasses in "TRAINING" mode. + """ + pass + + @abstractmethod + def _tp_partition(self, params_list: List[torch.Tensor]): + """ + Partitions the parameters for tensor parallelism. + It is necessary to ensure that this function only involves the logic of params partitioning. + """ + pass + + def config_requires_grad(self, weight): + if weight is not None: + if self.is_training_mode(): + if weight.requires_grad is None: + weight.requires_grad = True + else: + weight.requires_grad = False + + def config_tp_params(self, weight): + """ + Configures the weight tensor for training with tensor parallelism. This includes enabling gradients + and associating necessary methods for parameter gathering and partitioning. + + Args: + weight (Optional[torch.Tensor]): The weight tensor to configure for tensor parallelism. + If None, no action is taken. + """ + # # The RNG states have already been synchronized in init_inference. + if self.is_training_mode(): + assert self.support_training, "No implementation of backward." + if weight is not None: + self.config_requires_grad(weight) + weight.gather_params = self.gather_params + weight._tp_partition = self._tp_partition + setattr(weight, DS_TENSOR_MODEL_PARALLEL, True) + setattr(weight, DS_IS_REPLACED_MODULE, True) + + def is_training_mode(self): + global DEEPSPEED_AUTOTP_MODE + return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING + + def __deepcopy__(self, memo): + # This function is designed for + # 'mp_group' (a 'ProcessGroup') cannot be pickled during deepcopy in some usage. + cls = self.__class__ + new_obj = cls.__new__(cls) + + for key, value in vars(self).items(): + if key == 'mp_group': + new_obj.mp_group = self.mp_group + else: + setattr(new_obj, key, deepcopy(value, memo)) + + memo[id(self)] = new_obj + return new_obj + + def extra_repr(self): + out_features, in_features = None, None + if self.weight is not None: + out_features, in_features = self.weight.ds_shape[-2:] if is_zero_param( + self.weight) else self.weight.shape[-2:] + dtype = self.weight.dtype if self.weight is not None else None + return "in_features={}, out_features={}, bias={}, dtype={}".format(in_features, out_features, self.bias + is not None, dtype) + + def move(self, tensor): + # TODO: consider the timing of deletion + # to save host resources when DP > 1。 + + # keep_module_on_host is used to keep the module on the host. Checkpoints are loaded to the host first (in some + # cases it can be done from the disk even to prevent filling host's memory), thus no need to create a new copy. + if tensor.is_meta: + # Keep tensor in meta device if tensor is meta. + return tensor + else: + device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name() + return_new_copy = not self.__class__.keep_module_on_host + + # Using new tensors help in freeing memory (after split for example) was done before by calling clone(). + # Using copy=True instead of clone() will help in case of cpu --> cpu. + # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced. + cloned_tensor = tensor.to(device, copy=return_new_copy) + + if return_new_copy: + # free the memory of the original tensor to reduce memory peak + # Equivalent to directly deleting the tensor reference outside the function. + # see https://github.com/microsoft/DeepSpeed/pull/4353 + tensor.data = torch.empty(0, device=tensor.device) + return cloned_tensor + + +def configure_tensor_parallel_runtime(config): + runtime_keys = ['tp_overlap_comm'] + for key in runtime_keys: + if hasattr(config, key): + setattr(TensorParallel_Layer, key, getattr(config, key)) + + +class GatherReplacedLayerParams: + """ + A context manager for gathering parameters of a replaced layer, enabling partitioning and gathering functionality + based on the configuration of the model. + """ + + def __init__(self, + params: Union[Iterable[torch.Tensor], torch.Tensor], + module: torch.nn.Module, + enabled: bool = True): + """ + Initialize the context manager to handle parameter gathering and partitioning for a replaced layer. + + Args: + params (Iterable or torch.Tensor): A collection or single parameter to manage. + module (torch.nn.Module): The module that these parameters belong to. + enabled (bool): Flag indicating whether the parameter management is enabled (default: True). + """ + self.enabled = enabled + self.module = module + if not enabled: + return + + # Ensure params is a list, whether it's a single param or iterable (e.g., model.parameters()) + if isinstance(params, Iterable) and not isinstance(params, torch.Tensor): + self.params: List[torch.Tensor] = list(params) # Convert generators to a list for multiple iterations + else: + self.params: List[torch.Tensor] = [params] # Wrap single parameter in a list for uniform processing + + # Check if the parameters belong to a replaced layer (indicated by a specific attribute) + if not any(self._is_replaced_module_weight(p) for p in params): + self.enabled = False + return + + def _is_replaced_module_weight(self, param: torch.Tensor) -> bool: + """ + Helper function to determine if a parameter belongs to a replaced module. + + Args: + param (torch.Tensor): The parameter to check. + + Returns: + bool: True if the parameter belongs to a replaced module, False otherwise. + """ + return getattr(param, DS_IS_REPLACED_MODULE, False) + + def __enter__(self) -> None: + """ + Enter the context manager. If enabled, gather parameters for the replaced module. + """ + if self.enabled: + self.params[0].gather_params(self.params) + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """ + Exit the context manager. If enabled, partition the parameters for the replaced module. + """ + #TODO : Check whether there are any missing attributes. + if self.enabled: + self.params[0]._tp_partition(self.params) + + +class LinearAllreduce(TensorParallel_Layer): + + def __init__(self, module, mp_group, **kwargs): + super(LinearAllreduce, self).__init__(mp_group, **kwargs) + self.weight = module.weight + self.bias = module.bias + + self._tp_partition([self.weight, self.bias]) + self.support_training = True + self.config_tp_params(self.weight) + if self.bias is not None: + # bias here is not tp params + self.config_requires_grad(self.bias) + + def forward(self, input): + output = torch.matmul(input, self.weight.transpose(-1, -2)) + output = RowParallel.apply(self.mp_group, output, not self.is_training_mode()) + if self.bias is not None: + output = add_bias(output, self.bias) + return output + + @torch.no_grad() + def gather_params(self, params_list): + + for idx, param in enumerate(params_list): + if param is None or idx > 0: + # don't gather bias + return + params_list[idx].data_partition = param.data + param = param.transpose(0, 1).contiguous() + + output_param = torch.empty(self.tp_world_size * param.shape[0], + param.shape[1], + dtype=param.dtype, + device=param.device) + dist.all_gather_into_tensor(output_param, param, group=self.mp_group) + params_list[idx].data = output_param.transpose(0, 1).contiguous() + return + + @torch.no_grad() + def _tp_partition(self, params_list): + + if not self.is_training_mode(): + self.uneven_partition(params_list) + return + + else: + for idx, param in enumerate(params_list): + if param is None: + # don't slipt bias + return + if idx > 0: # move bias to device at initialization + _partition = self.move(param).detach() + params_list[idx].data = _partition + return + + _partition = torch.chunk(param, self.tp_world_size, dim=-1)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + def uneven_partition(self, params_list): + for idx, param in enumerate(params_list): + if param is None or idx > 0: + # don't slipt bias + return + assert self.name is not None, "The module name must be provided in the initialization." + _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size, + self.name), + dim=1)[self.tp_index] + + _partition = self.move(_partition).detach() + params_list[idx].data = _partition + + +#remove kwargs from partition. +class LinearLayer(TensorParallel_Layer): + + def __init__(self, module, mp_group=None, skip_partition=False, **kwargs): + super(LinearLayer, self).__init__(mp_group, **kwargs) + self.weight = module.weight + self.bias = module.bias + if not skip_partition: + self._tp_partition([self.weight, self.bias]) + self.support_training = True + self.config_tp_params(self.weight) + if self.bias is not None: + self.config_tp_params(self.bias) + + def forward(self, input): + if not self.__class__.tp_overlap_comm: + if getattr(self, 'mp_group', None) is not None: + input = ColumnParallel.apply(self.mp_group, input) + output = torch.matmul(input, self.weight.transpose(-1, -2)) + if self.bias is not None: + output = add_bias(output, self.bias) + else: + output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias) + + return output + + @torch.no_grad() + def gather_params(self, params_list): + # Does not support uneven shard. + for idx, param in enumerate(params_list): + + params_list[idx].data_partition = param.data + output_param = torch.empty((self.tp_world_size * param.shape[0], *param.shape[1:]), + dtype=param.dtype, + device=param.device) + dist.all_gather_into_tensor(output_param, param, group=self.mp_group) + params_list[idx].data = output_param.contiguous() + + @torch.no_grad() + def _tp_partition(self, params_list): + + if not self.is_training_mode(): + self.uneven_partition(params_list) + return + for idx, param in enumerate(params_list): + if param is None: + return + #split bias if provide + _partition = torch.chunk(param, self.tp_world_size, dim=0)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + def uneven_partition(self, params_list): + + for idx, param in enumerate(params_list): + if param is None: + #split bias if provide + return + assert self.name is not None, "The module name must be provided in the initialization." + _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size, + self.name), + dim=0)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + # for bwc + @classmethod + def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None): + if weight is not None: + in_features = weight.shape[1] + out_features = weight.shape[0] + linear = nn.Linear(in_features, out_features, bias=(bias is not None)) + linear.weight.data = weight + if bias is not None: + linear.bias.data = bias + else: + in_features = weight_shape[1] + out_features = weight_shape[0] + linear = nn.Linear(in_features, out_features, bias=(bias is not None)) + return cls(linear, skip_partition=True) + + +class FusedModuleWrapper: + + def __init__(self, fused_module: nn.Module): + self.fused_module = fused_module + + def __getattr__(self, module): + return self.fused_module + + +class fused_LinearLayer(LinearLayer): + + def __init__(self, module, mp_group, skip_partition=False, **kwargs): + assert kwargs.get('fused_module') is not None, "'fused_module' is required but not provided" + # Use the warp class to avoid module circular references. + self.fused_module = FusedModuleWrapper(kwargs.get('fused_module')) + super().__init__(module, mp_group, skip_partition, **kwargs) + + @torch.no_grad() + def _tp_partition(self, params_list): + for idx, param in enumerate(params_list): + if param is None: + return + + _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index) + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + +class conv_LinearLayer(LinearLayer): + + @torch.no_grad() + def _tp_partition(self, params_list): + weight = None + bias = None + if len(params_list) == 1: + weight = params_list[0] + elif len(params_list) == 2: + weight, bias = params_list[0], params_list[1] + _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name), + dim=1)[self.tp_index] + _partition = self.move(_partition).detach() + weight.data = _partition + + if bias is not None: + _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name), + dim=0)[self.tp_index] + _partition = self.move(_partition).detach() + + bias.data = _partition + + +#override the subclasses related to weight splitting. +class Yuan_LinearAllreduce(LinearAllreduce): + + #Yuan2 + @torch.no_grad() + def _tp_partition(self, params_list): + weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index, + self.tp_world_size, False) + params_list[0].data = weight + if bias is not None: + params_list[1].data = bias + + +class Yuan_LinearLayer(LinearLayer): + #Yuan2 + @torch.no_grad() + def _tp_partition(self, params_list): + weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index, + self.tp_world_size, True) + params_list[0].data = self.move(weight).detach() + if bias is not None: + params_list[1].data = self.move(bias).detach() + + +class GateUpPack_LinearLayer(LinearLayer): + # chatGLM2, chatGLM2 + @torch.no_grad() + def _tp_partition(self, params_list): + weight, bias = shard_chunk_mlp(params_list[0].data, params_list[1], self.tp_index, self.tp_world_size) + params_list[0].data = self.move(weight).detach() + if bias is not None: + params_list[1].data = self.move(bias).detach() + + +class Conv_LinearALlreduce(LinearAllreduce): + + @torch.no_grad() + def _tp_partition(self, params_list): + for idx, param in enumerate(params_list): + if param is None: + return + param.data = param.data.transpose(-1, -2).contiguous() + + _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name), + dim=1)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + +#override the subclasses related to fwd/bwd. +class LmHeadLinearAllreduce(LinearAllreduce): + + def __init__(self, module, mp_group, **kwargs): + # set the fixed name before partition + self.name = "lm_head" + + # In some tied_embedding cases, only the lm head is sharded, while the word embedding is not. + # Reinitialization is used to decouple them and prevent the word embedding from being sharded. + # This should also be effective for cases where both are sharded in tied_embedding scenarios. + + # TODO: Training scenario-related tests, is it necessary to re-implement the vocab parallel module? + module.weight = nn.Parameter(module.weight.clone().detach()) + if hasattr(module, 'bias') and module.bias is not None: + module.bias = nn.Parameter(module.bias.clone().detach()) + super().__init__(module, mp_group, **kwargs) + + def forward(self, input): + input_shard_size = get_shard_size(input.shape[-1], self.tp_world_size, "lm_head") + input_shard_offset = sum(get_shard_size_list(input.shape[-1], self.tp_world_size, "lm_head")[0:self.tp_index]) + output = torch.matmul(input[:, :, input_shard_offset:input_shard_offset + input_shard_size], + self.weight.transpose(-1, -2)) + if self.mp_group is not None: + dist.inference_all_reduce(output, group=self.mp_group) + if self.bias is not None: + output = add_bias(output, self.bias) + return output + + +class TensorParallelConv2d(nn.Module): + + def __init__(self, conv, rank, world_size, shard_by_oc): + super().__init__() + self.rank = rank + self.world_size = world_size + self.shard_by_oc = shard_by_oc + self.shard_weights(conv) + + # Split along the input/output channel depending on whether it is the last conv layer. + def shard_weights(self, conv): + if self.shard_by_oc: + total_size = conv.weight.shape[0] + else: + total_size = conv.weight.shape[1] + bias_data = None + cols_per_rank = [0] + for i in range(self.world_size - 1, -1, -1): + cols = total_size // self.world_size + if i < total_size % self.world_size: + cols += 1 + cols_per_rank.append(cols_per_rank[-1] + cols) + weight_data = conv.weight.data + if self.shard_by_oc: + # not last conv layer, split output channel + weight_data = weight_data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] + if conv.bias is not None: + bias_data = conv.bias.data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] + else: + # last conv layer, split input channel + weight_data = weight_data[:, cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] + if conv.bias is not None: + bias_data = conv.bias.data / float(self.world_size) + self.conv = nn.Conv2d(weight_data.shape[1], weight_data.shape[0], conv.kernel_size, conv.stride, conv.padding, + conv.dilation, conv.groups, conv.bias is not None, conv.padding_mode) + self.conv.weight = torch.nn.Parameter(weight_data) + if conv.bias is not None: + self.conv.bias = torch.nn.Parameter(bias_data) + del conv + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return self.conv(input) + + +class TensorParallelOcShardConv2d(TensorParallelConv2d): + + def __init__(self, conv, rank, world_size): + super().__init__(conv, rank, world_size, True) + + +class TensorParallelIcShardConv2d(TensorParallelConv2d): + + def __init__(self, conv, rank, world_size): + super().__init__(conv, rank, world_size, False) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + out = self.conv(input) + if self.world_size > 1: + dist.inference_all_reduce(out) + return out + + +class Normalize(nn.Module): + + def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None, bias=None): + super(Normalize, self).__init__() + if weight is not None: + self.weight = weight + self.bias = bias + else: + self.norm = nn.LayerNorm(dim, eps=eps).to(dtype).to(get_accelerator().current_device_name()) + self.weight = self.norm.weight + self.bias = self.norm.bias + + self.eps = eps + + def forward(self, input): + return nn.functional.layer_norm(input, input.shape[-1:], self.weight, self.bias, eps=self.eps) + + +class EmbeddingLayer(nn.Module): + + def __init__(self, weight_shape=None, dtype=torch.half, weight=None, bias=None): + super(EmbeddingLayer, self).__init__() + if weight is None: + self.weight = Parameter( + torch.empty(weight_shape[0], + weight_shape[1], + dtype=dtype, + device=get_accelerator().current_device_name())) + else: + self.weight = weight + + def forward(self, input): + return F.embedding(input, self.weight) + + +class OPTEmbedding(EmbeddingLayer): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, weight_shape=None, weight=None, bias=None): + # OPT is set up so that if padding_idx is specified then offset the embedding ids by 2 + # and adjust num_embeddings appropriately. Other models don't have this hack + self.offset = 2 + super().__init__(weight_shape, weight=weight) + + def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, position_ids: int = 0): + """`input_ids_shape` is expected to be [bsz x seqlen].""" + attention_mask = attention_mask.long() + + # create positions depending on attention_mask + positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1 + + # cut positions if `past_key_values_length` is > 0 + positions = positions[:, past_key_values_length:] + + return super().forward(positions + self.offset) + + +class RMSNormalize(nn.Module): + + def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None): + super(RMSNormalize, self).__init__() + if weight is not None: + self.weight = weight + else: + self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=get_accelerator().current_device_name())) + + self.eps = eps + + def forward(self, hidden_states): + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return hidden_states * self.weight diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/load_checkpoint.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/load_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..862628fa7b4b46ad2cd8bf2c3f5706233bda0627 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/load_checkpoint.py @@ -0,0 +1,285 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch import nn +from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference +from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference +from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference +from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference + +import deepspeed.ops.transformer as transformer_inference +from .layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding, RMSNormalize +import torch +import gc +from deepspeed.accelerator import get_accelerator +import re + + +def load_model_with_checkpoint(r_module, + sd, + mp_replace, + ckpt_type, + ckpt_mp_size, + weight_quantizer=None, + rank=0, + container=None): + error_msgs = [] + + def prefix_check(): + # if keys start with 'model.' or 'transformer.', don't skip level 0 prefix + for key in sd[0].keys(): + # OPT models + if re.match("^model[.]", key): + return False + # BLOOM models + if re.match("^transformer[.]", key): + return False + return True + + skip_level_0_prefix = prefix_check() and container.policy.use_load_prefix + + def transpose(data): + with torch.no_grad(): + data = data.contiguous() + data1 = data.transpose(-1, -2).reshape(-1) + data.reshape(-1).copy_(data1) + data1 = None + return data.reshape(data.shape[-1], data.shape[-2]) + + def load(module, prefix): + args = (sd[0], prefix, {}, True, [], [], error_msgs) + + if hasattr(module, 'weight'): + module.weight = mp_replace.copy(module.weight.data, sd[0][prefix + 'weight']) + if prefix + 'bias' in sd[0].keys(): + if module.bias.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), + requires_grad=module.bias.data.requires_grad) + module.bias = mp_replace.copy(module.bias.data, sd[0][prefix + 'bias']) + args = None + gc.collect() + + def load_transformer_layer(module, prefix): + if ckpt_type == "tp": + + def load_parameters(module, prefix): + for n, p in module.named_parameters(): + if prefix + n in sd[0] and len(n.split('.')) == 1: + if type(sd[0][prefix + n]) is list: + tmp_data, scale = sd[0][prefix + n] + tmp_data = tmp_data + scale = scale.to(get_accelerator().current_device_name()) + # set the quantizer number of groups using the checkpoint scale shape + weight_quantizer.num_groups = scale.shape[0] + else: + tmp_data = sd[0][prefix + n].to(get_accelerator().current_device_name()) + scale = None + src_shape = tmp_data.shape + dst_shape = p.shape + inner_dim = 1 if tmp_data.dtype == torch.int8 else 0 + outer_dim = 0 if tmp_data.dtype == torch.int8 else 1 + if (len(src_shape) == 2 and len(dst_shape) == 2): + if (src_shape[inner_dim] == dst_shape[0] and src_shape[outer_dim] == dst_shape[1]): + if tmp_data.dtype != torch.int8: + p = weight_quantizer.quantize( + transpose(tmp_data) if weight_quantizer.q_int8 else tmp_data) + else: + p = torch.nn.parameter.Parameter(tmp_data, requires_grad=False) + p.scale = scale + setattr(module, n, p) + else: + dim = inner_dim if src_shape[inner_dim] != dst_shape[0] else outer_dim + dim1 = 0 if src_shape[inner_dim] != dst_shape[0] else 1 + if src_shape[dim] > dst_shape[dim1]: + weight_partition = torch.split(tmp_data, dst_shape[dim1], dim=dim)[rank].to( + get_accelerator().current_device_name()) + assert tmp_data.dtype != torch.int8 or scale.numel() > weight_quantizer.num_groups * (rank+1), \ + '''ERROR: We require the quantization scales for larger TP-size when loading INT8 checkpoint!\ + Please use the FP16 checkpoint to generate INT8 checkpoint with the sharding parameters!''' + scale = scale.view(-1)[weight_quantizer.num_groups * (rank + 1):].reshape( + weight_quantizer.num_groups, -1).contiguous() + else: + assert tmp_data.dtype != torch.int8, \ + '''Merging of the checkpoints are not supported when using INT8 checkpoint! \ + Please use a as many GPUs as TP-size for the checkpoint''' + all_data = [ + sd[j][prefix + n] if type(sd[j][prefix + n]) is list else sd[j][prefix + n].to( + get_accelerator().current_device_name()) for j in range(len(sd)) + ] + # Check if the weight tensor is for the QKV parameter + if src_shape[1] == (3 * src_shape[0]) // ckpt_mp_size: + qkv_size = src_shape[outer_dim] // 3 + src_split = [ + torch.split(src[0].data, qkv_size, dim=outer_dim) for src in all_data + ] + + weight_partition = torch.cat([ + torch.cat([qkv_s[i] for qkv_s in src_split], axis=outer_dim) + for i in range(len(src_split[0])) + ], + dim=dim) + else: + weight_partition = torch.cat([ + ad[0].to(get_accelerator().current_device_name()) + if type(ad) is list else ad for ad in all_data + ], + dim=dim) + if tmp_data.dtype == torch.int8: + scale = torch.cat( + [ad[1].to(get_accelerator().current_device_name()) for ad in all_data], + dim=dim) + + if tmp_data.dtype != torch.int8: + weight_partition = weight_quantizer.quantize( + transpose(weight_partition), \ + parallel_dim=(0 if dim == 1 else 1)) if weight_quantizer.q_int8 else \ + weight_quantizer.quantize(weight_partition) + else: + weight_partition = torch.nn.parameter.Parameter(weight_partition, + requires_grad=False) + weight_partition.scale = scale + setattr(module, n, weight_partition) + else: + if src_shape[0] == dst_shape[0]: + p.data.copy_(tmp_data) + else: + if src_shape[0] > dst_shape[0]: + bias_split = torch.split(tmp_data, dst_shape[-1])[rank].to( + get_accelerator().current_device_name()).contiguous() + p.data.copy_(bias_split) + else: + # Check if the weight tensor is for the QKV parameter + if src_shape[0] == (3 * r_module.config.hidden_size) // ckpt_mp_size: + qkv_size = src_shape[0] // 3 + src_split = [ + torch.split(sd[j][prefix + n], qkv_size, dim=0) for j in range(len(sd)) + ] + + p.data.copy_( + torch.cat([ + torch.cat([qkv_s[i] for qkv_s in src_split], axis=0) + for i in range(len(src_split[0])) + ], + dim=0).to(get_accelerator().current_device_name()).contiguous()) + else: + p.data.copy_( + torch.cat([sd[j][prefix + n] for j in range(len(sd))], + dim=0).to(get_accelerator().current_device_name()).contiguous()) + + load_parameters(module, prefix) + for n, child in module.named_children(): + load_parameters(child, prefix + n + '.') + else: + container.load_params(module, sd[0], weight_quantizer, mp_replace, prefix) + + try: + import transformers + OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding + if hasattr(transformers.models, "llama"): + LlamaRMSNorm = transformers.models.llama.modeling_llama.LlamaRMSNorm + else: + LlamaRMSNorm = None + except: + OPTLearnedPositionalEmbedding = None + try: + from fairscale.nn.model_parallel.layers import ( + ColumnParallelLinear, + ParallelEmbedding, + RowParallelLinear, + ) + except: + ColumnParallelLinear = None + ParallelEmbedding = None + RowParallelLinear = None + try: + from llama.model import RMSNorm + except: + RMSNorm = None + layer_policies = { + nn.Linear: load, + nn.Embedding: load, + nn.LayerNorm: load, + EmbeddingLayer: load, + LinearLayer: load, + Normalize: load, + transformer_inference.DeepSpeedTransformerInference: load_transformer_layer, + DeepSpeedBloomInference: load_transformer_layer, + DeepSpeedGPTInference: load_transformer_layer, + DeepSpeedBERTInference: load_transformer_layer, + DeepSpeedMegatronGPTInference: load_transformer_layer, + DeepSpeedOPTInference: load_transformer_layer, + DeepSpeedLlama2Inference: load_transformer_layer, + OPTLearnedPositionalEmbedding: load, + OPTEmbedding: load, + LlamaRMSNorm: load, + RMSNormalize: load, + ColumnParallelLinear: load, + ParallelEmbedding: load, + RowParallelLinear: load, + RMSNorm: load + } + + all_ds_ids = {} + + def load_module_recursive(module, prefix='', level=0): + for name, child in module.named_children(): + if child.__class__ in layer_policies: + checking_key = prefix + name + '.' + if not any(checking_key in item for item in sd[0].keys()): + if hasattr(child, 'weight') and \ + (hasattr(child.weight, 'ds_id') and \ + child.weight.ds_id in all_ds_ids): + prefix1 = all_ds_ids[child.weight.ds_id] + if child.__class__ is nn.Linear: + child = LinearLayer.from_weights(weight=all_ds_ids[child.weight.ds_id]) + setattr(module, name, child) + continue + child_params = list(child.parameters()) + if len(child_params) > 0 and (child_params[0].numel() == 0 or child_params[0].is_meta): + if child.weight.is_meta: + ds_shape = child.weight.shape + else: + ds_shape = child.weight.ds_shape + if child.__class__ is nn.LayerNorm: + child = Normalize(dim=ds_shape[-1], dtype=child.weight.dtype, eps=child.eps) + setattr(module, name, child) + elif child.__class__ in [nn.Linear, ColumnParallelLinear, RowParallelLinear]: + child = LinearLayer.from_weights(weight_shape=child.weight.shape, + dtype=child.weight.dtype, + bias=child.bias) + setattr(module, name, child) + elif child.__class__ is OPTLearnedPositionalEmbedding: + child = OPTEmbedding(weight_shape=ds_shape) + setattr(module, name, child) + elif child.__class__ in [LlamaRMSNorm, RMSNorm]: + child = RMSNormalize(dim=ds_shape[-1], + dtype=child.weight.dtype, + eps=child.eps if hasattr(child, 'eps') else child.variance_epsilon) + setattr(module, name, child) + else: + ds_id = None + if hasattr(child.weight, 'ds_id'): + ds_id = child.weight.ds_id + child = EmbeddingLayer(weight_shape=ds_shape, dtype=child.weight.dtype) + if ds_id is not None: + all_ds_ids[ds_id] = child.weight + setattr(module, name, child) + layer_policies[child.__class__](child, prefix + name + '.') + else: + load_module_recursive( + child, + prefix if (level == 0 and ckpt_type == 'pp') and skip_level_0_prefix else \ + prefix + name + '.', + level + 1) + + load_module_recursive(r_module) + + for sd_ in sd: + del sd_ + sd = None + gc.collect() diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/module_quantize.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/module_quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..1f5b2f8a1d2861a50878e262a61b783a9324ad4e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/module_quantize.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + + +def quantize_transformer_layer(orig_layer_impl, model, megatron=False, preln=False): + """ Quantize bert-style transformer layers with DeepSpeed's transformer layer + Arguments: + orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for, + e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer + model (torch.nn.Module): user's nn.module representing their model + + megatron (bool): megatron model-parallel implementation (this is supported for inference only) + preln (bool): does the original layer implementation do pre or post layer norm? + + Note: For Bert kind of models, we inject based on the DeepSpeed-Example models, if not setting huggingface flag. + + Returns: + Updated nn.module with quantized transformer layers + """ + + def quantize_weight(weight): + return weight.to(torch.int8) + + def megatron_layer_quantize(layer): + layer.attention.query_key_value.weight.data = quantize_weight(layer.attention.query_key_value.weight.data) + layer.attention.dense.weight.data = quantize_weight(layer.attention.dense.weight.data) + layer.mlp.dense_h_to_4h.weight.data = quantize_weight(layer.mlp.dense_h_to_4h.weight.data) + layer.mlp.dense_4h_to_h.weight.data = quantize_weight(layer.mlp.dense_4h_to_h.weight.data) + + def bert_layer_quantize(layer): + layer.attention.self.query.weight.data = quantize_weight(layer.attention.self.query.weight.data) + layer.attention.self.key.weight.data = quantize_weight(layer.attention.self.key.weight.data) + layer.attention.self.value.weight.data = quantize_weight(layer.attention.self.value.weight.data) + layer.attention.output.dense.weight.data = quantize_weight(layer.attention.output.dense.weight.data) + if preln: + layer.intermediate.dense_act.weight.data = quantize_weight(layer.intermediate.dense_act.weight.data) + else: + layer.intermediate.dense.weight.data = quantize_weight(layer.intermediate.dense.weight.data) + layer.output.dense.weight.data = quantize_weight(layer.output.dense.weight.data) + + def quantize_fn(child): + if megatron: + # Quantize megatron GPT2 / GPT3 trained model + megatron_layer_quantize(child) + else: + # Quantize either DeepSpeed or HuggingFace trained model + bert_layer_quantize(child) + + return child + + return quantize_module(model=model, orig_class=orig_layer_impl, quantize_fn=quantize_fn) + + +def quantize_module(model, orig_class, quantize_fn): + policy = {orig_class: quantize_fn} + return _quantize_module(model, policy) + + +def _quantize_module(model, policies): + for name, child in model.named_children(): + if child.__class__ in policies: + orig = repr(child) + setattr(model, name, policies[child.__class__](child)) + new = getattr(model, name) + else: + _quantize_module(child, policies) + + return model diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/policy.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..41df2b85dc0c16fb1b1c21942f5c963fcdc02ad1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/policy.py @@ -0,0 +1,224 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from deepspeed.utils.types import ActivationFuncType, NormType +import torch +from deepspeed.accelerator import get_accelerator + +transformer_param_names = ( + 'attn_qkvw', \ + 'attn_qkvb', \ + 'attn_ow' , \ + 'attn_ob', \ + 'inter_w', \ + 'inter_b', \ + 'output_w', \ + 'output_b', \ + 'attn_nw', \ + 'attn_nb', \ + 'norm_w', \ + 'norm_b') + + +class DSPolicy(ABC): + _orig_layer_class = None + + def __init__(self): + self.cuda_graph_supported = False + + @abstractmethod + def attention(self): + """ + Returns attention qkv and dense parameters + weight: (3*hidden, hidden) and (hidden, hidden) + bias: (3*hidden) and (hidden) + """ + raise NotImplementedError + + +class TransformerPolicy(DSPolicy): + # a static class variable containing the HuggingFace model configuration. + # see e.g., transformers.models.opt.configuration_opt.OPTConfig + hf_model_config = None + + def __init__( + self, + inference=True, + linear_layer=True, + scale_attention=True, + megatron_v2=False, + use_mup=False, + # the type of activation function used in MLP + mlp_act_func_type=ActivationFuncType.GELU, + # applies layer norm before attention if `pre_attn_norm` is set to True + pre_attn_norm=True, + # this flag shows whether or not using prefix in loading the checkpoint + use_load_prefix=False, + # whether or not the qkv is stored in the split-format + split_qkv=True, + # Type of normalization to perform + norm_type=NormType.LayerNorm): + super().__init__() + self.cuda_graph_supported = False + self.inference = inference + self.linear_layer = linear_layer + self.scale_attention = scale_attention + self.is_megatron_v2 = megatron_v2 + self.use_mup = use_mup + self.mlp_act_func_type = mlp_act_func_type + self.pre_attn_norm = pre_attn_norm + self.use_load_prefix = use_load_prefix + self.split_qkv = split_qkv + self.norm_type = norm_type + + @abstractmethod + def attention(self): + """ + Returns attention qkv and dense parameters + weight: (3*hidden, hidden) and (hidden, hidden) + bias: (3*hidden) and (hidden) + """ + raise NotImplementedError + + @abstractmethod + def get_hidden_heads(self): + """ + return hidden_size and number of heads + """ + raise NotImplementedError + + @abstractmethod + def mlp(self): + """ + Returns mlp intermediate and output + weight: (intermediate, hidden) and (hidden, intermediate) + bias: (intermediate) and (hidden) + """ + raise NotImplementedError + + @abstractmethod + def layernorm(self): + """ + Returns LayerNorms used in transformer layer + Post-Attention and pre/post layer norm + gamma and beta with shape: (hidden) + """ + raise NotImplementedError + + +# TODO (lekurile): This function exists in base container as well, consolidate as some point +def transpose(data): + with torch.no_grad(): + data = data.contiguous() + data1 = data.transpose(-1, -2).reshape(-1) + data.reshape(-1).copy_(data1) + data1 = None + return data.reshape(data.shape[-1], data.shape[-2]) + + +# TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point +def _transpose(x, heads=1, mp_replace=None): + heads = heads // mp_replace.mp_size # type: ignore + outer_dim = -1 + attention_head_size = x.shape[outer_dim] // heads + new_x_shape = x.size()[:outer_dim] + (heads, attention_head_size) + x_1 = x.view(*new_x_shape) + (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=-1) + if len(q.shape) > 2: + new_shape = (q.shape[0], ) + (-1, ) + return torch.cat((q.reshape(new_shape), k.reshape(new_shape), v.reshape(new_shape)), + dim=outer_dim).reshape(x.shape) + else: + return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) + + +# This checks if the parameter exits in the checkpoint file and maybe copies it into the corresponding destination tensor. +# Note that not all parameters are saved in one checkpoint, that's why we always need to check if they exist! +def maybe_copy(module, + sd, + weight_quantizer, + mp_replace, + dst_name, + src_name, + qkv=False, + megatron_v2=False, + split_qkv=False, + heads=1): + if src_name in sd: + dst = getattr(module, dst_name) + tmp = sd[src_name] + if len(dst.shape) == 1: + if split_qkv: + dst = mp_replace.strided_copy(dst, tmp, num_splits=3) + else: + dst = mp_replace.copy(dst, tmp) + if qkv and megatron_v2: + dst = torch.nn.parameter.Parameter(_transpose(dst, heads=heads, mp_replace=mp_replace).contiguous()) + else: + if split_qkv: + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ + (transpose(tmp).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) + else: + if qkv and megatron_v2: + tmp = _transpose(transpose(tmp), heads=heads, mp_replace=mp_replace).contiguous() + if weight_quantizer.q_int8: + tmp = transpose(tmp) + dst = mp_replace.copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ + transpose(tmp)), int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + +# Extending the maybe_copy function for when the q, k, and v are in separate parameters! +def maybe_copy_qkv(module, sd, weight_quantizer, mp_replace, dst_name, src_names, split_qkv=False): + if src_names[0] in sd: + q = sd[src_names[0]] + k = sd[src_names[1]] + v = sd[src_names[2]] + qkv_data = torch.cat((q, k, v), dim=0) + dst = getattr(module, dst_name) + if len(dst.shape) == 1: + if split_qkv: + dst = mp_replace.strided_copy(dst, qkv_data.contiguous(), num_splits=3) + else: + dst = mp_replace.copy(dst, qkv_data) + else: + if split_qkv: + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + ((transpose(qkv_data)).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) + else: + dst = mp_replace.copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + transpose(qkv_data)), int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + +# Extending the `maybe_copy` function for when mlp1 is in separate parameters for GeGLU +def maybe_copy_geglu(module, sd, weight_quantizer, mp_replace, dst_name, src_names): + if src_names[0] in sd: + reg_proj = sd[src_names[0]] + gate_proj = sd[src_names[1]] + + mlp1_data = torch.cat((reg_proj, gate_proj), dim=0) + dst = getattr(module, dst_name) + + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(mlp1_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + transpose(mlp1_data)), num_splits=2, int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + +def pack_lora_weights(p): + return [ + p.lora_right_weight, \ + p.lora_left_weight, \ + p.lora_scaling + ] + + +def maybe_get_lora(p): + if hasattr(p, 'lora_right_weight'): + lora_param = pack_lora_weights(p) + else: + lora_param = [] + return lora_param diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_module.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_module.py new file mode 100644 index 0000000000000000000000000000000000000000..ed94a5021feecaac6be0383b104892e80f489b31 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_module.py @@ -0,0 +1,723 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +import tqdm +import deepspeed +import deepspeed.ops.transformer as transformer_inference +from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention +from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock +from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig +from deepspeed.accelerator import get_accelerator +from .replace_policy import replace_policies, generic_policies +from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading +from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d +from deepspeed.module_inject.layers import is_autotp_training_mode +from deepspeed import comm as dist +from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads, set_tp_grain_size + +from .load_checkpoint import load_model_with_checkpoint +import time + +from .utils import policy_to_ds_container +import gc + + +def get_transformer_name(replaced_module): + from .containers import supported_models + from torch.nn import ModuleList + transformer_name = '' + for n, c in replaced_module.named_children(): + if c.__class__ in supported_models: + transformer_name += n + '.' + for name, child in c.named_children(): + if child.__class__ is ModuleList: + transformer_name += name + break + break + return transformer_name + + +class GroupQuantizer: + + def __init__(self, q_int8=True, group_size=1, num_bits=8, num_groups=0): + self.group_size = group_size + self.num_bits = num_bits + self.q_int8 = q_int8 + + self.num_groups = num_groups + + def quantize(self, inputs, qkv=True, count=1, parallel_dim=0): + if not self.q_int8 or not qkv: + inputs = torch.nn.Parameter(inputs, requires_grad=False) + inputs.scale = torch.empty(1) + return inputs + q_range = 2**self.num_bits + num_groups = self.num_groups if self.num_groups > 0 else inputs.shape[0] // self.group_size + inputs = inputs.to(get_accelerator().current_device_name()) + input_flat = inputs.reshape(num_groups, -1).contiguous() + input_min = torch.min(input_flat, dim=1, keepdim=True)[0].float() + input_max = torch.max(input_flat, dim=1, keepdim=True)[0].float() + scale = torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range) + input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1) + inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous() + out = torch.nn.Parameter(inputs_q, requires_grad=False) + inputs_split = inputs.split(inputs.shape[parallel_dim] // 2, dim=parallel_dim) + input_flat = [inputs_split[i].reshape(num_groups, -1).contiguous() for i in range(2)] + input_min = [torch.min(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)] + input_max = [torch.max(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)] + scale1 = [(torch.max(input_min[i].abs(), input_max[i].abs()) * 2.0 / (q_range)).squeeze().unsqueeze(0) + for i in range(2)] + + out.scale = torch.cat([scale.squeeze().unsqueeze(0), scale1[0], scale1[1]], dim=0).reshape(num_groups, + -1).contiguous() + return out + + +def _module_match(module): + for policy in generic_policies: + policy = policy() + if policy.match(module): + return policy + return None + + +def generic_injection(module, dtype=None, enable_cuda_graph=True): + + def replace_attn(child, policy): + policy_attn = policy.attention(child) + if policy_attn is None: + return child + if len(policy_attn) == 5: + qkvw, attn_ow, attn_ob, hidden_size, heads = policy_attn + else: + qw, kw, vw, attn_ow, attn_ob, hidden_size, heads = policy_attn + + config = transformer_inference.DeepSpeedInferenceConfig( + hidden_size=hidden_size, + heads=heads, + dtype=dtype, + triangular_masking=False, + max_out_tokens=4096, + ) + attn_module = DeepSpeedDiffusersAttention(config) + + def transpose(data): + data = data.contiguous() + data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) + data = data.reshape(data.shape[-1], data.shape[-2]) + data.to(get_accelerator().current_device_name()) + return data + + if len(policy_attn) == 5: + attn_module.attn_qkvw.data = transpose(qkvw.data) + else: + attn_module.attn_qkvw = None + attn_module.attn_qw.data = transpose(qw.data) + attn_module.attn_kw.data = transpose(kw.data) + attn_module.attn_vw.data = transpose(vw.data) + + attn_module.attn_qkvb = None + attn_module.attn_ow.data = transpose(attn_ow.data) + attn_module.attn_ob.data.copy_(attn_ob.data.to(get_accelerator().current_device_name())) + return attn_module + + def replace_attn_block(child, policy): + config = Diffusers2DTransformerConfig() + return DeepSpeedDiffusersTransformerBlock(child, config) + + if isinstance(module, torch.nn.Module): + pass + else: + if dtype not in [torch.float16, torch.half]: + raise ValueError("Generic injection only supported with FP16") + + try: + import diffusers + if hasattr(diffusers.models.attention, 'CrossAttention'): + cross_attention = diffusers.models.attention.CrossAttention + else: + cross_attention = diffusers.models.attention_processor.Attention + attention_block = diffusers.models.attention.BasicTransformerBlock + new_policies = { + cross_attention: replace_attn, + attention_block: replace_attn_block, + } + except ImportError: + new_policies = {} + + #replace_transformer_layer(None, + # module.text_encoder, + # training=False, + # replace_with_kernel_inject=True, + # triangular_masking=True, + # max_out_tokens=8192) + from ..model_implementations.transformers.clip_encoder import DSClipEncoder + cg_encoder = DSClipEncoder(module.text_encoder, enable_cuda_graph=enable_cuda_graph) + setattr(module, 'text_encoder', cg_encoder) + for name in module.__dict__.keys(): + sub_module = getattr(module, name) + policy = _module_match(sub_module) + + if policy is not None: + + def _replace_module(module, policy): + for name, child in module.named_children(): + _replace_module(child, policy) + if child.__class__ in new_policies: + replaced_module = new_policies[child.__class__](child, policy) + setattr(module, name, replaced_module) + + _replace_module(sub_module, policy) + new_module = policy.apply(sub_module, enable_cuda_graph=enable_cuda_graph) + print(f"**** found and replaced {name} w. {type(new_module)}") + setattr(module, name, new_module) + + +container_g = None + + +def replace_transformer_layer(orig_layer_impl, model, checkpoint_dict, config, model_config): + """ Replace bert-style transformer layers with DeepSpeed's transformer layer + Arguments: + orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for, + e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer + model (torch.nn.Module): user's nn.module representing their model + checkpoint_dict: Dictionary for checkpoint passed from the Inference Engine + config: top-level DS Inference config defined in inference/config.py + model_config: HuggingFace model config passed from the inference/engine.py + Returns: + Updated nn.module with replaced transformer layers + """ + # defining globals as internally defined functions inherit these everywhere + quantize = (config.dtype == torch.int8) + # todo: Refactor later. In future, let's minimize the style used above and use config.** instead + + linear_layer_setting = None + ''' + linear_layer_setting (tuple of modules) [Optional]: shows which two classes are used for linear layers and embedding layers + ''' + micro_batch_size = -1 + seed = -1 + local_rank = -1 + + mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group, + mp_size=config.tensor_parallel.tp_size) #, out_dim=0, in_dim=1) + + def replace_with_policy(child, policy_cls, triangular_masking, inference=False, layer_id=0): + policy = policy_cls(child, inference=inference) + if not policy.cuda_graph_supported: + # policy says cuda graph is not supported raise an error if set + assert not config.enable_cuda_graph, "cuda graph is not supported with this model, please disable" + + from deepspeed.moe.layer import MoE + moe = False + if hasattr(child, 'mlp') and isinstance(child.mlp, MoE): + num_experts = child.mlp.num_experts + moe = True + + # 1. Create a model-specific container object using the policy object. + _container = policy_to_ds_container(policy=policy, + config=config, + model_config=model_config, + layer_id=layer_id, + child=child) + _container.set_moe(moe) + + # 2. Set the tensor parallelism config + _container.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group) + + # 3. Initialize tensors + _container.initialize_tensors() + + # 4. deal with data types -- needs refactor to use dtype instead of fp16 + if config.dtype in [torch.float16, torch.bfloat16, torch.int8]: + _container.convert_to_required_dtype() + + # 5. Set the quantization config + quantizer = GroupQuantizer(q_int8=quantize) + _container.set_quantization_config(quantizer) + + # 6. create a DS Inference config object + _container.create_ds_model_config() + + # 7. use the config and create the module + _container.create_module() + + # 8. transpose the weights and bias if needed + _container.transpose() + + # 9. deal with tensor parallelism. + _container.apply_tensor_parallelism(mp_replace) + + # 10. copy the tensors from the model-specific container to the new module + _container.copy_data_to_new_module() + + # 11. set global for generic checkpoint loading + global container_g + + if container_g is None: + container_g = _container + + return _container.module + + def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None): + #mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group) + + # 1. Create AutoTP object + _autotp = AutoTP(module, all_reduce_linears, prefix, state_dict, linear_layer_setting, orig_layer_impl, + config.keep_module_on_host) + + # 2. Set the tensor parallelism config + _autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group) + + # 3. Try to get num_key_heads from model_config.num_key_value_heads + if hasattr(model_config, "vision_config"): + if "MllamaVisionEncoderLayer" in str(module): + num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config) + elif hasattr(model_config, "text_config"): + num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config) + else: + num_kv_heads = _autotp.get_model_num_kv_heads(model_config) + else: + num_kv_heads = _autotp.get_model_num_kv_heads(model_config) + + # 4. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division + set_num_kv_heads(num_kv_heads) + + # 4.1 Get n_embd + n_embd = None + multi_query_n_embd_names = ['n_embd', 'hidden_size'] + for name in multi_query_n_embd_names: + if hasattr(model_config, name): + n_embd = getattr(model_config, name) + if n_embd != None: + break + + # 4.2 set n_embd + set_n_embd(n_embd) + + # 4.3 set attention_heads + if hasattr(model_config, 'num_attention_heads'): + set_num_attention_heads(getattr(model_config, 'num_attention_heads')) + + # 4.4 set tp_grain_size + set_tp_grain_size(config.tensor_parallel.tp_grain_size) + + # 5. Set linear policies + _autotp.update_linear_policies() + + # 6. Replace modules + if "lm_head" in all_reduce_linears or "embed_out" in all_reduce_linears: + return _autotp._replace_last_linear_module(module) + return _autotp._replace_module(module) + + def replace_fn(child, _policy, layer_id=0, prefix="", state_dict=None): + training = False # todo: refactor this part to go in the config + if training: + # copy relevant state from child -> new module + new_module = replace_with_policy(child, _policy, config.triangular_masking) + + else: + # copy relevant state from child -> new module + if not is_autotp_training_mode() and config.replace_with_kernel_inject: + new_module = replace_with_policy(child, + _policy, + config.triangular_masking, + inference=True, + layer_id=layer_id) + else: + new_module = replace_wo_policy(child, _policy, prefix=prefix, state_dict=state_dict) + + return new_module + + def set_lm_head(module): + if is_autotp_training_mode(): + # we need to handle autoTP training mode separately. + return + + embedding_weight = None + for n, p in module.named_parameters(): + if "word_embeddings." in n or "embed_tokens." in n or "wte." in n: + embedding_weight = p + if embedding_weight is not None and hasattr(module, "lm_head") and hasattr( + module.lm_head, "weight") and module.lm_head.weight.is_meta: + module.lm_head.weight = embedding_weight + # enable tensor parallel for the last linear + if hasattr(module, "lm_head") and hasattr(module.lm_head, "weight") and isinstance( + module.lm_head, torch.nn.Linear): + module = replace_wo_policy(module, ("lm_head", ), 0, "lm_head") + elif hasattr(module, "embed_out") and hasattr(module.embed_out, "weight") and isinstance( + module.embed_out, torch.nn.Linear): + module = replace_wo_policy(module, ("embed_out", ), 0, "embed_out") + elif hasattr(module, "language_model") and hasattr(module.language_model, "lm_head"): + module = replace_wo_policy(module.language_model, ("lm_head", ), 0, "lm_head") + return module + + def conv2d_parallel_shard_weights(model, rank, world_size): + # add conv policy + shard_oc_name = ["conv1"] + shard_ic_name = ["conv2"] + for name, sub_m in model.named_children(): + for l_name, l_sub_m in sub_m.named_children(): + if l_name in shard_oc_name: + TPConv2d = TensorParallelOcShardConv2d( + l_sub_m, + rank, + world_size, + ) + setattr(sub_m, l_name, TPConv2d) + if l_name in shard_ic_name: + TPConv2d = TensorParallelIcShardConv2d( + l_sub_m, + rank, + world_size, + ) + setattr(sub_m, l_name, TPConv2d) + conv2d_parallel_shard_weights(sub_m, rank, world_size) + + if checkpoint_dict is not None and not config.replace_with_kernel_inject: + # AutoTP shard loading + checkpoint = checkpoint_dict["checkpoints"] + pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards") + for i in range(len(checkpoint)): + checkpoint_file = os.path.join(config.base_dir, checkpoint[i]) + replaced_module = replace_module(model=model, + orig_class=orig_layer_impl, + replace_fn=replace_fn, + _replace_policy=config.injection_policy_tuple, + checkpoint=checkpoint_file) + pbar.update(1) + gc.collect() + # conv2d tp module replace + # Now is for yuan model. Add model list and conv policy to decide whether to replace conv. + if 'Yuan' in str(replaced_module): + conv2d_parallel_shard_weights(replaced_module, dist.get_rank(), dist.get_world_size()) + else: + replaced_module = replace_module(model=model, + orig_class=orig_layer_impl, + replace_fn=replace_fn, + _replace_policy=config.injection_policy_tuple) + # AutoTP default set lm_head tp + if not config.replace_with_kernel_inject: + replaced_module = set_lm_head(replaced_module) + + quantizer = GroupQuantizer(q_int8=quantize) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + if checkpoint_dict is not None and config.replace_with_kernel_inject: + assert container_g.ckpt_load_enabled, \ + f"Meta Tensor checkpoint loading not supported in {container_g.__class__.__name__} container" + start_time = time.time() + checkpoint = checkpoint_dict['checkpoints'] + ckpt_list = checkpoint["tp"] if type(checkpoint) is dict else checkpoint + ckpt_type = checkpoint_dict.get('parallelization', 'pp') + ckpt_mp_size = checkpoint_dict.get('tp_size', len(ckpt_list)) + ckpt_mp_size = checkpoint_dict.get('mp_size', ckpt_mp_size) + base_dir1 = checkpoint_dict.get('base_dir', config.base_dir) + + if ckpt_type == 'pp' and type(checkpoint) is list: + pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards") + + for i in range(len(checkpoint)): + sd = [torch.load(os.path.join(base_dir1, checkpoint[i]), map_location='cpu', weights_only=False)] + load_model_with_checkpoint(replaced_module, + sd, + mp_replace, + ckpt_type, + ckpt_mp_size, + quantizer, + container=container_g) + pbar.update(1) + else: + num_checkpoints = len(ckpt_list) // ckpt_mp_size + tp_split_size = (world_size / ckpt_mp_size) + sd_offset = int(rank / tp_split_size) + sd_count = int((rank + max(1, tp_split_size)) / tp_split_size) - sd_offset + pbar = tqdm.tqdm(total=num_checkpoints, desc=f"Loading {num_checkpoints} checkpoint shards") + for i in range(num_checkpoints): + pbar.update(1) + ckpt_index = i * ckpt_mp_size + sd_offset + ckpt_files = [ + os.path.join(base_dir1, ckpt_list[ckpt_index + j]) if base_dir1 else ckpt_list[ckpt_index + j] + for j in range(sd_count) + ] + sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False) for ckpt_file in ckpt_files] + load_model_with_checkpoint(replaced_module, + sds, + mp_replace, + ckpt_type, + ckpt_mp_size, + quantizer, + int(rank % tp_split_size), + container=container_g) + sds = [None for _ in sds] + gc.collect() + + if "non_tp" in checkpoint: + pbar = tqdm.tqdm(total=len(checkpoint["non_tp"]), + desc=f"Loading {len(checkpoint['non_tp'])} checkpoint shards") + + for i in range(len(checkpoint["non_tp"])): + pbar.update(1) + ckpt_file = os.path.join(base_dir1, + checkpoint["non_tp"][i]) if base_dir1 else checkpoint["non_tp"][i] + sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False)] + load_model_with_checkpoint(replaced_module, + sds, + mp_replace, + ckpt_type, + ckpt_mp_size, + quantizer, + int(rank % tp_split_size), + container=container_g) + sds = [None for _ in sds] + gc.collect() + set_lm_head(replaced_module) + print(f"checkpoint loading time at rank {rank}: {time.time()-start_time} sec") + + if not is_autotp_training_mode() and config.save_mp_checkpoint_path is not None: + from collections import OrderedDict + import json + num_partitions = 8 + + if checkpoint_dict is None: + ckpt_name = "ds_model" + try: + from transformers.models.bloom.modeling_bloom import BloomForCausalLM + if isinstance(model, BloomForCausalLM): + ckpt_name = "bloom" + except ImportError: + ckpt_name = "ds_model" + else: + ckpt_name = checkpoint_dict['type'] + if dist.is_initialized(): + dist.barrier() + transformer_name = get_transformer_name(replaced_module) + non_tp_ckpt_name = f'non-tp.pt' + ckpt_files = [non_tp_ckpt_name] + os.makedirs(config.save_mp_checkpoint_path, exist_ok=True) + + if not dist.is_initialized() or dist.get_rank() == 0: + print("Saving tp-sharded checkpoints") + torch.save( + OrderedDict({ + k: v + for k, v in dict(replaced_module.state_dict()).items() if transformer_name not in k + }), f'{config.save_mp_checkpoint_path}/{non_tp_ckpt_name}') + + dtype_reprs = { + torch.float32: 'float32', + torch.float16: 'float16', + torch.int8: 'int8', + torch.bfloat16: 'bfloat16' + } + + ckpt_config = json.dumps({ + 'type': ckpt_name, + 'base_dir': f'{config.save_mp_checkpoint_path}', + 'checkpoints': { + "non_tp": ckpt_files, + "tp": [f'tp_{r:0>2d}_{m:0>2d}.pt' for m in range(num_partitions) for r in range(world_size)] + }, + 'version': 1.0, + 'parallelization': 'tp', + 'tp_size': world_size, + 'dtype': dtype_reprs[config.dtype] + }) + with open(f"{config.save_mp_checkpoint_path}/ds_inference_config.json", "w") as cfg: + cfg.write(ckpt_config) + + rep_sd = replaced_module.state_dict() + for n, p in replaced_module.named_parameters(): + if hasattr(p, 'scale'): + rep_sd[n] = [p, p.scale] + keys = list(rep_sd.keys()) + partition_size = (len(keys) // num_partitions + 1) + for m in range(num_partitions): + torch.save( + OrderedDict({ + k: [rep_sd[k], rep_sd[k].scale] if hasattr(rep_sd[k], 'scale') else rep_sd[k] + for k in keys[m * partition_size:(m + 1) * partition_size] if transformer_name in k + }), f'{config.save_mp_checkpoint_path}/tp_{rank:0>2d}_{m:0>2d}.pt') + + return replaced_module + + +def revert_transformer_layer(orig_layer_impl, model, config, preln=False): + """ Revert DeepSpeed's transformer layer back to original bert-style transformer layer + Arguments: + orig_layer_impl (torch.nn.Module): the original transformer layer implementation that was replaced, + e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer + model (torch.nn.Module): user's nn.module representing their model + config (dict): model config containing hidden size, attention heads, etc. + Returns: + Updated nn.module with original bert-style transformer layers + """ + + def replace_fn(child, _replace_policy, layer_id): + #from turing.nvidia_modelingpreln import BertLayer + orig_module = orig_layer_impl(config) + + # copy relevant state from child -> original module + qkvw = child.attn_qkvw.data + qkvb = child.attn_qkvb.data + + qw, kw, vw = torch.chunk(qkvw, 3, axis=0) + qb, kb, vb = torch.chunk(qkvb, 3, axis=0) + + orig_module.attention.self.query.weight.data = qw + orig_module.attention.self.query.bias.data = qb + orig_module.attention.self.key.weight.data = kw + orig_module.attention.self.key.bias.data = kb + orig_module.attention.self.value.weight.data = vw + orig_module.attention.self.value.bias.data = vb + + orig_module.attention.output.dense.weight.data = child.attn_ow.data + orig_module.attention.output.dense.bias.data = child.attn_ob.data + + attn_ln_w = child.attn_nw.data + attn_ln_b = child.attn_nb.data + if preln: + orig_module.PostAttentionLayerNorm.weight.data = attn_ln_w + orig_module.PostAttentionLayerNorm.bias.data = attn_ln_b + else: + orig_module.attention.output.LayerNorm.weight.data = attn_ln_w + orig_module.attention.output.LayerNorm.bias.data = attn_ln_b + + inter_ff_w = child.inter_w.data + inter_ff_b = child.inter_b.data + if preln: + orig_module.intermediate.dense_act.weight.data = inter_ff_w + orig_module.intermediate.dense_act.bias.data = inter_ff_b + else: + orig_module.intermediate.dense.weight.data = inter_ff_w + orig_module.intermediate.dense.bias.data = inter_ff_b + + orig_module.output.dense.weight.data = child.output_w.data + orig_module.output.dense.bias.data = child.output_b.data + + transformer_ln_w = child.norm_w.data + transformer_ln_b = child.norm_b.data + if preln: + orig_module.PreAttentionLayerNorm.weight.data = transformer_ln_w + orig_module.PreAttentionLayerNorm.bias.data = transformer_ln_b + else: + orig_module.output.LayerNorm.weight.data = transformer_ln_w + orig_module.output.LayerNorm.bias.data = transformer_ln_b + return orig_module + + return replace_module(model=model, + orig_class=deepspeed.DeepSpeedTransformerLayer, + replace_fn=replace_fn, + _replace_policy=None) + + +def replace_module(model, orig_class, replace_fn, _replace_policy, checkpoint=None): + """ Scan the model for instances of ``orig_clas:`` to replace using ``replace_fn``. + Arguments: + model (torch.nn.Module): the model to augment + orig_class (torch.nn.Module): the module to search for + replace_fn (method): a method to convert instances of ``orig_class`` to the + desired type and return a new instance. + Returns: + A modified ``model``. + """ + sd = None + if checkpoint is not None: + if checkpoint.endswith(".safetensors"): + from safetensors.torch import load_file + sd = load_file(checkpoint) + else: + sd = torch.load(checkpoint, map_location='cpu', weights_only=False) + + policy = {} + if orig_class is not None: + policy.update({orig_class: (replace_fn, _replace_policy)}) + else: + for plcy in replace_policies: + # instantiate a throw-away policy in order to populate the _orig_layer_class + _ = plcy(None) + if isinstance(plcy._orig_layer_class, list): + for orig_layer_class in plcy._orig_layer_class: + policy.update({orig_layer_class: (replace_fn, plcy)}) + elif plcy._orig_layer_class is not None: + policy.update({plcy._orig_layer_class: (replace_fn, plcy)}) + assert len(policy.items()) > 0,\ + "No default policy found! Please specify your policy injection_policy (like {BertLayer:HFBEertLayerPolicy})." +\ + "You can find some samples here: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/replace_policy.py" + + replaced_module, _ = _replace_module(model, policy, state_dict=sd) + return replaced_module + + +from ..pipe import PipelineModule + +import re + + +def skip_level_0_prefix(model, state_dict): + model = str(model) + key = re.search(r": (.*?)Model", model) + if key is None: + key = re.search(r": (.*?)Stack", model) + if key is None: + key = re.match(r"(.*?)Model", model) + # if keys start with 'model.', don't skip level 0 prefix + if state_dict is not None: + for item in state_dict.keys(): + if re.match("^model[.]", item): + return False + if key is not None and key.group(1).lower() in ["bloom", "opt"]: + return True + return False + + +def _replace_module(model, policies, prefix='', layer_id=0, level_id=0, state_dict=None): + """ Traverse model's children recursively and apply any transformations in ``policies``. + Arguments: + model (torch.nn.Module): model to augment + policies (dict): Mapping of source class to replacement function. + Returns: + Modified ``model``. + """ + for name, child in model.named_children(): + if child.__class__ in policies: + replaced_module = policies[child.__class__][0](child, + policies[child.__class__][-1], + layer_id, + prefix=prefix + name, + state_dict=state_dict) + setattr(model, name, replaced_module) + if isinstance(model, PipelineModule): + assert hasattr(model, 'forward_funcs'),\ + "we require pipe-module to have the list of fwd_functions" + model.forward_funcs[model.fwd_map[name]] = replaced_module + layer_id += 1 + else: + checking_key = prefix + name + '.' + if Loading.is_load_module(child) and state_dict is not None: + if any(checking_key in item for item in state_dict): + Loading.load( + child, + state_dict, + checking_key, + ) + else: + continue + if len(child._buffers) != 0 and state_dict is not None: + Loading.load_buffer(child, state_dict, checking_key) + _, layer_id = _replace_module(child, + policies, + prefix if level_id == 0 and skip_level_0_prefix(model, state_dict) else \ + prefix + name + '.', + layer_id=layer_id, + level_id=level_id + 1, + state_dict=state_dict) + + # Add the reset_cache func to the model, so that it can be called in the beginning of text-generation. + model.reset_cache = transformer_inference.DeepSpeedTransformerInference.reset_cache + return model, layer_id diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_policy.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..2c06e31aaa41ae61b5091a7766b709fac2c534f5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_policy.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .containers import HFGPT2LayerPolicy +from .containers import HFBertLayerPolicy +from .containers import BLOOMLayerPolicy +from .containers import HFGPTJLayerPolicy +from .containers import HFGPTNEOLayerPolicy +from .containers import GPTNEOXLayerPolicy +from .containers import HFOPTLayerPolicy +from .containers import MegatronLayerPolicy +from .containers import HFDistilBertLayerPolicy +from .containers import HFCLIPLayerPolicy +from .containers import LLAMALayerPolicy +from .containers import UNetPolicy +from .containers import VAEPolicy +from .containers import LLAMA2LayerPolicy +from .containers import InternLMLayerPolicy + +# transformer-based policies +replace_policies = [ + HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy, + HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy, + LLAMALayerPolicy, LLAMA2LayerPolicy, InternLMLayerPolicy +] + +# non-transformer-based policies +generic_policies = [UNetPolicy, VAEPolicy] diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/tp_shard.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/tp_shard.py new file mode 100644 index 0000000000000000000000000000000000000000..ded262edcf61bbb39d5d3e0c1575430aa8d9b080 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/tp_shard.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed import comm as dist +global num_kv_heads + + +def set_num_kv_heads(num): + global num_kv_heads + num_kv_heads = num + + +def set_num_attention_heads(num): + global num_attention_heads + num_attention_heads = num + + +def set_n_embd(num): + global n_embd + n_embd = num + + +def set_tp_grain_size(num): + global tp_grain_size + tp_grain_size = num + + +def get_num_kv_heads(): + global num_kv_heads + if 'num_kv_heads' in globals(): + return num_kv_heads + return None + + +def get_num_attention_heads(): + global num_attention_heads + return num_attention_heads + + +def get_shard_size(total_size, mp_size, name=None, rank=None): + global num_kv_heads + last_linear = ["lm_head", "embed_out"] + # MoE MLP layer use near even division will get better perf. + moe_mlp_layer = ["gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"] + not_moe_mlp_layer = True + if name != None and any(s in str(name) for s in moe_mlp_layer): + not_moe_mlp_layer = False + # When we have num_kv_heads defined, uneven division is possible, otherwise enforce near even division + if rank == None: + rank = dist.get_rank() + if num_kv_heads != None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and str( + name) not in last_linear and not_moe_mlp_layer: + my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0) + return total_size * my_slices // num_kv_heads + else: + if total_size >= tp_grain_size: + grain_size = total_size // tp_grain_size + return (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size + else: + return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0) + + +def get_n_embd(): + global n_embd + return n_embd + + +def get_shard_size_list(total_size, mp_size, name=None): + shard_sizes = [] + for i in range(mp_size): + shard_sizes.append(get_shard_size(total_size, mp_size, name, i)) + return shard_sizes diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/utils.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..42822128f9e11c84660dcd57a2473753b3bc6642 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/module_inject/utils.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.utils import log_dist + + +# helper function to map between DS policies and DS containers +def policy_to_ds_container(**kwargs): + from .containers import HFGPT2LayerPolicy, DS_GPT2Container + from .containers import HFBertLayerPolicy, DS_BERTContainer + from .containers import BLOOMLayerPolicy, DS_BloomContainer + from .containers import HFGPTJLayerPolicy, DS_GPTJContainer + from .containers import HFGPTNEOLayerPolicy, DS_GPTNEOContainer + from .containers import GPTNEOXLayerPolicy, DS_GPTNEOXContainer + from .containers import HFOPTLayerPolicy, DS_OPTContainer + from .containers import MegatronLayerPolicy, DS_MegatronGPTContainer + from .containers import HFDistilBertLayerPolicy, DS_DistilBERTContainer + from .containers import LLAMALayerPolicy, DS_LLAMAContainer + from .containers import LLAMA2LayerPolicy, DS_LLAMA2Container + from .containers import InternLMLayerPolicy, DS_InternLMContainer + + policy_to_container = { + HFGPT2LayerPolicy: DS_GPT2Container, + HFBertLayerPolicy: DS_BERTContainer, + BLOOMLayerPolicy: DS_BloomContainer, + HFGPTJLayerPolicy: DS_GPTJContainer, + HFGPTNEOLayerPolicy: DS_GPTNEOContainer, + GPTNEOXLayerPolicy: DS_GPTNEOXContainer, + HFOPTLayerPolicy: DS_OPTContainer, + MegatronLayerPolicy: DS_MegatronGPTContainer, + HFDistilBertLayerPolicy: DS_DistilBERTContainer, + LLAMALayerPolicy: DS_LLAMAContainer, + LLAMA2LayerPolicy: DS_LLAMA2Container, + InternLMLayerPolicy: DS_InternLMContainer + } + + container = None + policy = kwargs['policy'] + assert policy is not None, "Policy cannot be None" + policy_type = type(policy) + + if policy_type not in policy_to_container: + log_dist(f"Policy type {policy_type} not supported", [0]) + else: + container = policy_to_container[policy_type](**kwargs) + + return container diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/moe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/moe/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d99cae3bfa6284be26c8634dc8be986b6d12c7b4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/experts.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/experts.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d87ce8b98d1e3481bf68b16f751ab04966f44283 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/experts.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/layer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f019848b9ea7c715d7444807aa41d9060e900cba Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/layer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/mappings.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/mappings.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2eab9b3b9c03d4b496b95104e8fced542a42f72 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/mappings.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..206b49b79e173fb78caf9fff8418f17b329a9353 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f49a2bb61dc48ea5552b25d04c884a0df96672a5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/experts.py b/venv/lib/python3.10/site-packages/deepspeed/moe/experts.py new file mode 100644 index 0000000000000000000000000000000000000000..0863221d7edf93770af126e3fa1a78be3c7f4289 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/moe/experts.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy +from typing import List, Optional + +import torch +from torch import nn + + +class Experts(nn.Module): + + def __init__(self, expert: nn.Module, num_local_experts: int = 1, expert_group_name: Optional[str] = None) -> None: + super(Experts, self).__init__() + + self.deepspeed_experts = nn.ModuleList([copy.deepcopy(expert) for _ in range(num_local_experts)]) + self.num_local_experts = num_local_experts + + # TODO: revisit allreduce for moe.gate... + for expert in self.deepspeed_experts: + # TODO: Create param groups to handle expert + data case (e.g. param.group = moe_group) + for param in expert.parameters(): + param.allreduce = False + param.group_name = expert_group_name + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + chunks = inputs.chunk(self.num_local_experts, dim=1) + expert_outputs: List[torch.Tensor] = [] + + for chunk, expert in zip(chunks, self.deepspeed_experts): + out = expert(chunk) + if isinstance(out, tuple): + out = out[0] # Ignore the bias term for now + expert_outputs += [out] + + return torch.cat(expert_outputs, dim=1) diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/layer.py b/venv/lib/python3.10/site-packages/deepspeed/moe/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..6777788ab885a618544ced9ed9448e08c5159751 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/moe/layer.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Optional, Tuple + +import torch +from torch import nn +from torch.nn import functional as F + +from deepspeed.utils import groups, log_dist +from .experts import Experts +from .sharded_moe import MOELayer, TopKGate + + +class MoE(nn.Module): + """Initialize an MoE layer. + + Arguments: + hidden_size (int): the hidden dimension of the model, importantly this is also the input and output dimension. + expert (nn.Module): the torch module that defines the expert (e.g., MLP, torch.linear). + num_experts (int, optional): default=1, the total number of experts per layer. + ep_size (int, optional): default=1, number of ranks in the expert parallel world or group. + k (int, optional): default=1, top-k gating value, only supports k=1 or k=2. + capacity_factor (float, optional): default=1.0, the capacity of the expert at training time. + eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time. + min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor. + use_residual (bool, optional): default=False, make this MoE layer a Residual MoE (https://arxiv.org/abs/2201.05596) layer. + noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample' or 'None'. + drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent to infinite capacity). + use_rts (bool, optional): default=True, whether to use Random Token Selection. + use_tutel (bool, optional): default=False, whether to use Tutel optimizations (if installed). + enable_expert_tensor_parallelism (bool, optional): default=False, whether to use tensor parallelism for experts + top2_2nd_expert_sampling (bool, optional): default=True, whether to perform sampling for 2nd expert + """ + + def __init__(self, + hidden_size: int, + expert: nn.Module, + num_experts: int = 1, + ep_size: int = 1, + k: int = 1, + capacity_factor: float = 1.0, + eval_capacity_factor: float = 1.0, + min_capacity: int = 4, + use_residual: bool = False, + noisy_gate_policy: Optional[str] = None, + drop_tokens: bool = True, + use_rts: bool = True, + use_tutel: bool = False, + enable_expert_tensor_parallelism: bool = False, + top2_2nd_expert_sampling: bool = True) -> None: + + super(MoE, self).__init__() + + self.use_residual = use_residual + self.enable_expert_tensor_parallelism = enable_expert_tensor_parallelism + assert num_experts % ep_size == 0, f"Number of experts ({num_experts}) should be divisible by expert parallel size ({ep_size})" + self.ep_size = ep_size + self.expert_group_name = f"ep_size_{self.ep_size}" + self.num_experts = num_experts + self.num_local_experts = num_experts // self.ep_size + + log_dist( + f'Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {self.ep_size}', + [0]) + + assert noisy_gate_policy is None or noisy_gate_policy in ['None', 'Jitter', 'RSample'], \ + 'Unsupported noisy_gate_policy: ' + noisy_gate_policy + + experts = Experts(expert, self.num_local_experts, self.expert_group_name) + self.deepspeed_moe = MOELayer(TopKGate(hidden_size, num_experts, k, capacity_factor, eval_capacity_factor, + min_capacity, noisy_gate_policy, drop_tokens, use_rts, None, + top2_2nd_expert_sampling), + experts, + self.expert_group_name, + self.ep_size, + self.num_local_experts, + use_tutel=use_tutel) + if self.use_residual: + self.mlp = expert + # coefficient is used for weighted sum of the output of expert and mlp + self.coefficient = nn.Linear(hidden_size, 2) + + def set_deepspeed_parallelism(self, use_data_before_expert_parallel_: bool = False) -> None: + self._create_process_groups(use_data_before_expert_parallel_=use_data_before_expert_parallel_) + + def _create_process_groups(self, use_data_before_expert_parallel_: bool = False) -> None: + # Create process group for a layer if needed + if self.expert_group_name not in groups._get_expert_parallel_group_dict(): + print(f"No existing process group found, creating a new group named: {self.expert_group_name}") + if (groups.mpu is None) or (not self.enable_expert_tensor_parallelism): + # Condition 1 - no groups.mpu means no tensor parallelism + # Condition 2 - disabling expert tensor parallelism on purpose + groups._create_expert_and_data_parallel( + self.ep_size, use_data_before_expert_parallel_=use_data_before_expert_parallel_) + else: + # expert tensor parallelism is enabled + groups._create_expert_data_and_model_parallel( + self.ep_size, mpu=groups.mpu, use_data_before_expert_parallel_=use_data_before_expert_parallel_) + # Set the group handle for the MOELayer (deepspeed_moe) object + self.deepspeed_moe._set_ep_group(groups._get_expert_parallel_group(self.expert_group_name)) + + def forward(self, + hidden_states: torch.Tensor, + used_token: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ MoE forward + + Arguments: + hidden_states (Tensor): input to the layer + used_token (Tensor, optional): default: None, mask only used tokens + + Returns: + A tuple including output, gate loss, and expert count. + + * output (Tensor): output of the model + + * l_aux (Tensor): gate loss value + + * exp_counts (Tensor): expert count + """ + output = self.deepspeed_moe(hidden_states, used_token) + if self.use_residual: + # Residual MoE + output_mlp = self.mlp(hidden_states) + if isinstance(output_mlp, tuple): + output_mlp = output_mlp[0] # Ignore the bias term for now + coef = self.coefficient(hidden_states) + coef = F.softmax(coef, dim=-1) + output = output * coef[..., 0:1] + output_mlp * coef[..., 1:] + return output, self.deepspeed_moe.l_aux, self.deepspeed_moe.exp_counts diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/mappings.py b/venv/lib/python3.10/site-packages/deepspeed/moe/mappings.py new file mode 100644 index 0000000000000000000000000000000000000000..e57f66b85193d86734c186a916baa6da033a90a1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/moe/mappings.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# The file has been adapted from the following Megatron-LM file: +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/mappings.py +# Git commit hash: 9dc3c42a84aa656f583703cf8b6b4f79f712b796 +# We retain the following copyright from the original files: + +# Copyright (c) 2020, 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. + +import torch +import deepspeed +from deepspeed.utils.bwc import (bwc_tensor_model_parallel_world_size, bwc_tensor_model_parallel_rank, + bwc_tensor_model_parallel_group) + + +def _gather_tokens(input_, dim=0): + """Gather tensors and concatenate them along a dimension""" + mpu = deepspeed.utils.groups.mpu + + input_ = input_.contiguous() + world_size = bwc_tensor_model_parallel_world_size(mpu) + if world_size == 1: + return input_ + + gather_buffer = torch.empty(world_size * input_.numel(), dtype=input_.dtype, device=input_.device) + deepspeed.comm.all_gather_into_tensor(gather_buffer, input_, group=bwc_tensor_model_parallel_group(mpu)) + if dim == 0: + shape = list(input_.size()) + shape[0] = shape[0] * world_size + output = gather_buffer.view(shape) + else: + tensor_list = [ + gather_buffer.narrow(0, + input_.numel() * i, input_.numel()).view_as(input_) for i in range(world_size) + ] + # Note: torch.cat already creates a contiguous tensor. + output = torch.cat(tensor_list, dim=dim).contiguous() + + return output + + +def _drop_tokens(input_, dim=0): + """Divide a tensor among the tensor parallel ranks""" + mpu = deepspeed.utils.groups.mpu + + total_chunks = bwc_tensor_model_parallel_world_size(mpu) + if total_chunks == 1: + return input_ + this_chunk = bwc_tensor_model_parallel_rank(mpu) + assert input_.shape[ + dim] % total_chunks == 0, f"input dimension {dim} ({input_.shape[dim]}) is not divisible by tensor parallel world size ({total_chunks})" + chunk_size = input_.shape[dim] // total_chunks + + return torch.narrow(input_, dim, this_chunk * chunk_size, chunk_size) + + +class _GatherTokens(torch.autograd.Function): + """All gather tokens among the tensor parallel ranks""" + + @staticmethod + def symbolic(graph, input_, dim): + return _gather_tokens(input_, dim) + + @staticmethod + def forward(ctx, input_, dim): + ctx.dim = dim + return _gather_tokens(input_, dim) + + @staticmethod + def backward(ctx, grad_output): + return _drop_tokens(grad_output, ctx.dim), None + + +class _DropTokens(torch.autograd.Function): + "Divide tokens equally among the tensor parallel ranks" + + @staticmethod + def symbolic(graph, input_, dim): + return _drop_tokens(input_, dim) + + @staticmethod + def forward(ctx, input_, dim): + ctx.dim = dim + return _drop_tokens(input_, dim) + + @staticmethod + def backward(ctx, input_): + return _gather_tokens(input_, ctx.dim), None + + +def gather_tokens(input_, dim=0): + mpu = deepspeed.utils.groups.mpu + if mpu is None or bwc_tensor_model_parallel_world_size(mpu) == 1: + # no tensor parallelism for non-experts + return input_ + return _GatherTokens.apply(input_, dim) + + +def drop_tokens(input_, dim=0): + mpu = deepspeed.utils.groups.mpu + if mpu is None or bwc_tensor_model_parallel_world_size(mpu) == 1: + # no tensor parallelism for non-experts + return input_ + return _DropTokens.apply(input_, dim) diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/sharded_moe.py b/venv/lib/python3.10/site-packages/deepspeed/moe/sharded_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..5aacd1a11a1ad12c8d125c55851c13f6c19df747 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/moe/sharded_moe.py @@ -0,0 +1,678 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +The file has been adapted from two fairscale files: + (1) https://github.com/facebookresearch/fairscale/blob/master/fairscale/nn/moe/moe_layer.py + (2) https://github.com/facebookresearch/fairscale/blob/master/fairscale/nn/moe/top2gate.py + Git commit hash: 34df606902a240567a0d898037ece55c2f1336cf + We retain the following license from the original files: +""" + +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +from deepspeed.utils.timer import SynchronizedWallClockTimer +from deepspeed.utils import logger +from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size +from typing import Callable, Dict, TYPE_CHECKING, Any, Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.nn import Module +import torch.nn.functional as F +from deepspeed.utils import groups +from .mappings import drop_tokens, gather_tokens + +if TYPE_CHECKING: + Base = Module[Tensor] +else: + Base = Module + +TOPK_GATE_TIMER = 'topk_gate' +MOE_TIMER = 'moe' +FIRST_ALLTOALL_TIMER = '1st_a2a' +SECOND_ALLTOALL_TIMER = '2nd_a2a' + +uniform_map: Dict[torch.device, Callable] = {} +gumbel_map: Dict[torch.device, Callable] = {} +exp_selection_uniform_map: Dict[torch.device, Callable] = {} + +try: + # To enable Tutel MoE optimizations: + # python3 -m pip install --user --upgrade git+https://github.com/deepspeedai/tutel@v0.1.x + from tutel import moe as tutel_moe + TUTEL_INSTALLED = True +except: + # Fail silently so we don't spam logs unnecessarily if user isn't using tutel + TUTEL_INSTALLED = False + pass + + +def multiplicative_jitter(x, device: torch.device, epsilon=1e-2): + """ + Modified from switch transformer paper. mesh transformers + Multiply values by a random number between 1-epsilon and 1+epsilon. + Makes models more resilient to rounding errors introduced by bfloat16. + This seems particularly important for logits. + Args: + x: a torch.tensor + device: torch.device + epsilon: a floating point value + Returns: + a jittered x. + """ + if epsilon == 0: + return x + uniform = uniform_map.get(device) + if uniform is None: + uniform = torch.distributions.uniform.Uniform(low=torch.tensor(1.0 - epsilon, device=device), + high=torch.tensor(1.0 + epsilon, + device=device)).rsample # type: ignore + uniform_map[device] = uniform + return x * uniform(x.shape) + + +def gumbel_rsample(shape: Tuple, device: torch.device) -> Tensor: + gumbel = gumbel_map.get(device) + if gumbel is None: + one = torch.tensor(1.0, device=device) + zero = torch.tensor(0.0, device=device) + gumbel = torch.distributions.gumbel.Gumbel(zero, one).rsample # type: ignore + gumbel_map[device] = gumbel + return gumbel(shape) + + +from deepspeed import comm as dist + +# einsum dimensions: (g)roup, (s)equence, (e)xpert, (m)odel, (c)apacity +# See https://arxiv.org/pdf/2006.16668.pdf for details. + + +# Based on https://github.com/pytorch/pytorch/pull/40762 +class _AllToAll(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor: # type: ignore + ctx.group = group + input = input.contiguous() + output = torch.empty_like(input) + dist.all_to_all_single(output, input, group=group) + return output + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor]: + return (None, _AllToAll.apply(ctx.group, *grad_output)) + + +# einsum rewrites are on par or more performant +# switch can be bubbled up in future +USE_EINSUM = True + + +# einsum dimensions: (g)roup, (s)equence, (e)xpert, (m)odel, (c)apacity +# See https://arxiv.org/pdf/2006.16668.pdf for details. +def einsum(rule, a, b): + if USE_EINSUM: + return torch.einsum(rule, a, b) + elif rule == 's,se->se': + return a.reshape(a.shape[0], -1) * b + elif rule == 'se,sc->sec': + return a.unsqueeze(2) * b.unsqueeze(1) + elif rule == 'se,se->s': + return torch.bmm(a.unsqueeze(1), b.unsqueeze(2)).reshape(-1) + elif rule == 'se,sec->sec': + return a.unsqueeze(2) * b + elif rule == 'sec,sm->ecm': + s = a.shape[0] + e = a.shape[1] + c = a.shape[2] + m = b.shape[1] + return torch.matmul(a.reshape(s, -1).t(), b).reshape(e, c, m) + elif rule == 'sec,ecm->sm': + return torch.matmul(a.reshape(a.shape[0], -1), b.reshape(-1, b.shape[-1])) + elif rule == 'ks,ksm->sm': + k = b.shape[0] + s = b.shape[1] + m = b.shape[2] + # [k, s] -> [s, k] -> [s, 1, k] + a = a.t().unsqueeze(1) + # [k,s,m] -> [k, sm] -> [sm, k] -> [s, m, k] + b = b.reshape(k, -1).t().reshape(s, m, k) + # bmm([s, 1, k], [s, m, k]^t) -> [s, m, 1] + return torch.bmm(a, b.transpose(1, 2)).squeeze(2) + else: + return torch.einsum(rule, a, b) + + +# The following functions are extracted and scripted +# because otherwise during a torch.jit.trace, the non-Tensor +# values used in the calculations get recorded as constants. +# torch.jit.script coerces them into Tensors and preserves +# their dynamic shapes. This enables ONNX export. +# We can't script the entire top1gating function because it +# includes stateful caching logic which is incompatible with ONNX. + + +@torch.jit.script +def _capacity(gates: Tensor, capacity_factor: Tensor, min_capacity: Tensor) -> Tensor: + # gates has shape of SE + num_tokens = gates.shape[0] + num_experts = gates.shape[1] + # to(torch.int64) works around a bug in torch.onnx.export: + # it should cast k to int64 when converting torch.topk but it doesn't. + capacity = torch.ceil((num_tokens / num_experts) * capacity_factor).to(torch.int64) + if capacity < min_capacity: + capacity = min_capacity.to(torch.int64) + return capacity + + +@torch.jit.script +def _top_idx(source, k): + return torch.topk(source, k=k, dim=0)[1] + + +@torch.jit.script +def _one_hot_to_float(x, num_classes): + return F.one_hot(x, num_classes=num_classes).float() + + +def top1gating(logits: Tensor, + capacity_factor: float, + min_capacity: int, + used_token: Tensor = None, + noisy_gate_policy: Optional[str] = None, + drop_tokens: bool = True, + use_rts: bool = True, + ep_group: Union[torch.distributed.ProcessGroup, None] = None, + use_tutel: bool = False) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Implements Top1Gating on logits.""" + if noisy_gate_policy == 'RSample': + logits_w_noise = logits + gumbel_rsample(logits.shape, device=logits.device) + # everything is in fp32 in this function + + gates = F.softmax(logits, dim=1) + capacity = _capacity(gates, torch.tensor(capacity_factor), torch.tensor(min_capacity)) + + # Create a mask for 1st's expert per token + # noisy gating + indices1_s = torch.argmax(logits_w_noise if noisy_gate_policy == 'RSample' else gates, dim=1) + num_experts = int(gates.shape[1]) + mask1 = F.one_hot(indices1_s, num_classes=num_experts) + + # mask only used tokens + if used_token is not None: + mask1 = einsum("s,se->se", used_token, mask1) + + # gating decisions + exp_counts = torch.sum(mask1, dim=0).detach().to(logits.device) + + # if we don't want to drop any tokens + if not drop_tokens: + new_capacity = torch.max(exp_counts).to(logits.device) + # Communicate across expert processes to pick the maximum capacity. + if ep_group is not None: + dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=ep_group) + if groups._get_expert_model_parallel_world_size() == 1: + # If the non-expert is tensor-parallel, we need to pad the capacity to 'tp'. + # This is since we are going to activate drop_tokens() to drop duplicate tokens. + tp = 1 if groups.mpu is None else bwc_tensor_model_parallel_world_size(mpu=groups.mpu) + new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype) + # Make sure the capacity value does not exceed the number of tokens. + capacity = min(new_capacity, torch.tensor(mask1.size(0)).to(new_capacity.device)) + + # Compute l_aux + me = torch.mean(gates, dim=0) + ce = torch.mean(mask1.float(), dim=0) + l_aux = torch.sum(me * ce) * num_experts + + # Random Token Selection + if use_rts: + uniform = exp_selection_uniform_map.get(logits.device) + if uniform is None: + uniform = torch.distributions.uniform.Uniform(low=torch.tensor(0.0, device=logits.device), + high=torch.tensor(1.0, device=logits.device)).rsample + exp_selection_uniform_map[logits.device] = uniform + + mask1_rand = mask1 * uniform(mask1.shape) + else: + mask1_rand = mask1 + + assert logits.shape[ + 0] >= min_capacity, "No. of tokens (batch-size) should be greater than min_capacity. Either set min_capacity to 0 or increase your batch size." + + top_idx = _top_idx(mask1_rand, capacity) + + new_mask1 = mask1 * torch.zeros_like(mask1).scatter_(0, top_idx, 1) + mask1 = new_mask1 + + if use_tutel: + # Tutel doesn't support index values masked with zero + # so we need to replace masked indices with -1 + indices_mask = mask1.sum(dim=1) * num_experts - 1 + indices1_s = torch.min(indices1_s, indices_mask) + + # Compute locations in capacity buffer + if use_tutel: + locations1 = tutel_moe.fast_cumsum_sub_one(mask1) + else: + locations1 = torch.cumsum(mask1, dim=0) - 1 + + if use_tutel: + gates1_s = (gates * mask1).sum(dim=1) + locations1_s = torch.sum(locations1 * mask1, dim=1) + return l_aux, capacity, num_experts, [ + indices1_s, + ], [ + locations1_s, + ], [ + gates1_s, + ], exp_counts + + # Store the capacity location for each token + locations1_s = torch.sum(locations1 * mask1, dim=1) + + # Normalize gate probabilities + mask1_float = mask1.float() + gates = gates * mask1_float + + locations1_sc = _one_hot_to_float(locations1_s, capacity) + combine_weights = einsum("se,sc->sec", gates, locations1_sc) + + dispatch_mask = combine_weights.bool() + + return l_aux, combine_weights, dispatch_mask, exp_counts + + +def top2gating(logits: Tensor, + capacity_factor: float, + min_capacity: int, + drop_tokens: bool = True, + ep_group: Union[torch.distributed.ProcessGroup, None] = None, + top2_2nd_expert_sampling: bool = True) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Implements Top2Gating on logits.""" + # everything is in fp32 in this function + gates = F.softmax(logits, dim=1) + + # Create a mask for 1st's expert per token + indices1_s = torch.argmax(gates, dim=1) + num_experts = int(gates.shape[1]) + mask1 = F.one_hot(indices1_s, num_classes=num_experts) + + if top2_2nd_expert_sampling: + # Create a mask for 2nd's expert per token using Gumbel-max trick + # https://timvieira.github.io/blog/post/2014/07/31/gumbel-max-trick/ + logits += gumbel_rsample(logits.shape, device=logits.device) + + # Replace top-expert with min value + logits_except1 = logits.masked_fill(mask1.bool(), float("-inf")) + indices2_s = torch.argmax(logits_except1, dim=1) + mask2 = F.one_hot(indices2_s, num_classes=num_experts) + + # Compute locations in capacity buffer + locations1 = torch.cumsum(mask1, dim=0) - 1 + locations2 = torch.cumsum(mask2, dim=0) - 1 + # Update 2nd's location by accounting for locations of 1st + locations2 += torch.sum(mask1, dim=0, keepdim=True) + + # Compute l_aux + me = torch.mean(gates, dim=0) + ce = torch.mean(mask1.float(), dim=0) + l_aux = torch.mean(me * ce) * num_experts * num_experts + + # gating decisions + exp_counts = torch.sum(mask1 + mask2, dim=0).detach().to(logits.device) + + if drop_tokens: + # Calculate configured capacity and remove locations outside capacity from mask + capacity = _capacity(gates, torch.tensor(capacity_factor * 2), torch.tensor(min_capacity)) + mask1 *= torch.lt(locations1, capacity) + mask2 *= torch.lt(locations2, capacity) + else: + # Do not drop tokens - set capacity according to current expert assignments + new_capacity = torch.max(exp_counts) + if ep_group is not None: + dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=ep_group) + if groups._get_expert_model_parallel_world_size() == 1: + # If the non-expert is tensor-parallel, we need to pad the capacity to 'tp'. + # This is since we are going to activate drop_tokens() to drop duplicate tokens. + tp = 1 if groups.mpu is None else bwc_tensor_model_parallel_world_size(mpu=groups.mpu) + new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype) + capacity = new_capacity + + # Store the capacity location for each token + locations1_s = torch.sum(locations1 * mask1, dim=1) + locations2_s = torch.sum(locations2 * mask2, dim=1) + + # Normalize gate probabilities + mask1_float = mask1.float() + mask2_float = mask2.float() + gates1_s = einsum("se,se->s", gates, mask1_float) + gates2_s = einsum("se,se->s", gates, mask2_float) + denom_s = gates1_s + gates2_s + # Avoid divide-by-zero + denom_s = torch.clamp(denom_s, min=torch.finfo(denom_s.dtype).eps) + gates1_s /= denom_s + gates2_s /= denom_s + + # Calculate combine_weights and dispatch_mask + gates1 = einsum("s,se->se", gates1_s, mask1_float) + gates2 = einsum("s,se->se", gates2_s, mask2_float) + locations1_sc = _one_hot_to_float(locations1_s, capacity) + locations2_sc = _one_hot_to_float(locations2_s, capacity) + combine1_sec = einsum("se,sc->sec", gates1, locations1_sc) + combine2_sec = einsum("se,sc->sec", gates2, locations2_sc) + combine_weights = combine1_sec + combine2_sec + dispatch_mask = combine_weights.bool() + + return l_aux, combine_weights, dispatch_mask, exp_counts + + +def topkgating( + logits: Tensor, + k: int, + capacity_factor: float, + min_capacity: int, + drop_tokens: bool = True, + ep_group: Union[torch.distributed.ProcessGroup, None] = None, + drop_policy: str = "probs", +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Implements TopKGating on logits.""" + + # everything is in fp32 in this function + # get topk gates + top_gate, top_idx = torch.topk(logits, k=k, dim=1) + # gating decisions + gates = F.softmax(logits, dim=1) + num_experts = int(gates.shape[1]) + + # get topk mask + topk_masked_gates = torch.zeros_like(logits).scatter(1, top_idx, top_gate) + + mask = torch.zeros_like(gates, dtype=torch.bool).scatter_(1, top_idx, 1) + + exp_counts = torch.sum(mask, dim=0).detach().to(logits.device) + + # Compute l_aux + me = torch.mean(gates, dim=0) + ce = torch.mean(mask.float(), dim=0) + l_aux = torch.mean(me * ce) * num_experts * num_experts / k + + if drop_tokens: + # Calculate configured capacity and remove locations outside capacity from mask + capacity = _capacity(gates, torch.tensor(capacity_factor * k), torch.tensor(min_capacity)) + # update mask and locations by capacity + + if drop_policy == 'probs': + capacity_probs, capacity_indices = torch.topk(topk_masked_gates, k=capacity, dim=0, sorted=False) + capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1) + mask = torch.logical_and(mask, capacity_mask) + locations = torch.cumsum(mask, dim=0) - 1 + + elif drop_policy == "position": + locations = torch.cumsum(mask, dim=0) - 1 + mask *= torch.lt(locations, capacity) + else: + raise ValueError(f"Invalid drop_policy: {drop_policy}") + + else: + # Do not drop tokens - set capacity according to current expert assignments + new_capacity = torch.max(exp_counts) + if ep_group is not None: + dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=ep_group) + if groups._get_expert_model_parallel_world_size() == 1: + # If the non-expert is tensor-parallel, we need to pad the capacity to 'tp'. + # This is since we are going to activate drop_tokens() to drop duplicate tokens. + tp = 1 if groups.mpu is None else bwc_tensor_model_parallel_world_size(mpu=groups.mpu) + new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype) + capacity = new_capacity + locations = torch.cumsum(mask, dim=0) - 1 + + # normalize gates + gates_masked = gates * mask + gates_s = torch.sum(gates_masked, dim=-1, keepdim=True) + denom_s = torch.clamp(gates_s, min=torch.finfo(gates_masked.dtype).eps) + gates_masked = gates_masked / denom_s + + # dispatch_mask + locations_sc = _one_hot_to_float((locations * mask), capacity) + + combine_weights = torch.einsum("se,sec->sec", gates_masked, locations_sc) + + dispatch_mask = combine_weights.bool() + + return l_aux, combine_weights, dispatch_mask, exp_counts + + +class TopKGate(Module): + """Gate module which implements Top2Gating as described in Gshard_. + :: + + gate = TopKGate(model_dim, num_experts) + l_aux, combine_weights, dispatch_mask = gate(input) + + .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf + + Args: + model_dim (int): + size of model embedding dimension + num_experts (int): + number of experts in model + """ + + wg: torch.nn.Linear + + def __init__(self, + model_dim: int, + num_experts: int, + k: int = 1, + capacity_factor: float = 1.0, + eval_capacity_factor: float = 1.0, + min_capacity: int = 8, + noisy_gate_policy: Optional[str] = None, + drop_tokens: bool = True, + use_rts: bool = True, + ep_group: Union[torch.distributed.ProcessGroup, None] = None, + top2_2nd_expert_sampling: bool = True) -> None: + super().__init__() + + self.wg = torch.nn.Linear(model_dim, num_experts, bias=False) + self.ep_group = ep_group + self.k = k + self.capacity_factor = capacity_factor + self.eval_capacity_factor = eval_capacity_factor + self.min_capacity = min_capacity + self.noisy_gate_policy = noisy_gate_policy + self.timers = SynchronizedWallClockTimer() + self.wall_clock_breakdown = False + self.gate_time = 0.0 + self.drop_tokens = drop_tokens + self.use_rts = use_rts + self.top2_2nd_expert_sampling = top2_2nd_expert_sampling + + def _set_ep_group(self, ep_group): + assert self.ep_group is None, f'Attempting to override an existing ep_group' + self.ep_group = ep_group + + def forward(self, + input: torch.Tensor, + used_token: torch.Tensor = None, + use_tutel: bool = False) -> Tuple[Tensor, Tensor, Tensor]: # type: ignore + + if self.wall_clock_breakdown: + self.timers(TOPK_GATE_TIMER).start() + + input_fp32 = input.float() + # input jittering + if self.noisy_gate_policy == 'Jitter' and self.training: + input_fp32 = multiplicative_jitter(input_fp32, device=input.device) + logits = torch.nn.functional.linear(input_fp32, weight=self.wg.weight.float(), bias=None) + + if self.k == 1: + gate_output = top1gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor, + self.min_capacity, used_token, self.noisy_gate_policy if self.training else None, + self.drop_tokens, self.use_rts, self.ep_group, use_tutel) + + elif self.k == 2: + gate_output = top2gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor, + self.min_capacity, self.drop_tokens, self.ep_group, self.top2_2nd_expert_sampling) + else: + gate_output = topkgating(logits, self.k, + self.capacity_factor if self.training else self.eval_capacity_factor, + self.min_capacity, self.drop_tokens, self.ep_group) + + if self.wall_clock_breakdown: + self.timers(TOPK_GATE_TIMER).stop() + self.gate_time = self.timers(TOPK_GATE_TIMER).elapsed(reset=False) + + return gate_output + + +class MOELayer(Base): + """MOELayer module which implements MixtureOfExperts as described in Gshard_. + :: + + gate = TopKGate(model_dim, num_experts) + moe = MOELayer(gate, expert) + output = moe(input) + l_aux = moe.l_aux + + .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf + + Args: + gate (torch.nn.Module): + gate network + expert (torch.nn.Module): + expert network + """ + + def __init__(self, + gate: Module, + experts: Module, + ep_group_name, + ep_size, + num_local_experts: int, + use_tutel: bool = False) -> None: + super().__init__() + self.gate = gate + self.experts = experts + self.ep_group = None + self.ep_size = ep_size + self.ep_group_name = ep_group_name + self.num_local_experts = num_local_experts + self.time_falltoall = 0.0 + self.time_salltoall = 0.0 + self.time_moe = 0.0 + self.timers = SynchronizedWallClockTimer() + self.wall_clock_breakdown = False + + self.use_tutel = use_tutel and TUTEL_INSTALLED and gate.k == 1 + + if self.use_tutel: + logger.info('Using Tutel optimizations.') + elif use_tutel and not TUTEL_INSTALLED: + logger.warning("Tutel optimization requested but not installed. " + "Proceeding without Tutel.") + elif use_tutel and TUTEL_INSTALLED and gate.k != 1: + logger.warning("To enable Tutel optimization, use top-1 instead of top-2 gate. " + "Proceeding without Tutel.") + + def _set_ep_group(self, ep_group): + self.ep_group = ep_group + self.gate._set_ep_group(ep_group) + + def forward(self, *input: Tensor, **kwargs: Any) -> Tensor: + + if self.wall_clock_breakdown: + self.timers(MOE_TIMER).start() + + # Implement Algorithm 2 from GShard paper. + d_model = input[0].shape[-1] + + # Initial implementation -> Reshape into S tokens by dropping sequence dimension. + # Reshape into G groups so that each group can distribute tokens equally + # group_size = kwargs['group_size'] if 'group_size' in kwargs.keys() else 1 + reshaped_input = input[0].reshape(-1, d_model) + + if self.use_tutel: + self.l_aux, C, E, indices_, locations_, gates_, self.exp_counts = self.gate(reshaped_input, input[1], True) + S, M = reshaped_input.size(0), reshaped_input.size(1) + + if not hasattr(self, '_tutel_dispatcher'): + self._tutel_dispatcher = tutel_moe.fast_dispatcher(E, C, M, dispatch_dtype=reshaped_input.dtype) + self._tutel_dispatcher.update(indices_, locations_, gates_, capacity=C) + dispatched_input = self._tutel_dispatcher.encode(reshaped_input) + else: + self.l_aux, combine_weights, dispatch_mask, self.exp_counts = self.gate(reshaped_input, input[1]) + dispatched_input = einsum("sec,sm->ecm", dispatch_mask.type_as(input[0]), reshaped_input) + + if self.wall_clock_breakdown: + self.timers(FIRST_ALLTOALL_TIMER).start() + + tensor_model_world_size = bwc_tensor_model_parallel_world_size(groups.mpu) + if tensor_model_world_size > 1: + # If the non-expert is tensor-parallel, + # Whether expert is tensor-parallel or not , it will create + # duplicate tokens on the tensor-parallel ranks. + # drop duplicate tokens also doubles up as a communication + # optimization as we are reducing the all-to-all communication volume. + # 1: for not tensor-parallel expert,drop duplicate tokens to ensure + # both correctness and reduce all-to-all communication. + # 2: for tensor-parallel expert,drop duplicate tokens to reduce all-to-all + # communication volume,before expert execution, it is necessary to perform + # an allgather to ensure correctness, + dispatched_input = drop_tokens(dispatched_input, dim=1) + + dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input) + + if self.wall_clock_breakdown: + self.timers(FIRST_ALLTOALL_TIMER).stop() + self.time_falltoall = self.timers(FIRST_ALLTOALL_TIMER).elapsed(reset=False) + + if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1: + # if both expert and non-expert are tensor-parallel + # the dropped duplicate tokens need to be gathered on each + # tensor parallel rank again to ensure correctness + dispatched_input = gather_tokens(dispatched_input, dim=1) + + # Re-shape after all-to-all: ecm -> gecm + dispatched_input = dispatched_input.reshape(self.ep_size, self.num_local_experts, -1, d_model) + expert_output = self.experts(dispatched_input) + # Re-shape before drop_tokens: gecm -> ecm + expert_output = expert_output.reshape(self.ep_size * self.num_local_experts, -1, d_model) + if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1: + # if both expert and non-expert are tensor-parallel + # drop duplicate tokens to ensure both correctness + # and reduce all-to-all communication. + expert_output = drop_tokens(expert_output, dim=1) + + if self.wall_clock_breakdown: + self.timers(SECOND_ALLTOALL_TIMER).start() + + expert_output = _AllToAll.apply(self.ep_group, expert_output) + + if self.wall_clock_breakdown: + self.timers(SECOND_ALLTOALL_TIMER).stop() + self.time_salltoall = self.timers(SECOND_ALLTOALL_TIMER).elapsed(reset=False) + + if tensor_model_world_size > 1: + # the dropped duplicate tokens need to be gathered on each + # tensor parallel rank again for the tensor-parallel + # non-expert of the next layer. + expert_output = gather_tokens(expert_output, dim=1) + + if self.use_tutel: + combined_output = self._tutel_dispatcher.decode(expert_output.view(E * C, M)) + else: + combined_output = einsum("sec,ecm->sm", combine_weights.type_as(input[0]), expert_output) + + a = combined_output.reshape(input[0].shape) + + if self.wall_clock_breakdown: + self.timers(MOE_TIMER).stop() + self.time_moe = self.timers(MOE_TIMER).elapsed(reset=False) + + return a diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/utils.py b/venv/lib/python3.10/site-packages/deepspeed/moe/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..20866378efac72c96a7e2f56bdf97e5b7b4effb9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/moe/utils.py @@ -0,0 +1,182 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from collections import defaultdict +from typing import Any, Dict, List, Set, Tuple, Union, cast + +import torch +from torch import nn + +from .layer import MoE + + +def has_moe_layers(m: nn.Module) -> Tuple[bool, int]: + has_moe = False + num_experts = 0 + + for module in m.modules(): + if isinstance(module, MoE): + has_moe = True + num_experts = module.num_experts + break + return has_moe, num_experts + + +def is_moe_param(param: torch.Tensor) -> bool: + if hasattr(param, "allreduce") and not param.allreduce: + return True + return False + + +def split_params_into_shared_and_expert_params( + params: List[torch.nn.Parameter]) -> Tuple[List[torch.nn.Parameter], List[torch.nn.Parameter]]: + shared_params: List[nn.Parameter] = [] + expert_params: List[nn.Parameter] = [] + + for p in params: + if is_moe_param(p): + expert_params.append(p) + else: + shared_params.append(p) + return shared_params, expert_params + + +def split_params_grads_into_shared_and_expert_params( + group: List[torch.nn.Parameter]) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: + """Split grad of parameters into grads of non-expert params + and grads of expert params. This is useful while computing + grad-norms for clipping and overflow detection + + group (List[torch.nn.Parameter]): + Args: + The group of parameters to split + + Returns: + Tuple[List[torch.Tensor], List[torch.Tensor]]: + list of gradients for non MoE params, list of gradients of MoE params + """ + expert_grads: List[torch.Tensor] = [] + shared_grads: List[torch.Tensor] = [] + + for p in group: + if p.grad is not None: + if is_moe_param(p): + expert_grads.append(p.grad.to(p.dtype)) + else: + shared_grads.append(p.grad.to(p.dtype)) + return shared_grads, expert_grads + + +def split_params_into_different_moe_groups_for_optimizer( + param_groups: Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[str, Any]]], + max_group_size: Union[int, float] = 178956971) -> List[Dict[str, Any]]: + """Split parameters into different MoE groups for optimizer + + Args: + param_groups (Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[str, Any]]]) + The list of parameter groups to split + + Returns: + List[Dict[str, Any]]: + list of MoE/non-MoE groups for optimizer + """ + if isinstance(param_groups, tuple): + param_groups = list(param_groups) # Tuple cannot be modified + elif isinstance(param_groups, dict): + param_groups = [param_groups] + elif not isinstance(param_groups, list): + raise ValueError(f"Unknown param group type of {type(param_groups)}") + + # gather all data parallel group names + data_parallel_group_names: Set[str] = set() + for param_group in param_groups: + for param in cast(List[nn.Parameter], param_group["params"]): + if is_moe_param(param): + data_parallel_group_names.add(param.group_name) + + # Create the param MoE groups, leave param assign to next step + group_moe: Dict[str, Dict[str, Dict[str, Any]]] = defaultdict(lambda: defaultdict(dict)) + for param_group in param_groups: + for key in data_parallel_group_names: + group_moe[param_group['name']][key] = { + **param_group, + 'name': key, + 'moe': True, + 'params': [], + } + + # Assign param + for param_group in param_groups: + new_params: List[nn.Parameter] = [] + + for param in cast(List[nn.Parameter], param_group['params']): + if is_moe_param(param): + group_moe[param_group['name']][param.group_name]['params'].append(param) + else: + new_params.append(param) + param_group['params'] = new_params + + # Flatten the moe groups + if max_group_size is not None: + for moe_group in group_moe.values(): + for param_group in moe_group.values(): + cur_group: List[nn.Parameter] = [] + all_groups: List[List[nn.Parameter]] = [] + size_of_cur_group = 0 + + for param in cast(List[nn.Parameter], param_group['params']): + if size_of_cur_group + param.numel() <= max_group_size: + cur_group.append(param) + size_of_cur_group += param.numel() + else: + all_groups.append(cur_group) + cur_group = [param] + size_of_cur_group = param.numel() + + if cur_group: + all_groups.append(cur_group) + + for group in all_groups: + param_groups.append({**param_group, 'params': group}) + else: + for moe_group in group_moe.values(): + for param_group in moe_group.values(): + param_groups.append(param_group) + + return param_groups + + +def is_moe_param_group(param_group): + return param_group.get('moe', False) + + +def configure_moe_param_groups(model_parameters: List): + assert isinstance(model_parameters, list), "model_parameters must be a list" + + for p in model_parameters: + # match torch.optim.Optimizer expectations, + # see: https://github.com/pytorch/pytorch/blob/2ffab6e663b9c6951048b8c8ba82d2cc5ca5c2fc/torch/optim/optimizer.py#L270-L272 + if not isinstance(p, (torch.Tensor, dict)): + raise TypeError("param argument that would be given to the optimizer should be " + f"an iterable of Tensors or dicts, but got {type(p)}") + + # peak at the first element to determine how to proceed + first = model_parameters[0] + + # Case 1: model_parameters is a list of torch.nn.Parameter + # -> need to create moe compatible param groups + if isinstance(first, torch.nn.Parameter): + param_group = {'params': model_parameters, 'name': 'dense-params'} + return split_params_into_different_moe_groups_for_optimizer(param_group) + + # Case 2: model_parameters is a list of param groups List[dict] + # -> moe compatible param groups might already exist, if not create them + elif isinstance(first, dict): + #there are no moe groups created + if not any(['moe' in param_group for param_group in model_parameters]): + return split_params_into_different_moe_groups_for_optimizer(model_parameters) + else: + # moe groups exist, nothing to do + return model_parameters diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee220003517d39176acde8ad9c1176048a2242c7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/comet.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/comet.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..213191f0bdff5bdba17dcf634303bcd64f19c1b7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/comet.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb6c34ac797ef10c2e9d3b9e433ae506e4fb6c61 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50bd44d010e98247fc67c83615aba4399f3c1739 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2aafc2ef90c68744692d64c79ae5d4b663a74fa8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..726283ba773ea26b38dde1ed8bc903db0633c6c1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9cbb9e045371efeed8da59eaf1a1443a239305b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05cf6516620132fcd98416180ec997b5fe2fe8dd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/comet.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/comet.py new file mode 100644 index 0000000000000000000000000000000000000000..d8bc4017800f3914fd1a1881b0510c3cdee75324 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/comet.py @@ -0,0 +1,92 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import TYPE_CHECKING, Any, Tuple, List, Dict, Optional + +from .utils import check_comet_availability +from .monitor import Monitor + +import deepspeed.comm as dist + +if TYPE_CHECKING: + import comet_ml + from .config import CometConfig + +Name = str +Value = Any +GlobalSamples = int +Event = Tuple[Name, Value, GlobalSamples] + + +class CometMonitor(Monitor): + + def __init__(self, comet_config: "CometConfig"): + super().__init__(comet_config) + check_comet_availability() + import comet_ml + + self.enabled = comet_config.enabled + self._samples_log_interval = comet_config.samples_log_interval + self._experiment: Optional["comet_ml.ExperimentBase"] = None + + if self.enabled and dist.get_rank() == 0: + self._experiment = comet_ml.start( + api_key=comet_config.api_key, + project=comet_config.project, + workspace=comet_config.workspace, + experiment_key=comet_config.experiment_key, + mode=comet_config.mode, + online=comet_config.online, + ) + + if comet_config.experiment_name is not None: + self._experiment.set_name(comet_config.experiment_name) + + self._events_log_scheduler = EventsLogScheduler(comet_config.samples_log_interval) + + @property + def experiment(self) -> Optional["comet_ml.ExperimentBase"]: + return self._experiment + + @property + def samples_log_interval(self) -> int: + return self._samples_log_interval + + def write_events(self, event_list: List[Event]) -> None: + if not self.enabled or dist.get_rank() != 0: + return None + + for event in event_list: + name = event[0] + value = event[1] + engine_global_samples = event[2] + + if self._events_log_scheduler.needs_logging(name, engine_global_samples): + self._experiment.__internal_api__log_metric__( + name=name, + value=value, + step=engine_global_samples, + ) + + +class EventsLogScheduler: + + def __init__(self, samples_log_interval: int): + self._samples_log_interval = samples_log_interval + self._last_logged_events_samples: Dict[str, int] = {} + + def needs_logging(self, name: str, current_sample: int) -> bool: + if name not in self._last_logged_events_samples: + self._last_logged_events_samples[name] = current_sample + return True + + last_logged_sample = self._last_logged_events_samples[name] + samples_delta = current_sample - last_logged_sample + + if samples_delta >= self._samples_log_interval: + self._last_logged_events_samples[name] = current_sample + return True + + return False diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/config.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/config.py new file mode 100644 index 0000000000000000000000000000000000000000..960ce1ba997a231edec53e84855942f3d6ece1d8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/config.py @@ -0,0 +1,144 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Optional + +from pydantic import model_validator +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + + +def get_monitor_config(param_dict): + monitor_dict = {key: param_dict.get(key, {}) for key in ("tensorboard", "wandb", "csv_monitor", "comet")} + return DeepSpeedMonitorConfig(**monitor_dict) + + +class TensorBoardConfig(DeepSpeedConfigModel): + """Sets parameters for TensorBoard monitor.""" + + enabled: bool = False + """ Whether logging to Tensorboard is enabled. Requires `tensorboard` package is installed. """ + + output_path: str = "" + """ + Path to where the Tensorboard logs will be written. If not provided, the + output path is set under the training script’s launching path. + """ + + job_name: str = "DeepSpeedJobName" + """ Name for the current job. This will become a new directory inside `output_path`. """ + + +class WandbConfig(DeepSpeedConfigModel): + """Sets parameters for WandB monitor.""" + + enabled: bool = False + """ Whether logging to WandB is enabled. Requires `wandb` package is installed. """ + + group: Optional[str] = None + """ Name for the WandB group. This can be used to group together runs. """ + + team: Optional[str] = None + """ Name for the WandB team. """ + + project: str = "deepspeed" + """ Name for the WandB project. """ + + +class CSVConfig(DeepSpeedConfigModel): + """Sets parameters for CSV monitor.""" + + enabled: bool = False + """ Whether logging to local CSV files is enabled. """ + + output_path: str = "" + """ + Path to where the csv files will be written. If not provided, the output + path is set under the training script’s launching path. + """ + + job_name: str = "DeepSpeedJobName" + """ Name for the current job. This will become a new directory inside `output_path`. """ + + +class CometConfig(DeepSpeedConfigModel): + """ + Sets parameters for Comet monitor. For logging data Comet uses + experiment object. + https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/ + """ + + enabled: bool = False + """ Whether logging to Comet is enabled. Requires `comet_ml` package is installed. """ + + samples_log_interval: int = 100 + """ Metrics will be submitted to Comet after processing every `samples_log_intervas` samples""" + + project: Optional[str] = None + """ + Comet project name. Can be set through .comet.config file or environment variable COMET_PROJECT_NAME + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + workspace: Optional[str] = None + """ + Comet workspace name. Can be set through .comet.config file or environment variable COMET_WORKSPACE + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + api_key: Optional[str] = None + """ + Comet API key. Can be set through .comet.config file or environment variable COMET_API_KEY + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + experiment_name: Optional[str] = None + """ + The name for comet experiment to be used for logging. + Can be set through .comet.config file or environment variable COMET_EXPERIMENT_NAME + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + experiment_key: Optional[str] = None + """ + The key for comet experiment to be used for logging. Must be an alphanumeric string whose length is between 32 and 50 characters. + Can be set through .comet.config or environment variable COMET_EXPERIMENT_KEY + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + online: Optional[bool] = None + """ + If True, the data will be logged to Comet server, otherwise it will be stored locally in offline experiment + Defaults to True. + """ + + mode: Optional[str] = None + """ + Control how the Comet experiment is started, 3 options are possible.: + - "get": Continue logging to an existing experiment identified by the `experiment_key` value. + - "create": Always creates of a new experiment, useful for HPO sweeps. + - "get_or_create" (default): Starts a fresh experiment if required, or persists logging to an existing one. + """ + + +class DeepSpeedMonitorConfig(DeepSpeedConfigModel): + """Sets parameters for various monitoring methods.""" + + tensorboard: TensorBoardConfig = {} + """ TensorBoard monitor, requires `tensorboard` package is installed. """ + + comet: CometConfig = {} + """ Comet monitor, requires `comet_ml` package is installed """ + + wandb: WandbConfig = {} + """ WandB monitor, requires `wandb` package is installed. """ + + csv_monitor: CSVConfig = {} + """ Local CSV output of monitoring data. """ + + @model_validator(mode="after") + def check_enabled(self): + enabled = self.tensorboard.enabled or self.wandb.enabled or self.csv_monitor.enabled or self.comet.enabled + self.__dict__["enabled"] = enabled + return self diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/csv_monitor.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/csv_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..c7a19b14ad8227fc8187bbe161667285b7d0c717 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/csv_monitor.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .monitor import Monitor +import os + +import deepspeed.comm as dist + + +class csvMonitor(Monitor): + + def __init__(self, csv_config): + super().__init__(csv_config) + self.filenames = [] + self.enabled = csv_config.enabled + self.output_path = csv_config.output_path + self.job_name = csv_config.job_name + self.log_dir = self.setup_log_dir() + + def setup_log_dir(self, base=os.path.join(os.path.expanduser("~"), "csv_monitor")): + if self.enabled and dist.get_rank() == 0: + if self.output_path is not None: + log_dir = os.path.join(self.output_path, self.job_name) + # NOTE: This code path currently is never used since the default tensorboard_output_path is an empty string and not None. Saving it in case we want this functionality in the future. + else: + if "DLWS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLWS_JOB_ID"] + elif "DLTS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLTS_JOB_ID"] + else: + infra_job_id = "unknown-job-id" + + csv_monitor_dir_name = os.path.join(infra_job_id, "logs") + log_dir = os.path.join(base, csv_monitor_dir_name, self.job_name) + os.makedirs(log_dir, exist_ok=True) + return log_dir + + def write_events(self, event_list): + if self.enabled and dist.get_rank() == 0: + import csv + # We assume each event_list element is a tensorboard-style tuple in the format: (log_name: String, value, step: Int) + for event in event_list: + log_name = event[0] + value = event[1] + step = event[2] + + # Set the header to the log_name + # Need this check because the deepspeed engine currently formats log strings to separate with '/' + if '/' in log_name: + record_splits = log_name.split('/') + header = record_splits[len(record_splits) - 1] + else: + header = log_name + + # sanitize common naming conventions into filename + filename = log_name.replace('/', '_').replace(' ', '_') + fname = self.log_dir + '/' + filename + '.csv' + + # Open file and record event. Insert header if this is the first time writing + with open(fname, 'a+') as csv_monitor_file: + csv_monitor_writer = csv.writer(csv_monitor_file) + if filename not in self.filenames: + self.filenames.append(filename) + csv_monitor_writer.writerow(['step', header]) + csv_monitor_writer.writerow([step, value]) diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/monitor.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..e7e26dc483d9707bfc04febda348bff30d25e53b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/monitor.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Support different forms of monitoring such as wandb and tensorboard +""" + +from abc import ABC, abstractmethod +import deepspeed.comm as dist + + +class Monitor(ABC): + + @abstractmethod + def __init__(self, monitor_config): + self.monitor_config = monitor_config + + @abstractmethod + def write_events(self, event_list): + pass + + +from .wandb import WandbMonitor +from .tensorboard import TensorBoardMonitor +from .csv_monitor import csvMonitor +from .comet import CometMonitor + + +class MonitorMaster(Monitor): + + def __init__(self, monitor_config): + super().__init__(monitor_config) + self.tb_monitor = None + self.wandb_monitor = None + self.csv_monitor = None + self.comet_monitor = None + self.enabled = monitor_config.enabled + + if dist.get_rank() == 0: + if monitor_config.tensorboard.enabled: + self.tb_monitor = TensorBoardMonitor(monitor_config.tensorboard) + if monitor_config.wandb.enabled: + self.wandb_monitor = WandbMonitor(monitor_config.wandb) + if monitor_config.csv_monitor.enabled: + self.csv_monitor = csvMonitor(monitor_config.csv_monitor) + if monitor_config.comet.enabled: + self.comet_monitor = CometMonitor(monitor_config.comet) + + def write_events(self, event_list): + if dist.get_rank() == 0: + if self.tb_monitor is not None: + self.tb_monitor.write_events(event_list) + if self.wandb_monitor is not None: + self.wandb_monitor.write_events(event_list) + if self.csv_monitor is not None: + self.csv_monitor.write_events(event_list) + if self.comet_monitor is not None: + self.comet_monitor.write_events(event_list) diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/tensorboard.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/tensorboard.py new file mode 100644 index 0000000000000000000000000000000000000000..985c9ed44b6f5be28785699bca4a3638ed6063bd --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/tensorboard.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .utils import check_tb_availability +from .monitor import Monitor +import os + +import deepspeed.comm as dist + + +class TensorBoardMonitor(Monitor): + + def __init__(self, tensorboard_config): + super().__init__(tensorboard_config) + check_tb_availability() + + self.summary_writer = None + self.enabled = tensorboard_config.enabled + self.output_path = tensorboard_config.output_path + self.job_name = tensorboard_config.job_name + + if self.enabled and dist.get_rank() == 0: + self.get_summary_writer() + + def get_summary_writer(self, base=os.path.join(os.path.expanduser("~"), "tensorboard")): + if self.enabled and dist.get_rank() == 0: + from torch.utils.tensorboard import SummaryWriter + if self.output_path is not None: + log_dir = os.path.join(self.output_path, self.job_name) + # NOTE: This code path currently is never used since the default output_path is an empty string and not None. Saving it in case we want this functionality in the future. + else: + if "DLWS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLWS_JOB_ID"] + elif "DLTS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLTS_JOB_ID"] + else: + infra_job_id = "unknown-job-id" + + summary_writer_dir_name = os.path.join(infra_job_id, "logs") + log_dir = os.path.join(base, summary_writer_dir_name, self.output_path) + os.makedirs(log_dir, exist_ok=True) + self.summary_writer = SummaryWriter(log_dir=log_dir) + return self.summary_writer + + def write_events(self, event_list, flush=True): + if self.enabled and self.summary_writer is not None and dist.get_rank() == 0: + for event in event_list: + self.summary_writer.add_scalar(*event) + if flush: + self.summary_writer.flush() + + def flush(self): + if self.enabled and self.summary_writer is not None and dist.get_rank() == 0: + self.summary_writer.flush() diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/utils.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f5530e8532e1a6982ce8acdf20b98aa8fdbb5e81 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/utils.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from packaging import version as pkg_version + + +def check_tb_availability(): + try: + # torch.utils.tensorboard will fail if `tensorboard` is not available, + # see their docs for more details: https://pytorch.org/docs/1.8.0/tensorboard.html + import tensorboard # noqa: F401 # type: ignore + except ImportError: + print('If you want to use tensorboard logging, please `pip install tensorboard`') + raise + + +def check_wandb_availability(): + try: + import wandb # noqa: F401 # type: ignore + except ImportError: + print( + 'If you want to use wandb logging, please `pip install wandb` and follow the instructions at https://docs.wandb.ai/quickstart' + ) + raise + + +def check_comet_availability(): + try: + import comet_ml + comet_version = pkg_version.parse(comet_ml.__version__) + if comet_version < pkg_version.Version("3.41.0"): + raise ImportError("`comet_ml` must have at least version 3.41.0") + except ImportError: + print('If you want to use comet logging, please `pip install "comet_ml>=3.41.0"`') + raise diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/wandb.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/wandb.py new file mode 100644 index 0000000000000000000000000000000000000000..174a2eb2d3b7e71d37ac5b61c7fdb44352455f07 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/monitor/wandb.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .utils import check_wandb_availability +from .monitor import Monitor + +import deepspeed.comm as dist + + +class WandbMonitor(Monitor): + + def __init__(self, wandb_config): + super().__init__(wandb_config) + check_wandb_availability() + import wandb + + self.enabled = wandb_config.enabled + self.group = wandb_config.group + self.team = wandb_config.team + self.project = wandb_config.project + + if self.enabled and dist.get_rank() == 0: + wandb.init(project=self.project, group=self.group, entity=self.team) + + def log(self, data, step=None, commit=None): + if self.enabled and dist.get_rank() == 0: + import wandb + return wandb.log(data, step=step, commit=commit) + + def write_events(self, event_list): + if self.enabled and dist.get_rank() == 0: + for event in event_list: + label = event[0] + value = event[1] + step = event[2] + self.log({label: value}, step=step) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/nebula/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nebula/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8045d0a09285f4a9133159be15e5de540455d934 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30483f724d6cfd67506f7ca131abf51e500a76e5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..61b782db27e5b0249dbde3fe939b8d6d83e2e0c8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/config.py b/venv/lib/python3.10/site-packages/deepspeed/nebula/config.py new file mode 100644 index 0000000000000000000000000000000000000000..dc49185738c92a3173f2fd5c68f1d6ab5a32dd92 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nebula/config.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject +from deepspeed.nebula.constants import * + + +class DeepSpeedNebulaConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedNebulaConfig, self).__init__() + + self.enabled = None + self.persistent_storage_path = None + self.persistent_time_interval = None + self.num_of_version_in_retention = None + self.enable_nebula_load = None + + if NEBULA in param_dict.keys(): + nebula_dict = param_dict[NEBULA] + else: + nebula_dict = {} + + self._initialize(nebula_dict) + + def _initialize(self, nebula_dict): + self.enabled = get_scalar_param(nebula_dict, NEBULA_ENABLED, NEBULA_ENABLED_DEFAULT) + + self.load_path = get_scalar_param(nebula_dict, NEBULA_LOAD_PATH, NEBULA_LOAD_PATH_DEFAULT) + + self.enable_nebula_load = get_scalar_param(nebula_dict, NEBULA_ENABLE_NEBULA_LOAD, + NEBULA_ENABLE_NEBULA_LOAD_DEFAULT) + + self.persistent_storage_path = get_scalar_param(nebula_dict, NEBULA_PERSISTENT_STORAGE_PATH, + NEBULA_PERSISTENT_STORAGE_PATH_DEFAULT) + + self.persistent_time_interval = get_scalar_param(nebula_dict, NEBULA_PERSISTENT_TIME_INTERVAL, + NEBULA_PERSISTENT_TIME_INTERVAL_DEFAULT) + + self.num_of_version_in_retention = get_scalar_param(nebula_dict, NEBULA_NUM_OF_VERSION_IN_RETENTION, + NEBULA_NUM_OF_VERSION_IN_RETENTION_DEFAULT) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/constants.py b/venv/lib/python3.10/site-packages/deepspeed/nebula/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..9fa5769b55979e4dcd5c80ead06d2117dcc2ec40 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nebula/constants.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######################################### +# nebula +######################################### +# Nebula. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +NEBULA_FORMAT = ''' +nebula should be enabled as: +"session_params": { + "nebula": { + "enabled": true, + "persistent_storage_path": "/foo/bar", + "persistent_time_interval": 100, + "num_of_version_in_retention": 2, + "enable_nebula_load": true + } +} +''' + +NEBULA = "nebula" + +NEBULA_ENABLED = "enabled" +NEBULA_ENABLED_DEFAULT = False + +# There is a case where customer want to load the checkpoint saved +# by raw torch. Because nebula cannot load torch checkpoint directly +# as they have different folder structures to bring the gap for +# loading(the data are totally same in bytes for torch and nebula +# saving). +# In this case, we must disable nebula load to use raw torch load. +# Customer can just set NEBULA_ENABLE_NEBULA_LOAD to False. Then use +# original way of deepspeed to load, i.e. set the value of "--load". +NEBULA_ENABLE_NEBULA_LOAD = "enable_nebula_load" +NEBULA_ENABLE_NEBULA_LOAD_DEFAULT = True + +# When you want to resume the previous checkpoint saved by nebula, +# you can set NEBULA_LOAD_PATH as the parent folder of checkpoint. +# If NEBULA_LOAD_PATH is None, the NEBULA_PERSISTENT_STORAGE_PATH +# will be the default path to load. +NEBULA_LOAD_PATH = "nebula_load_path" +NEBULA_LOAD_PATH_DEFAULT = None + +# Nebula will save the checkpoint under NEBULA_LOAD_PATH in the +# asynchronous way. +NEBULA_PERSISTENT_STORAGE_PATH = "persistent_storage_path" +NEBULA_PERSISTENT_STORAGE_PATH_DEFAULT = None + +# Time interval to trigger the nebula persistence. +NEBULA_PERSISTENT_TIME_INTERVAL = "persistent_time_interval" +NEBULA_PERSISTENT_TIME_INTERVAL_DEFAULT = 100 + +# Checkpoint number which will be kept in memory. Let us say, +# if the value is 2. Then we have checkpoints 1 and 2 are ready +# now. When it comes to checkpoint 3, the 1 will be removed if +# 1 has been persisted to disk. +NEBULA_NUM_OF_VERSION_IN_RETENTION = "num_of_version_in_retention" +NEBULA_NUM_OF_VERSION_IN_RETENTION_DEFAULT = 2 + +# Nebula envs +NEBULA_EXPORT_ENVS = [ + 'DLTS_JOB_ID', 'DLTS_NUM_WORKER', 'NEBULA_PERSISTENT_STORAGE_PATH', 'NEBULA_PERSISTENT_TIME_INTERVAL', + 'AML_RUN_ID', 'AZUREML_RUN_TOKEN', 'AZUREML_WORKSPACE_SCOPE', 'AZUREML_EXPERIMENT_SCOPE', + 'AZUREML_RUN_HISTORY_SERVICE_ENDPOINT', 'AZUREML_RUN_ID', 'NEBULA_MEMORY_BUFFER_SIZE', + 'AZUREML_PARAMETER_ITPJOB_NAME', 'FC_TASKROLE_NAME', 'FC_TASK_INDEX', 'MASTER_HOST', 'LOCAL_HOST', + 'AZUREML_BLOB_ACCOUNT_NAME', 'AZUREML_BLOB_ACCOUNT_KEY' +] + +# ITP env files +DLTS_POD_ENV_PATH = '/dlts-runtime/env/pod.env' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6d0de857cbd345169651f0133cee87a992ea907e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .perf_run_sweep import sweep_main, parse_sweep_arguments +from .perf_generate_param import generate_main +from .test_ds_aio import ds_io_main diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7f4f55efbf2e87bbf95a74a8ac9a9dcb21471ea Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afef6d87676ebb4051cd7167fba6b055200e7896 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d59209874a11edd77a86b0cf342c10a1b39a847 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b6ddf90739ba5ea094f38596f39f4ad060ec819 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3546b4f4c26d2c8e2acff714a46d80d5e283af17 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce739961549d224fd649441a464ee695f686acc6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d56dc8e0bb915333bc56ed35db97d48bf44b6f4e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..147ee6251c076435e6a283295db9a5910afe9b40 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5045ea233fca73122d632acdcee1ede6c95674d3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05fbca2a4dd3712a78fe8f88d6ce4196efbdff5f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a19b99eb2c2192a25ee994fa81015ff96c744bb Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ab553e181e434a0a8bc62d4edbdb6b512feacdd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f69fa66d4a3afe4b0f1ea34a28e994b702d7b421 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1df8cfc46405e00edfe017f50bf7a5b58e660e5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef0f874e0fd88264967896bbe219b1fb05bc4f24 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36ae3bc0330e52cad10958206f22a4bf28f5eef8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_args.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_args.py new file mode 100644 index 0000000000000000000000000000000000000000..95165d98e3ffeb31684592cd6a5a24018ee08c0a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_args.py @@ -0,0 +1,214 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +import argparse +import os +from .test_ds_aio_utils import refine_integer_value +from .ds_aio_constants import AIO_HANDLE, AIO_BASIC, TORCH_FAST_IO, TORCH_IO, VALID_ENGINES +from deepspeed.accelerator import get_accelerator + +MAPPING_DELIMITER = ':' + + +def refine_args(args): + if args.io_size and type(args.io_size) == str: + args.io_size = refine_integer_value(args.io_size) + + if args.block_size and type(args.block_size) == str: + args.block_size = refine_integer_value(args.block_size) + + if args.fast_io_size and type(args.fast_io_size) == str: + args.fast_io_size = refine_integer_value(args.fast_io_size) + + return args + + +def _get_mapping_dict(args): + if args.folder is not None: + d = {i: args.folder for i in range(args.multi_process)} + else: + d = {} + for m in args.folder_to_device_mapping: + fields = m.split(MAPPING_DELIMITER) + d[fields[1]] = fields[0] + + return d + + +def _validate_folder_mapping(args): + no_error = True + error_messages = [] + invalid_mappings = [m for m in args.folder_to_device_mapping if MAPPING_DELIMITER not in m] + if len(invalid_mappings) > 0: + error_messages.append( + f'Missing delimiter ({MAPPING_DELIMITER}) in folder_to_device_mapping {invalid_mappings}') + no_error = False + + folder_list = [m.split(MAPPING_DELIMITER)[0] for m in args.folder_to_device_mapping] + invalid_folders = [d for d in folder_list if not os.path.exists(d)] + if len(invalid_folders) > 0: + error_messages.append(f'Invalid folders in folder_to_device_mapping: {invalid_folders}') + no_error = False + + if args.gpu: + device_list = [int(m.split(MAPPING_DELIMITER)[1]) for m in args.folder_to_device_mapping] + invalid_device_list = [dev_id for dev_id in device_list if not dev_id < get_accelerator().device_count()] + if len(invalid_device_list) > 0: + error_messages.append(f'Invalid device ids in folder_to_device_mapping: {invalid_device_list}') + no_error = False + + return no_error, error_messages + + +def validate_args(args): + no_error = True + error_messages = [] + + if args.folder is not None and len(args.folder_to_device_mapping) > 0: + error_messages.append(f'--folder and --folder_to_device_mapping cannot be specified together.') + no_error = False + elif args.folder is None and len(args.folder_to_device_mapping) == 0: + error_messages.append(f'At least one of --folder or --folder_to_device_mapping must be specified.') + no_error = False + + # Validate --folder + if args.folder is not None and not os.path.exists(args.folder): + no_error = False + error_messages.append(f'Invalid folder in --folder: {args.folder} ') + + # Validate --folder_mapping_to_device + if len(args.folder_to_device_mapping) > 0: + no_mapping_error, mapping_error_messages = _validate_folder_mapping(args) + no_error = no_error and no_mapping_error + error_messages += mapping_error_messages + + # Validate --engine + if args.engine not in VALID_ENGINES: + no_error = False + error_messages.append(f'Invalid engine {args.engine}. Valid options = {VALID_ENGINES}') + + # Validate --engine=torch_io + if args.engine == TORCH_IO: + if args.read: + no_error = False + error_messages.append(f'Read not currently supported for --engine={TORCH_IO}') + + if not no_error: + print(f'Found {len(error_messages)} validation error(s)') + # Validate --gpu, --use_gds + if args.use_gds and not args.gpu: + error_messages.append(f'--gpu must be set to transfer with --use_gds') + no_error = False + + if not no_error: + print(f'Found {len(error_messages)} validation errors') + for i, msg in enumerate(error_messages): + print(f'{i+1}: {msg}') + + return no_error + + +def parse_arguments(): + parser = argparse.ArgumentParser() + + parser.add_argument('--folder', default=None, type=str, help='Folder to use for I/O.') + + parser.add_argument('--folder_to_device_mapping', + default=[], + nargs='+', + help='Specification of mapping of folder to (gpu) device id, (ignored for cpu accesses).' + 'Can be specified multiple times for multi-process runs,' + 'e.g. --folder_to_device_mapping /mnt/nvme0:0 --folder_to_device_mapping /mnt/nvme1:15 --gpu' + 'means access /mnt/nvme0 with gpu 0 and /mnt/nvme1 with gpu 15') + + parser.add_argument('--io_size', type=str, default=None, required=True, help='Number of bytes to read or write.') + + parser.add_argument('--fast_io_size', type=str, default='64M', help='Size of fast_io pinned buffer (bytes).') + + parser.add_argument('--read', action='store_true', help='Perform read I/O (default is write)') + + parser.add_argument('--multi_process', + type=int, + default=1, + help='Number of parallel processes doing I/O (default 1).') + + parser.add_argument('--block_size', + type=str, + default='1M', + help='I/O block size. Can use K, M, or G suffix (default 1M for 1 megabytes).') + + parser.add_argument('--queue_depth', type=int, default=32, help='I/O queue depth (default 32).') + + parser.add_argument('--single_submit', + action='store_true', + help='Submit I/O requests in singles (default is submit queue_depth amount at once.).') + + parser.add_argument( + '--sequential_requests', + action='store_true', + help= + 'Delay I/O request submission until completion of prior requests (default is overlap I/O submission and completion requests.).' + ) + + parser.add_argument('--validate', action='store_true', help='Perform validation of I/O transfer in library.') + + parser.add_argument( + '--engine', + type=str, + default=AIO_HANDLE, + help= + f'Engine to perform I/O. Options are [{AIO_HANDLE}, {AIO_BASIC}, {TORCH_IO}, {TORCH_FAST_IO}]. Default is aio_handle' + ) + + parser.add_argument('--loops', type=int, default=3, help='Count of operation repetitions') + + parser.add_argument('--io_parallel', type=int, default=None, help='Per iop parallelism') + + parser.add_argument('--gpu', action='store_true', help='Use GPU memory') + + parser.add_argument('--use_gds', action='store_true', help='Enable GDS AIO') + + parser.add_argument('--slow_bounce_buffer', + action='store_true', + help='For GPU memory transfers, measure impact of bounce buffer pinning on critical path.') + + parser.add_argument('--torch_legacy_save', action='store_true', help='Use torch legacy save approach') + + parser.add_argument('--use_accelerator_pin_memory', + action='store_true', + help='Obtain pinned (CPU page-locked) tensors from accelerator') + + parser.add_argument('--warmup_loops', type=int, default=1, help='Count of operation warmup repetitions') + + parser.add_argument('--include_warmup_time', action='store_true', help='Include warmup latency in results') + + parser.add_argument('--different_file_each_iteration', + action='store_true', + help='Read/write a different file on each iteration.') + + args = parser.parse_args() + print(f'args = {args}') + return args + + +def get_validated_args(): + args = parse_arguments() + args = refine_args(args) + if not validate_args(args): + quit() + print(f'Successful validation of command line arguments') + args.total_loops = args.warmup_loops + args.loops + peer_tag = 'gpu' if args.gpu else 'process' + args.mapping_dict = _get_mapping_dict(args) + args.mapping_list = [(device_id, folder) for device_id, folder in args.mapping_dict.items()] + assert len(args.mapping_dict) == len(args.mapping_list) + print(f'Configuring {len(args.mapping_list)} {peer_tag} to folder mapping') + for i, (device_id, folder) in enumerate(args.mapping_list): + print(f'[{i}]: {peer_tag} {device_id} <----> {folder}') + + return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_basic.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..7be7bcd1cd67c7741d3a8008af900f3349f798a3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_basic.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +import os +import time +from deepspeed.ops.aio import AsyncIOBuilder +from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor +from .ds_aio_constants import * + + +class AIOBasic_Engine(object): + + def __init__(self, args, tid, read_op): + self.ctxt = self._create_context(args, tid, read_op) + + def fini(self): + self.ctxt[BUFFER].detach() + self.ctxt[BUFFER] = None + + def read(self, args, tid, loop_id): + start_time = time.time() + AsyncIOBuilder().load().aio_read(self.ctxt[BUFFER], self.ctxt[FILE], args.block_size, args.queue_depth, + args.single_submit, not args.sequential_requests, args.validate) + end_time = time.time() + self.ctxt[ELAPSED_SEC] += end_time - start_time + + def write(self, args, tid, loop_id): + # Avoid overwriting existing files as it could be artificially faster + if os.path.isfile(self.ctxt[FILE]): + os.remove(self.ctxt[FILE]) + + start_time = time.time() + AsyncIOBuilder().load().aio_write(self.ctxt[BUFFER], self.ctxt[FILE], args.block_size, args.queue_depth, + args.single_submit, not args.sequential_requests, args.validate) + end_time = time.time() + self.ctxt[ELAPSED_SEC] += end_time - start_time + + def _create_context(self, args, tid, read_op): + io_string = "Read" if read_op else "Write" + device_id, folder = args.mapping_list[tid] + filename = create_filename(folder, args.read, args.io_size, tid) + if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): + create_file(filename, args.io_size) + + task_log(tid, f'Allocate tensor of size {args.io_size} bytes') + + buffer = create_page_locked_tensor(args.io_size, True) + + task_log(tid, + f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}') + + task_log(tid, f'created deepspeed aio basic engine') + + ctxt = {} + ctxt[FILE] = filename + ctxt[NUM_BYTES] = args.io_size + ctxt[BUFFER] = buffer + ctxt[ELAPSED_SEC] = 0 + return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_constants.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..1b07ed8672ef4b4453523a2292141f18ba50d51b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_constants.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +AIO_HANDLE = 'aio_handle' +AIO_BASIC = 'aio_basic' +TORCH_IO = 'torch_io' +TORCH_FAST_IO = 'torch_fastio' +VALID_ENGINES = [AIO_HANDLE, AIO_BASIC, TORCH_IO, TORCH_FAST_IO] + +BUFFER = 'buffer' +BOUNCE_BUFFER = 'bounce_buffer' +NUM_BYTES = 'num_bytes' +FILE = 'file' +HANDLE = 'handle' +ELAPSED_SEC = 'elapsed_sec' +FAST_IO_BUFFER = 'fast_io_buffer' +USE_CPU_LOCKED_TENSOR = 'cpu_locked_tensor' +USE_GDS = 'gds' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_handle.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_handle.py new file mode 100644 index 0000000000000000000000000000000000000000..efc157078dd95b21691bd6e4821ce905ae69fde1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_handle.py @@ -0,0 +1,126 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +import torch +import os +import time +from deepspeed.ops.aio import AsyncIOBuilder +from deepspeed.ops.op_builder import GDSBuilder +from deepspeed.accelerator import get_accelerator +from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor +from .ds_aio_constants import * + + +class AIOHandle_Engine(object): + + def __init__(self, args, tid, read_op): + self.ctxt = self._create_context(args, tid, read_op) + + def fini(self): + for buf in [BUFFER, BOUNCE_BUFFER]: + if self.ctxt[buf] is not None: + if self.ctxt[USE_CPU_LOCKED_TENSOR]: + self.ctxt[HANDLE].free_cpu_locked_tensor(self.ctxt[buf]) + + self.ctxt[buf].detach() + self.ctxt[buf] = None + + def read(self, args, tid, loop_id): + handle = self.ctxt[HANDLE] + + start_time = time.time() + dest_buffer = BOUNCE_BUFFER if self.ctxt[BOUNCE_BUFFER] is not None else BUFFER + ret = handle.pread(self.ctxt[dest_buffer], self.ctxt[FILE][loop_id], args.validate, True) + assert ret != -1 + handle.wait() + if dest_buffer == BOUNCE_BUFFER: + self.ctxt[BUFFER].data.copy_(self.ctxt[BOUNCE_BUFFER].data) + end_time = time.time() + self.ctxt[ELAPSED_SEC].append(end_time - start_time) + + def write(self, args, tid, loop_id): + # Avoid overwriting existing files as it could be artificially faster + # if os.path.isfile(self.ctxt[FILE]): + # os.remove(self.ctxt[FILE]) + + handle = self.ctxt[HANDLE] + start_time = time.time() + if self.ctxt[BOUNCE_BUFFER] is not None: + source_buffer = BOUNCE_BUFFER + self.ctxt[BOUNCE_BUFFER].data.copy_(self.ctxt[BUFFER].data) + else: + source_buffer = BUFFER + ret = handle.pwrite(self.ctxt[source_buffer], self.ctxt[FILE][loop_id], args.validate, True) + assert ret != -1 + handle.wait() + end_time = time.time() + self.ctxt[ELAPSED_SEC].append(end_time - start_time) + + def _create_files(self, args, folder, tid): + if args.different_file_each_iteration: + filenames = [ + create_filename(folder, args.read, args.io_size, f'{tid}_{l}') for l in range(args.total_loops) + ] + else: + filenames = [ + create_filename(folder, args.read, args.io_size, f'{tid}_{0}') for _ in range(args.total_loops) + ] + + if args.read: + for f in filenames: + if not (os.path.isfile(f) and os.path.getsize(f) == args.io_size): + create_file(f, args.io_size) + else: + for f in filenames: + if os.path.isfile(f): + os.remove(f) + + return filenames + + def _create_context(self, args, tid, read_op): + io_string = "Read" if read_op else "Write" + device_id, folder = args.mapping_list[tid] + filenames = self._create_files(args, folder, tid) + + gds = True if args.use_gds else False + io_parallel = args.io_parallel if args.io_parallel else 1 + if gds: + handle = GDSBuilder().load().gds_handle(args.block_size, args.queue_depth, args.single_submit, + not args.sequential_requests, io_parallel) + else: + handle = AsyncIOBuilder().load().aio_handle(args.block_size, args.queue_depth, args.single_submit, + not args.sequential_requests, io_parallel) + task_log(tid, f'Created DeepNVMe handle engine') + + bounce_buffer = None + if args.gpu: + device_name = get_accelerator().device_name(device_id) + buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=device_name) + if gds: + handle.pin_device_tensor(buffer) + elif not args.slow_bounce_buffer: + bounce_buffer = create_page_locked_tensor(args.io_size, args.use_accelerator_pin_memory, handle) + else: + buffer = create_page_locked_tensor(args.io_size, args.use_accelerator_pin_memory, handle) + task_log(tid, f'Allocate tensor of size {args.io_size} bytes') + + ctxt = {} + ctxt[FILE] = filenames + ctxt[NUM_BYTES] = args.io_size + ctxt[HANDLE] = handle + ctxt[USE_GDS] = gds + ctxt[BUFFER] = buffer + ctxt[BOUNCE_BUFFER] = bounce_buffer + ctxt[ELAPSED_SEC] = [] + ctxt[USE_CPU_LOCKED_TENSOR] = not args.use_accelerator_pin_memory + + task_log(tid, + f'{io_string} file {filenames} of size {args.io_size} bytes from buffer on device {buffer.device}', + force=True) + + return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_job.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_job.py new file mode 100644 index 0000000000000000000000000000000000000000..0f9c8b5f1bcc5b475ef76740a88a10345f5d22d1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_job.py @@ -0,0 +1,50 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" +import subprocess +import shlex + + +class Job(object): + + def __init__(self, cmd_line, output_file=None, work_dir=None): + self.cmd_line = cmd_line + self.output_file = output_file + self.work_dir = work_dir + self.output_fd = None + + def cmd(self): + return self.cmd_line + + def get_stdout(self): + return self.output_fd + + def get_stderr(self): + return self.output_fd + + def get_cwd(self): + return self.work_dir + + def open_output_file(self): + if self.output_file is not None: + self.output_fd = open(self.output_file, 'w') + + def close_output_file(self): + if self.output_fd is not None: + self.output_fd.close() + self.output_fd = None + + +def run_job(job, verbose=False): + args = shlex.split(' '.join(job.cmd())) + if verbose: + print(f'args = {args}') + job.open_output_file() + proc = subprocess.run(args=args, stdout=job.get_stdout(), stderr=job.get_stderr(), cwd=job.get_cwd()) + job.close_output_file() + assert proc.returncode == 0, \ + f"This command failed: {job.cmd()}" diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/io_engine.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/io_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..aea66ec5587499dbc9ea899c1ed28607388d9ab1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/io_engine.py @@ -0,0 +1,126 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import time +from multiprocessing import Pool, Barrier + +from .ds_aio_constants import AIO_BASIC, TORCH_FAST_IO, TORCH_IO +from .test_ds_aio_utils import report_results, task_log, task_barrier +from .ds_aio_handle import AIOHandle_Engine +from .ds_aio_basic import AIOBasic_Engine +from .torch_io import TorchIO_Engine +from .torch_fastio_engine import Torch_FastIO_Engine + + +def prepare_operation(args, tid, read_op): + if args.engine == TORCH_IO: + io_engine = TorchIO_Engine(args, tid, read_op) + elif args.engine == AIO_BASIC: + io_engine = AIOBasic_Engine(args, tid, read_op) + elif args.engine == TORCH_FAST_IO: + io_engine = Torch_FastIO_Engine(args, tid, read_op) + else: + io_engine = AIOHandle_Engine(args, tid, read_op) + + return io_engine + + +def prepare_read(pool_params): + args, tid = pool_params + return prepare_operation(args, tid, True) + + +def prepare_write(pool_params): + args, tid = pool_params + return prepare_operation(args, tid, False) + + +def post_operation(pool_params): + _, _, io_engine = pool_params + io_engine.fini() + + +def read_operation(pool_params): + args, tid, loop_id, io_engine = pool_params + return io_engine.read(args, tid, loop_id) + + +def write_operation(pool_params): + args, tid, loop_id, io_engine = pool_params + return io_engine.write(args, tid, loop_id) + + +def get_schedule(args, read_op): + schedule = {} + if read_op: + schedule['pre'] = prepare_read + schedule['post'] = post_operation + schedule['main'] = read_operation + else: + schedule['pre'] = prepare_write + schedule['post'] = post_operation + schedule['main'] = write_operation + + return schedule + + +def io_engine_tasklet(pool_params): + args, tid, read_op = pool_params + num_processes = len(args.mapping_dict) + + # Create schedule + schedule = get_schedule(args, read_op) + task_log(tid, f'schedule = {schedule}') + task_barrier(aio_barrier, num_processes) + + # Run pre task + task_log(tid, f'running pre-task') + io_engine = schedule["pre"]((args, tid)) + task_barrier(aio_barrier, num_processes) + + # Run main tasks in a loop + io_engine.ctxt["main_task_sec"] = [] + for i in range(args.total_loops): + task_log(tid, f'running main task {i}') + start_time = time.time() + schedule["main"]((args, tid, i, io_engine)) + task_barrier(aio_barrier, num_processes) + stop_time = time.time() + io_engine.ctxt["main_task_sec"].append(stop_time - start_time) + + # Run post task + task_log(tid, f'running post-task') + schedule["post"]((args, tid, io_engine)) + task_barrier(aio_barrier, num_processes) + + ctxt = io_engine.ctxt + # return ctxt["main_task_sec"], ctxt["elapsed_sec"], ctxt["num_bytes"] * args.loops + if args.include_warmup_time: + e2e_latency_sec = sum(ctxt["main_task_sec"]) + task_latency_sec = sum(ctxt["elapsed_sec"]) + actual_loops = args.total_loops + else: + e2e_latency_sec = sum(ctxt["main_task_sec"][args.warmup_loops:]) + task_latency_sec = sum(ctxt["elapsed_sec"][args.warmup_loops:]) + actual_loops = args.loops + + l = ctxt["elapsed_sec"] + task_log(tid, f'task_latency_sec = {l}') + return e2e_latency_sec, task_latency_sec, ctxt["num_bytes"] * actual_loops + + +def _init_takslet(b): + global aio_barrier + aio_barrier = b + + +def io_engine_multiprocessing(args, read_op): + num_processes = len(args.mapping_dict) + b = Barrier(num_processes) + pool_params = [(args, p, read_op) for p in range(num_processes)] + with Pool(processes=num_processes, initializer=_init_takslet, initargs=(b, )) as p: + pool_results = p.map(io_engine_tasklet, pool_params) + + report_results(args, read_op, pool_results) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/parse_nvme_stats.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/parse_nvme_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..09c79ada5b3699676a5c6415bf30909e8686f548 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/parse_nvme_stats.py @@ -0,0 +1,148 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +import os +import argparse + +READ_SPEED = 'read_speed' +WRITE_SPEED = 'write_speed' + +PERF_METRICS = [READ_SPEED, WRITE_SPEED] + +METRIC_SEARCH = {READ_SPEED: 'E2E Read Speed', WRITE_SPEED: 'E2E Write Speed'} + + +def parse_arguments(): + parser = argparse.ArgumentParser() + + parser.add_argument('--log_dir', type=str, required=True, help='Folder of statistics logs') + + parser.add_argument('--metric', + type=str, + required=True, + help='Performance metric to report: [read_speed|write_speed]') + + args = parser.parse_args() + print(f'args = {args}') + + return args + + +def extract_value(key, file): + INVALID_PREFIXES = ["ds"] + for p in INVALID_PREFIXES: + if key.startswith(p): + return key + try: + if key[0] in ['t', 'd', 'p']: + return int(key[1:]) + if key.startswith("bs"): + if key.endswith('K'): + v = key[2:].split('K') + return int(v[0]) * 1024 + elif key.endswith('M'): + v = key[2:].split('M') + return int(v[0]) * 1024 * 1024 + else: + return int(key[2:]) + except: + print(f"{file}: extract_value fails on {key}") + return None + + return key + + +def get_file_key(file): + f, _ = os.path.splitext(os.path.basename(file)) + fields = f.split('_') + values = [extract_value(k, file) for k in fields] + return tuple(values) + + +def get_thread_count(file): + f, _ = os.path.splitext(os.path.basename(file)) + fields = f.split('_') + for key in fields: + if key[0] == 't': + return int(key[1:]) + return 1 + + +""" +Extract performance metric from log file. +Sample file lines are: +Task Read Latency = 0.031647682189941406 sec +Task Read Speed = 12.342926020792527 GB/sec +E2E Read Latency = 0.031697988510131836 sec +E2E Read Speed = 12.323337169333062 GB/sec + +For the above sample, -metric = "read_speed" corresponds to "E2E Read Speed", and 12.32 will be returned +""" + + +def get_metric(file, metric): + thread_count = get_thread_count(file) + with open(file) as f: + for line in f.readlines(): + if line.startswith(METRIC_SEARCH[metric]): + if metric in [READ_SPEED, WRITE_SPEED]: + fields = line.split() + return float(fields[-2]) + else: + fields = line.split('=') + return float(fields[-1]) + + return None + + +def validate_args(args): + if not args.metric in PERF_METRICS: + print(f'{args.metric} is not a valid performance metrics') + return False + + if not os.path.isdir(args.log_dir): + print(f'{args.log_dir} folder is not existent') + return False + + return True + + +def get_results(log_files, metric): + results = {} + for f in log_files: + file_key = get_file_key(f) + value = get_metric(f, metric) + results[file_key] = value + + return results + + +def get_sorted_results(log_dir, metric): + log_files = [f for f in os.listdir(log_dir) if os.path.isfile(os.path.join(log_dir, f))] + + log_files_path = [os.path.join(log_dir, f) for f in log_files] + results = get_results(log_files_path, metric) + result_keys = list(results.keys()) + sorted_keys = sorted(result_keys) + return sorted_keys, results + + +def main(): + print("Parsing aio statistics") + args = parse_arguments() + + if not validate_args(args): + quit() + + sorted_keys, results = get_sorted_results(args.log_dir, args.metric) + for k in sorted_keys: + print(f'{k} = {results[k]}') + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_generate_param.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_generate_param.py new file mode 100644 index 0000000000000000000000000000000000000000..99be75c2d9190637ee43d49d48cf7a911eae6fbf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_generate_param.py @@ -0,0 +1,101 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" +import os +import argparse +import json +from .parse_nvme_stats import READ_SPEED, WRITE_SPEED, get_sorted_results +from .perf_sweep_utils import BENCH_LOG_DIR, READ_LOG_DIR, WRITE_LOG_DIR + + +def parse_arguments(): + parser = argparse.ArgumentParser() + + parser.add_argument('--log_dir', + type=str, + default=BENCH_LOG_DIR, + help=f'Folder of performance sweep logs. Default is {os.path.join(".", BENCH_LOG_DIR)}') + parser.add_argument('--verbose', action='store_true', help='Print debugging information.') + + args = parser.parse_args() + if args.verbose: + print(f'args = {args}') + + return args + + +def validate_args(args): + for d in [READ_LOG_DIR, WRITE_LOG_DIR]: + log_dir = os.path.join(args.log_dir, d) + if not os.path.isdir(log_dir): + print(f'{log_dir} folder is not existent') + return False + + return True + + +def convert_to_param(key): + assert len(key) == 6 + return { + "single_submit": "true" if key[0] == "single" else "false", + "overlap_events": "true" if key[1] == "overlap" else "false", + "num_threads": int(key[5]), + "queue_depth": int(key[3]), + "block_size": int(key[4]) + } + + +def generate_aio_param(read_log_dir, write_log_dir): + _, read_results = get_sorted_results(read_log_dir, READ_SPEED) + _, write_results = get_sorted_results(write_log_dir, WRITE_SPEED) + + read_results_count = len(read_results.items()) + write_results_count = len(write_results.items()) + assert read_results_count == write_results_count, f"Mismatch in number of read & write results: {read_results_count=} != {write_results_count=}" + + combined_perf = {key[1:]: value for key, value in read_results.items()} + for key, value in write_results.items(): + new_key = key[1:] + if new_key in combined_perf: + combined_perf[new_key] += value + else: + combined_perf[new_key] = 0 + + optimal_key = None + optimal_perf = 0.0 + for key, value in combined_perf.items(): + if value > optimal_perf: + optimal_perf = value + optimal_key = key + + aio_param = {"aio": convert_to_param(optimal_key)} + + read_perf_keys = {key[1:]: key for key in read_results.keys()} + write_perf_keys = {key[1:]: key for key in write_results.keys()} + optimal_config_read = read_results.get(read_perf_keys[optimal_key], None) + optimal_config_write = write_results.get(write_perf_keys[optimal_key], None) + + print(f'Best performance (GB/sec): read = {optimal_config_read:5.2f}, write = {optimal_config_write:5.2f}') + print(json.dumps(aio_param, indent=3)) + + +def generate_main(log_dir): + read_log_dir = os.path.join(log_dir, READ_LOG_DIR) + write_log_dir = os.path.join(log_dir, WRITE_LOG_DIR) + generate_aio_param(read_log_dir, write_log_dir) + + +def main(): + args = parse_arguments() + if not validate_args(args): + quit() + print(f'Generate DeepNVMe configuration from {args.log_dir} logs') + generate_main(args.log_dir) + + +if __name__ == "__main__": + generate_main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_run_sweep.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_run_sweep.py new file mode 100644 index 0000000000000000000000000000000000000000..4560911bcb98d6f4a7af6911ea750738e02377cf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_run_sweep.py @@ -0,0 +1,319 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" +import os +import sys +import argparse +import json +import itertools +import shutil + +from deepspeed.ops.op_builder import AsyncIOBuilder, GDSBuilder +from .ds_aio_job import Job, run_job +from .perf_sweep_utils import READ_OP_DESC, WRITE_OP_DESC, BENCH_LOG_DIR, \ + READ_LOG_DIR, WRITE_LOG_DIR + +OTHER_OPTIONS = '--engine aio_handle' +PERF_SCRIPT = 'ds_io' +DEFAULT_SWEEP_CONFIG = { + "block_size": ["1M", "8M"], + "queue_depth": [32, 128], + "sequential_requests": [True, False], + "single_submit": [False, True], + "io_parallel": [1, 2, 4, 8], +} + + +class SweepConfig(object): + + def __init__(self, args): + self.folder_to_device_mapping = get_ftd_map(args.nvme_dir) + self.search_space = get_sweep_config_dict(args.sweep_config) + self.search_space.update(self.folder_to_device_mapping) + self.read = not args.no_read + self.write = not args.no_write + self.flush_cache = args.flush_page_cache + self.log_dir = args.log_dir + self.verbose = args.verbose + self.other_options = f'{OTHER_OPTIONS} --loops {args.loops} --io_size {args.io_size}' + if args.gpu: + self.other_options += ' --gpu' + if args.gds: + self.other_options += ' --use_gds' + + +def validate_arguments(args): + if not async_io_setup(): + error_msg = """ + Failing because environment is not properly configured for deepspeed async i/o module. + Possible fix: apt install libaio-dev. + """ + print(error_msg) + quit() + + if args.gds and not gds_io_setup(): + error_msg = """ + Failing because environment is not properly configured for deepspeed GDS I/O operator. + """ + print(error_msg) + quit() + + +def parse_sweep_arguments(): + parser = argparse.ArgumentParser() + + parser.add_argument('--nvme_dir', + nargs='+', + required=True, + help='Directory in which to perform I/O tests. A writeable directory on a NVMe device.') + + parser.add_argument('--sweep_config', type=str, default=None, help='Performance sweep configuration json file.') + + parser.add_argument('--no_read', action='store_true', help='Disable read performance measurements.') + + parser.add_argument('--no_write', action='store_true', help='Disable write performance measurements.') + + parser.add_argument('--io_size', + type=str, + default="400M", + help='Number of I/O bytes to read/write for performance measurements.') + + parser.add_argument('--gpu', action='store_true', help='Test tensor transfers between GPU device and NVME device.') + + parser.add_argument('--gds', action='store_true', help='Run the sweep over NVIDIA GPUDirectStorage operator') + + parser.add_argument( + '--flush_page_cache', + action='store_true', + help= + 'Page cache will not be flushed and reported read speeds may be higher than actual ***Requires sudo access***.' + ) + + parser.add_argument( + '--log_dir', + type=str, + default=BENCH_LOG_DIR, + help=f'Output directory for performance log files. Default is {os.path.join(".", BENCH_LOG_DIR)}') + + parser.add_argument('--loops', type=int, default=1, help='Count of operation repetitions') + + parser.add_argument('--verbose', action='store_true', help='Print debugging information.') + + args = parser.parse_args() + if args.verbose: + print(f'args = {args}') + validate_arguments(args) + + return args + + +def dump_cmd_lines(cmd_lines): + print(f'cmd line count = {len(cmd_lines)}') + for i, cmd in enumerate(cmd_lines): + print(f'{i}: {cmd}') + + +def get_ftd_map(nvme_dir_list): + ftd_list = [f'{dir}:{dev}' for dev, dir in enumerate(nvme_dir_list)] + ftd_arg = [' '.join(ftd for ftd in ftd_list)] + return {'folder_to_device_mapping': ftd_arg} + + +def get_sweep_config_dict(sweep_config_json): + if sweep_config_json is None: + return DEFAULT_SWEEP_CONFIG + + with open(sweep_config_json) as fp: + sweep_config = json.load(fp) + return sweep_config + + +def get_sweep_cmd_lines(sweep_config_dict): + + def flatten_options(key, value_list): + flat_list = [] + for v in value_list: + if not type(v) is bool: + flat_list.append(f'--{key} {v}') + elif v: + flat_list.append(f'--{key}') + else: + flat_list.append(' ') + + return flat_list + + flat_list = [flatten_options(key, value) for key, value in sweep_config_dict.items()] + cmd_list = list(itertools.product(*flat_list)) + cmd_list = [list(cmd) for cmd in cmd_list] + #dump_cmd_lines(cmd_list) + return cmd_list + + +def launch_sweep(sweep_jobs, sync_job, flush_cache_job, verbose): + for perf_job in sweep_jobs: + if flush_cache_job is not None: + run_job(sync_job, verbose) + run_job(flush_cache_job, verbose) + + run_job(perf_job, verbose) + + run_job(sync_job, verbose) + + +def create_cmd_tags(cmd_line): + tags = {} + for param_value in cmd_line: + fields = param_value.split() + if len(fields) == 1: + tags[fields[0]] = None + elif len(fields) == 2: + if fields[0] == '--folder_to_device_mapping': + tags[fields[0]] = len(fields[1:]) + else: + tags[fields[0]] = fields[1] + elif len(fields) > 2: + tags[fields[0]] = len(fields[1:]) + return tags + + +def get_log_file(io_op_desc, cmd_line): + QUEUE_DEPTH = "--queue_depth" + BLOCK_SIZE = "--block_size" + SINGLE_SUBMIT = "--single_submit" + SEQUENTIAL_REQUESTS = "--sequential_requests" + FTD_MAP = "--folder_to_device_mapping" + IO_PARALLEL = "--io_parallel" + + tag_map = { + QUEUE_DEPTH: "d", + BLOCK_SIZE: "bs", + SINGLE_SUBMIT: "single", + SEQUENTIAL_REQUESTS: "sequential", + FTD_MAP: "ftd", + IO_PARALLEL: "p" + } + + tag_default = { + QUEUE_DEPTH: 1, + BLOCK_SIZE: "1M", + SINGLE_SUBMIT: "block", + SEQUENTIAL_REQUESTS: "overlap", + FTD_MAP: 1, + IO_PARALLEL: 1 + } + + def get_default_value(tag): + value = tag_default[tag] + if tag in [SINGLE_SUBMIT, SEQUENTIAL_REQUESTS]: + return value + return f'{tag_map[tag]}{value}' + + def get_config_value(tag, value): + tag_key = tag_map[tag] + if value is None: + return tag_key + return f'{tag_key}{value}' + + tag_list = [SINGLE_SUBMIT, SEQUENTIAL_REQUESTS, FTD_MAP, QUEUE_DEPTH, BLOCK_SIZE, IO_PARALLEL] + log_tags = [io_op_desc] + cmd_tags = create_cmd_tags(cmd_line) + for tag in tag_list: + if tag in cmd_tags: + log_tags.append(get_config_value(tag, cmd_tags[tag])) + else: + log_tags.append(get_default_value(tag)) + + log_file = '_'.join(log_tags) + log_file += '.txt' + return log_file + + +def create_perf_jobs(io_op_desc, log_dir, cmd_lines): + py_cmd = [os.path.join(script_path(), PERF_SCRIPT)] + + perf_jobs = [] + for cmd in cmd_lines: + log_file = os.path.join(log_dir, get_log_file(io_op_desc, cmd)) + job = Job(cmd_line=py_cmd + cmd, output_file=log_file) + perf_jobs.append(job) + + return perf_jobs + + +def script_path(): + return os.path.dirname(os.path.realpath(sys.argv[0])) + + +def async_io_setup(): + return AsyncIOBuilder().is_compatible() + + +def gds_io_setup(): + return GDSBuilder().is_compatible() + + +def remove_folder(folder): + assert os.path.isdir(folder), f"Error: cannot remove {folder} - folder not found" + shutil.rmtree(folder) + + +def run_read_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines): + read_cmd_lines = [[f'--read {sweep_config.other_options}'] + cmd for cmd in cmd_lines] + # dump_cmd_lines(cmd_lines) + + log_folder = os.path.join(sweep_config.log_dir, f'{READ_LOG_DIR}') + os.makedirs(log_folder, exist_ok=True) + + perf_jobs = create_perf_jobs(io_op_desc=READ_OP_DESC, log_dir=log_folder, cmd_lines=read_cmd_lines) + + launch_sweep(sweep_jobs=perf_jobs, + sync_job=sync_job, + flush_cache_job=flush_cache_job, + verbose=sweep_config.verbose) + + +def run_write_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines): + write_cmd_lines = [[f'{sweep_config.other_options}'] + cmd for cmd in cmd_lines] + # dump_cmd_lines(write_cmd_lines) + + log_folder = os.path.join(sweep_config.log_dir, f'{WRITE_LOG_DIR}') + os.makedirs(log_folder, exist_ok=True) + + perf_jobs = create_perf_jobs(io_op_desc=WRITE_OP_DESC, log_dir=log_folder, cmd_lines=write_cmd_lines) + + launch_sweep(sweep_jobs=perf_jobs, + sync_job=sync_job, + flush_cache_job=flush_cache_job, + verbose=sweep_config.verbose) + + +def sweep_main(args): + sweep_config = SweepConfig(args) + cmd_lines = get_sweep_cmd_lines(sweep_config.search_space) + + if sweep_config.flush_cache: + flush_cache_job = Job(cmd_line=['sudo', 'bash -c', "'echo 1 > /proc/sys/vm/drop_caches'"]) + else: + flush_cache_job = None + + sync_job = Job(cmd_line=['sync']) + + if sweep_config.read: + run_read_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines) + + if sweep_config.write: + run_write_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines) + + +def main(): + args = parse_sweep_arguments() + print(f"Running DeepNVMe performance sweep on {args.nvme_dir}") + sweep_main(args) + + +if __name__ == "__main__": + sweep_main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_sweep_utils.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_sweep_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e6832c1baa492d8dfc695ac251a606e49db82f35 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_sweep_utils.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +SCRIPT_PREFIX = '_aio_bench' +WRITE_OP_DESC = 'write' +READ_OP_DESC = 'read' +READ_IO_DIR = f'{SCRIPT_PREFIX}_{READ_OP_DESC}_io' +WRITE_IO_DIR = f'{SCRIPT_PREFIX}_{WRITE_OP_DESC}_io' +BENCH_LOG_DIR = f'{SCRIPT_PREFIX}_logs' +READ_LOG_DIR = f'{SCRIPT_PREFIX}_{READ_OP_DESC}_logs' +WRITE_LOG_DIR = f'{SCRIPT_PREFIX}_{WRITE_OP_DESC}_logs' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio.py new file mode 100644 index 0000000000000000000000000000000000000000..a71ad3ee97d373e6feda2f43299ab31c51aec56c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +import multiprocessing as mp +from .ds_aio_args import get_validated_args +from .io_engine import io_engine_multiprocessing + + +def ds_io_main(): + print(f'Testing DeepNVMe python frontend') + + args = get_validated_args() + mp.set_start_method('spawn', force=True) + multiprocess_function = io_engine_multiprocessing + multiprocess_function(args, args.read) + + +if __name__ == "__main__": + ds_io_main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio_utils.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..90b994b0b5322893420037b442aa4763ea5c4af1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio_utils.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +import os +from .ds_aio_job import Job, run_job +import torch +from deepspeed.accelerator import get_accelerator + +BYTES_PER_GB = 1024**3 +BYTES_PER_MB = 1024**2 +BYTES_PER_KB = 1024 +LOG_TIDS = [0] + + +def task_log(tid, msg, force=False): + if force or tid in LOG_TIDS: + print(f'tid {tid}: {msg}') + + +def task_barrier(barrier, num_parties): + assert barrier.parties == num_parties + barrier.wait() + assert barrier.broken == False + + +def report_results(args, read_op, pool_results): + #print(f'pool_results = {pool_results}') + io_string = 'Read' if read_op else 'Write' + if None in pool_results: + print(f'Failure in one of {args.threads} {io_string} processes') + return + + total_bytes = sum([num_bytes for _, _, num_bytes in pool_results]) + + task_latency_sec = max([sec for _, sec, _ in pool_results]) + task_speed_GB = 0 if task_latency_sec == 0 else total_bytes / task_latency_sec / BYTES_PER_GB + print(f'Task {io_string} Latency = {task_latency_sec} sec') + print(f'Task {io_string} Speed = {task_speed_GB} GB/sec') + + e2e_latency_sec = max([sec for sec, _, _ in pool_results]) + e2e_speed_GB = 0 if e2e_latency_sec == 0 else total_bytes / e2e_latency_sec / BYTES_PER_GB + print(f'E2E {io_string} Latency = {e2e_latency_sec} sec') + print(f'E2E {io_string} Speed = {e2e_speed_GB} GB/sec') + + +def get_block_size_and_count(io_bytes): + if io_bytes > BYTES_PER_MB and io_bytes % BYTES_PER_MB == 0: + block_size = BYTES_PER_MB + block_size_string = '1M' + else: + assert io_bytes % BYTES_PER_KB == 0 + block_size = BYTES_PER_KB + block_size_string = '1K' + block_count = io_bytes / block_size + + return block_size_string, int(block_count) + + +def refine_integer_value(value): + unit_dict = {'K': 1024, 'M': 1024**2, 'G': 1024**3} + + if value[-1] in list(unit_dict.keys()): + int_value = int(value[:-1]) * unit_dict[value[-1]] + return int_value + return int(value) + + +def create_filename(folder, read_op, size, tid): + io_string = "read" if read_op else "write" + return os.path.join(folder, f'_aio_{io_string}_{size}.pt.{tid}') + + +def create_file(filename, num_bytes): + block_size, block_count = get_block_size_and_count(num_bytes) + dd_job = Job(cmd_line=[f'dd if=/dev/urandom of={filename} bs={block_size} count={block_count}']) + print(f'[Start] Create {filename} of {num_bytes} bytes by running {dd_job.cmd()} ....') + run_job(dd_job) + print(f'[Done] Create read file of {num_bytes} bytes by running {dd_job.cmd()} ....') + + +def create_page_locked_tensor(num_elem, use_accelerator, aio_handle=None): + if use_accelerator: + return get_accelerator().pin_memory(torch.randint(high=128, size=(num_elem, ), dtype=torch.uint8, + device='cpu')) + else: + return aio_handle.new_cpu_locked_tensor(num_elem, torch.empty(0, dtype=torch.uint8)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_fastio_engine.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_fastio_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..fd4918d84c396bf7d14827c42a6b338bed5e9f63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_fastio_engine.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import os +import time +from deepspeed.ops.aio import AsyncIOBuilder +from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor +from .ds_aio_constants import * +from deepspeed.io import FastFileWriter + + +class Torch_FastIO_Engine(object): + + def __init__(self, args, tid, read_op): + assert read_op is False, f'Read operation is not currently supported' + self.ctxt = self._create_context(args, tid, read_op) + self.zipfile_serialization = not args.torch_legacy_save + + def fini(self): + if self.ctxt[USE_CPU_LOCKED_TENSOR]: + for buf in [BUFFER, FAST_IO_BUFFER]: + self.ctxt[HANDLE].free_cpu_locked_tensor(self.ctxt[buf]) + + self.ctxt[BUFFER].detach() + self.ctxt[BUFFER] = None + + def read(self, args, tid): + start_time = time.time() + torch.load(f=self.ctxt[FILE], map_location=self.ctxt[BUFFER].device) + end_time = time.time() + self.ctxt[ELAPSED_SEC] += end_time - start_time + + def write(self, args, tid): + # Avoid overwriting existing files as it could be artificially faster + if os.path.isfile(self.ctxt[FILE]): + os.remove(self.ctxt[FILE]) + + ds_file_writer = FastFileWriter(file_path=self.ctxt[FILE], + aio_handle=self.ctxt[HANDLE], + pinned_tensor=self.ctxt[FAST_IO_BUFFER]) + + start_time = time.time() + torch.save(obj=self.ctxt[BUFFER], f=ds_file_writer, _use_new_zipfile_serialization=self.zipfile_serialization) + ds_file_writer.close() # Force flush to storage + end_time = time.time() + self.ctxt[ELAPSED_SEC] += end_time - start_time + ds_file_writer._dump_state() + + def _create_context(self, args, tid, read_op): + io_string = "Read" if read_op else "Write" + device_id, folder = args.mapping_list[tid] + filename = create_filename(folder, args.read, args.io_size, tid) + if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): + create_file(filename, args.io_size) + + io_parallel = args.io_parallel if args.io_parallel else 1 + aio_handle = AsyncIOBuilder().load().aio_handle(args.block_size, args.queue_depth, args.single_submit, + not args.sequential_requests, io_parallel) + + if args.gpu: + buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=f'cuda:{device_id}') + else: + buffer = create_page_locked_tensor(args.io_size, args.use_accelerator_pin_memory, aio_handle) + + task_log(tid, f'Allocate tensor of size {args.io_size} bytes') + + fast_io_buffer = create_page_locked_tensor(args.fast_io_size, args.use_accelerator_pin_memory, aio_handle) + + task_log(tid, f'created torch_fastio engine') + + ctxt = {} + ctxt[FILE] = filename + ctxt[NUM_BYTES] = args.io_size + ctxt[BUFFER] = buffer + ctxt[HANDLE] = aio_handle + ctxt[FAST_IO_BUFFER] = fast_io_buffer + ctxt[ELAPSED_SEC] = 0 + ctxt[USE_CPU_LOCKED_TENSOR] = not args.use_accelerator_pin_memory + + task_log(tid, + f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}', + force=True) + + return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_io.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_io.py new file mode 100644 index 0000000000000000000000000000000000000000..3371771c02f031d1612e55bfce6acc44003131c8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_io.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import os +import time +from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor +from .ds_aio_constants import * + + +class TorchIO_Engine(object): + + def __init__(self, args, tid, read_op): + self.ctxt = self._create_context(args, tid, read_op) + self.zipfile_serialization = not args.torch_legacy_save + + def fini(self): + self.ctxt[BUFFER].detach() + self.ctxt[BUFFER] = None + + def read(self, args, tid): + start_time = time.time() + torch.load(f=self.ctxt[FILE], map_location=self.ctxt[BUFFER].device) + end_time = time.time() + self.ctxt[ELAPSED_SEC] += end_time - start_time + + def write(self, args, tid): + # Avoid overwriting existing files as it could be artificially faster + if os.path.isfile(self.ctxt[FILE]): + os.remove(self.ctxt[FILE]) + + start_time = time.time() + torch.save(obj=self.ctxt[BUFFER], f=self.ctxt[FILE], _use_new_zipfile_serialization=self.zipfile_serialization) + end_time = time.time() + self.ctxt[ELAPSED_SEC] += end_time - start_time + + def _create_context(self, args, tid, read_op): + io_string = "Read" if read_op else "Write" + device_id, folder = args.mapping_list[tid] + filename = create_filename(folder, args.read, args.io_size, tid) + if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): + create_file(filename, args.io_size) + + task_log(tid, f'Allocate tensor of size {args.io_size} bytes') + + if args.gpu: + buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=f'cuda:{device_id}') + else: + buffer = create_page_locked_tensor(args.io_size, True) + + task_log(tid, + f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}', + force=True) + + task_log(tid, f'created torch_io engine') + + ctxt = {} + ctxt[FILE] = filename + ctxt[NUM_BYTES] = args.io_size + ctxt[BUFFER] = buffer + ctxt[ELAPSED_SEC] = 0 + return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/validate_async_io.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/validate_async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..10fb638347bcd9132abc365ba8f488551ea7cb01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/nvme/validate_async_io.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" +from deepspeed.ops.op_builder import AsyncIOBuilder +assert AsyncIOBuilder().is_compatible() +assert AsyncIOBuilder().load() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..15179984173c1f7f8bcdaaeb98037aa8a05ce1e9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from . import adam +from . import adagrad +from . import lamb +from . import lion +from . import sparse_attention +from . import transformer +from . import fp_quantizer +from .transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig + +from ..git_version_info import compatible_ops as __compatible_ops__ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..493ed737a4d90c705ae7055709e91f69a6bc24a8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5bf2e4721cd61ec3f548d652698f62b7a50772d5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .cpu_adagrad import DeepSpeedCPUAdagrad diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba249042d02572d83565a2c0d2a7f8a15847f0ab Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70a2d4aaffa1d314ef91f6f849c2cdee77c1ef8d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..dbde6d95f6520eca20b91d3b9328b741b99df074 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/cpu_adagrad.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.ops.op_builder import CPUAdagradBuilder +from deepspeed.utils.logging import should_log_le + + +class DeepSpeedCPUAdagrad(torch.optim.Optimizer): + optimizer_id = 0 + + def __init__(self, model_params, lr=1e-2, eps=1e-10, weight_decay=0, amsgrad=False, fp32_optimizer_states=True): + + default_args = dict(lr=lr, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad) + super(DeepSpeedCPUAdagrad, self).__init__(model_params, default_args) + + self.opt_id = DeepSpeedCPUAdagrad.optimizer_id + DeepSpeedCPUAdagrad.optimizer_id = DeepSpeedCPUAdagrad.optimizer_id + 1 + self.fp32_optimizer_states = fp32_optimizer_states + self.ds_opt_adagrad = CPUAdagradBuilder().load() + + self.ds_opt_adagrad.create_adagrad(self.opt_id, lr, eps, weight_decay, should_log_le("info")) + + def __del__(self): + # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize + # is used multiple times in the same process (notebook or pytest worker) + self.ds_opt_adagrad.destroy_adagrad(self.opt_id) + + def __setstate__(self, state): + super(DeepSpeedCPUAdagrad, self).__setstate__(state) + for group in self.param_groups: + group.setdefault('amsgrad', False) + + @torch.no_grad() + def step(self, closure=None): + """Update the model parameters. + + .. note:: + This method will be called internally by ZeRO-Offload. DeepSpeed + users should still use ``engine.step()`` as shown in the + `Getting Started + `_ guide. + + Args: + closure (callable, optional): closure to compute the loss. + Defaults to ``None``. + + Returns: + loss: if ``closure`` is provided. Otherwise ``None``. + """ + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # intended device for step + device = torch.device('cpu') + + for group_id, group in enumerate(self.param_groups): + for param_id, p in enumerate(group['params']): + + if p.grad is None: + continue + + assert p.device == device, f"CPUAdagrad param is on {p.device} and must be 'cpu', make " \ + "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config." + + state = self.state[p] + # State initialization + if len(state) == 0: + #print(f'group {group_id} param {param_id} = {p.numel()}') + state['step'] = 0 + + #use full precision by default unless self.fp32_optimizer_states is off + state_dtype = torch.float if self.fp32_optimizer_states else p.dtype + + #memory_format=torch.preserve_format) + # gradient variances + state['exp_avg_sq'] = torch.zeros_like(p.data, dtype=state_dtype, device='cpu') + #memory_format=torch.preserve_format) + + state['step'] += 1 + + if p.grad.is_sparse == True: + sparse_param = p.sparse_mask(p.grad) + sparse_exp_avg_sq = state['exp_avg_sq'].sparse_mask(p.grad) + self.ds_opt_adagrad.adagrad_update(self.opt_id, state['step'], group['lr'], group['eps'], + group['weight_decay'], sparse_param.values(), p.grad.values(), + sparse_exp_avg_sq.values()) + p[sparse_param.indices()] = sparse_param.values() + state['exp_avg_sq'][sparse_exp_avg_sq.indices()] = sparse_exp_avg_sq.values() + else: + self.ds_opt_adagrad.adagrad_update(self.opt_id, state['step'], group['lr'], group['eps'], + group['weight_decay'], p.data, p.grad.data, state['exp_avg_sq']) + return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a29bb9447d015d0992c18624dba41a15aa838866 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .cpu_adam import DeepSpeedCPUAdam +from .fused_adam import FusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88219ad6802c24c4a0dac7984770a9a7a37afb7e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..853b9f5d7821004d5872f84305e92f0d1137b174 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7072b45d80e0aad20ad6dfd3febe3840b1674898 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebc7312343790b7d86fc210c625334972a213bf8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..e0a72a4942577b4db2ac822ee0675a084df1b2bc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py @@ -0,0 +1,166 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from cpuinfo import get_cpu_info +from deepspeed.utils import logger +from deepspeed.utils.logging import should_log_le +from deepspeed.ops.op_builder import CPUAdamBuilder + + +class DeepSpeedCPUAdam(torch.optim.Optimizer): + optimizer_id = 0 + + def __init__(self, + model_params, + lr=1e-3, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + amsgrad=False, + adamw_mode=True, + fp32_optimizer_states=True): + """Fast vectorized implementation of two variations of Adam optimizer on CPU: + + * Adam: A Method for Stochastic Optimization: (https://arxiv.org/abs/1412.6980); + * AdamW: Fixing Weight Decay Regularization in Adam (https://arxiv.org/abs/1711.05101) + + DeepSpeed CPU Adam(W) provides between 5x to 7x speedup over torch.optim.adam(W). + In order to apply this optimizer, the model requires to have its master parameter (in FP32) + reside on the CPU memory. + + To train on a heterogeneous system, such as coordinating CPU and GPU, DeepSpeed offers + the ZeRO-Offload technology which efficiently offloads the optimizer states into CPU memory, + with minimal impact on training throughput. DeepSpeedCPUAdam plays an important role to minimize + the overhead of the optimizer's latency on CPU. Please refer to ZeRO-Offload tutorial + (https://www.deepspeed.ai/tutorials/zero-offload/) for more information on how to enable this technology. + + For calling step function, there are two options available: (1) update optimizer's states and (2) update + optimizer's states and copy the parameters back to GPU at the same time. We have seen that the second + option can bring 30% higher throughput than the doing the copy separately using option one. + + + .. note:: + We recommend using our `config + `_ + to allow :meth:`deepspeed.initialize` to build this optimizer + for you. + + + Arguments: + model_params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in DeepSpeed CPUAdam! + adamw_mode: select between Adam and AdamW implementations (default: AdamW) + fp32_optimizer_states: creates momentum and variance in full precision regardless of + the precision of the parameters (default: True) + """ + + default_args = dict(lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + bias_correction=bias_correction, + amsgrad=amsgrad) + super(DeepSpeedCPUAdam, self).__init__(model_params, default_args) + + cpu_info = get_cpu_info() + self.cpu_vendor = cpu_info["vendor_id_raw"].lower() if "vendor_id_raw" in cpu_info else "unknown" + if "amd" in self.cpu_vendor: + for group_id, group in enumerate(self.param_groups): + for param_id, p in enumerate(group['params']): + if p.dtype == torch.half: + logger.warning("FP16 params for CPUAdam may not work on AMD CPUs") + break + else: + continue + break + + self.opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id = DeepSpeedCPUAdam.optimizer_id + 1 + self.adam_w_mode = adamw_mode + self.fp32_optimizer_states = fp32_optimizer_states + self.ds_opt_adam = CPUAdamBuilder().load() + + self.ds_opt_adam.create_adam(self.opt_id, lr, betas[0], betas[1], eps, weight_decay, adamw_mode, + should_log_le("info")) + + def __del__(self): + # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize + # is used multiple times in the same process (notebook or pytest worker) + self.ds_opt_adam.destroy_adam(self.opt_id) + + def __setstate__(self, state): + super(DeepSpeedCPUAdam, self).__setstate__(state) + for group in self.param_groups: + group.setdefault('amsgrad', False) + + @torch.no_grad() + def step(self, closure=None): + """Update the model parameters. + + .. note:: + This method will be called internally by ZeRO-Offload. DeepSpeed + users should still use ``engine.step()`` as shown in the + `Getting Started + `_ guide. + + Args: + closure (callable, optional): closure to compute the loss. + Defaults to ``None``. + + Returns: + loss: if ``closure`` is provided. Otherwise ``None``. + """ + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # intended device for step + device = torch.device('cpu') + + for group_id, group in enumerate(self.param_groups): + for param_id, p in enumerate(group['params']): + + if p.grad is None: + continue + + assert p.device == device, f"CPUAdam param is on {p.device} and must be 'cpu', make " \ + "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config." + + state = self.state[p] + # State initialization + if len(state) == 0: + #print(f'group {group_id} param {param_id} = {p.numel()}') + state['step'] = 0 + + #use full precision by default unless self.fp32_optimizer_states is off + state_dtype = torch.float if self.fp32_optimizer_states else p.dtype + + # gradient momentums + state['exp_avg'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) + #memory_format=torch.preserve_format) + # gradient variances + state['exp_avg_sq'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) + #memory_format=torch.preserve_format) + + state['step'] += 1 + beta1, beta2 = group['betas'] + + self.ds_opt_adam.adam_update(self.opt_id, state['step'], group['lr'], beta1, beta2, group['eps'], + group['weight_decay'], group['bias_correction'], p.data, p.grad.data, + state['exp_avg'], state['exp_avg_sq']) + return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..53f859e9cc87bde8f16760a4c23394ea1207af0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/fused_adam.py @@ -0,0 +1,195 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit 6bd01c4 +""" + +import torch +from .multi_tensor_apply import MultiTensorApply + +multi_tensor_applier = MultiTensorApply(2048 * 32) +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import FusedAdamBuilder + + +class FusedAdam(torch.optim.Optimizer): + """Implements Adam algorithm. + + Currently GPU-only. Requires Apex to be installed via + ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``. + + This version of fused Adam implements 2 fusions. + + * Fusion of the Adam update's elementwise operations + * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches. + + :class:`apex.optimizers.FusedAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``, + or ``torch.optim.Adam`` with ``adam_w_mode=False``:: + + opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....) + ... + opt.step() + + :class:`apex.optimizers.FusedAdam` may be used with or without Amp. If you wish to use :class:`FusedAdam` with Amp, + you may choose any ``opt_level``:: + + opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....) + model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2") + ... + opt.step() + + In general, ``opt_level="O1"`` is recommended. + + + .. warning:: + A previous version of :class:`FusedAdam` allowed a number of additional arguments to ``step``. These additional arguments + are now deprecated and unnecessary. + + Adam was been proposed in `Adam: A Method for Stochastic Optimization`_. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in FusedAdam! + adam_w_mode (boolean, optional): Apply L2 regularization or weight decay + True for decoupled weight decay(also known as AdamW) (default: True) + set_grad_none (bool, optional): whether set grad to None when zero_grad() + method is called. (default: True) + + .. _Adam - A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + lr=1e-3, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + adam_w_mode=True, + weight_decay=0., + amsgrad=False, + set_grad_none=True): + + if amsgrad: + raise RuntimeError('FusedAdam does not support the AMSGrad variant.') + defaults = dict(lr=lr, bias_correction=bias_correction, betas=betas, eps=eps, weight_decay=weight_decay) + super(FusedAdam, self).__init__(params, defaults) + self.adam_w_mode = 1 if adam_w_mode else 0 + self.set_grad_none = set_grad_none + + fused_adam_cuda = FusedAdamBuilder().load() + # Skip buffer + self._dummy_overflow_buf = get_accelerator().IntTensor([0]) + self.multi_tensor_adam = fused_adam_cuda.multi_tensor_adam + + def zero_grad(self): + if self.set_grad_none: + for group in self.param_groups: + for p in group['params']: + p.grad = None + else: + super(FusedAdam, self).zero_grad() + + def step(self, closure=None, grads=None, output_params=None, scale=None, grad_norms=None, grad_scaler=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + + The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes. + """ + if any(p is not None for p in [grads, output_params, scale, grad_norms]): + raise RuntimeError( + 'FusedAdam has been updated. Simply initialize it identically to torch.optim.Adam, and call step() with no arguments.' + ) + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + if len(group['params']) == 0: + continue + bias_correction = 1 if group['bias_correction'] else 0 + beta1, beta2 = group['betas'] + + # assume same step across group now to simplify things + # per parameter step can be easily support by making it tensor, or pass list into kernel + if 'step' not in group: + group['step'] = 0 + + # create lists for multi-tensor apply + g_16, p_16, m_16, v_16 = [], [], [], [] + g_bf, p_bf, m_bf, v_bf = [], [], [], [] + g_32, p_32, m_32, v_32 = [], [], [], [] + + for p in group['params']: + if p.grad is None: + continue + if p.grad.data.is_sparse: + raise RuntimeError( + 'FusedAdam does not support sparse gradients, please consider SparseAdam instead') + + state = self.state[p] + # State initialization + if len(state) == 0: + # DeepSpeed ZeRO 3 processes each subgroup a time, so we need to keep tracking step count for each tensor separately. + # While this is not an issue for ZeRO 1 & 2, since they apply a single optimization step to the whole param group at the same time. + # In order to keep backward compatibility for the existing checkpoints, we use group['state'] to initialize state['step'] if it exists. + state['step'] = group.get('step', 0) + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + + if p.dtype == torch.float16: + g_16.append(p.grad.data) + p_16.append(p.data) + m_16.append(state['exp_avg']) + v_16.append(state['exp_avg_sq']) + elif p.dtype == torch.bfloat16: + g_bf.append(p.grad) + p_bf.append(p) + m_bf.append(state['exp_avg']) + v_bf.append(state['exp_avg_sq']) + elif p.dtype == torch.float32: + g_32.append(p.grad.data) + p_32.append(p.data) + m_32.append(state['exp_avg']) + v_32.append(state['exp_avg_sq']) + else: + raise RuntimeError('FusedAdam only support fp16, bf16 and fp32.') + + if len(g_16) > 0: + state['step'] += 1 + multi_tensor_applier(self.multi_tensor_adam, self._dummy_overflow_buf, [g_16, p_16, m_16, v_16], + group['lr'], beta1, beta2, group['eps'], state['step'], self.adam_w_mode, + bias_correction, group['weight_decay']) + + if len(g_bf) > 0: + state['step'] += 1 + multi_tensor_applier(self.multi_tensor_adam, self._dummy_overflow_buf, [g_bf, p_bf, m_bf, v_bf], + group['lr'], beta1, beta2, group['eps'], state['step'], self.adam_w_mode, + bias_correction, group['weight_decay']) + + if len(g_32) > 0: + state['step'] += 1 + multi_tensor_applier(self.multi_tensor_adam, self._dummy_overflow_buf, [g_32, p_32, m_32, v_32], + group['lr'], beta1, beta2, group['eps'], state['step'], self.adam_w_mode, + bias_correction, group['weight_decay']) + + return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/multi_tensor_apply.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/multi_tensor_apply.py new file mode 100644 index 0000000000000000000000000000000000000000..0ba228505cef747eea4fec62f3e68707fa4daa0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/multi_tensor_apply.py @@ -0,0 +1,17 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from NVIDIA/apex, commit a109f85 +""" + + +class MultiTensorApply(object): + + def __init__(self, chunk_size): + self.chunk_size = chunk_size + + def __call__(self, op, noop_flag_buffer, tensor_lists, *args): + return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6aea264237e366b9978c8bd096a2fbc28945705e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..op_builder import AsyncIOBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..199dbff1fd824c073030a17d24759bd4efb9fc75 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e38d56359fea64b067373b0e56cfed1dca5dab9a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..op_builder import DeepCompileBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e45f60e10aa7984b09a9b2b42e9e0382da30a6db Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e276ad0856dd620ecbe47231c5ebc4a8a753cb1f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp @@ -0,0 +1,215 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_adagrad.h" +#include +#include +#include +#include +#include +#include +#include + +using namespace std::string_literals; +static std::unordered_map> s_optimizers; + +// C++ interface + +template +void Adagrad_Optimizer::Step_1(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<1>(&rounded_size, _params, grads, _exp_avg_sq, _param_size); +#endif + if (_param_size > rounded_size) { + float step_size = -1 * _alpha; + for (size_t t = rounded_size; t < _param_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > _param_size) copy_size = _param_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t k = t; k < offset; k++) { + float grad = (float)grads[k]; + float param = (float)_params[k]; + float momentum = grads[k]; + float variance = _exp_avg_sq[k]; + if (_weight_decay > 0) { grad = param * _weight_decay + grad; } + + variance += grad * grad; + + grad = sqrt(variance); + grad += _eps; + grad = momentum / grad; + param = grad * step_size + param; + _params[k] = param; + // STORE UPDATE TERM TO GRAD'S MEMORY + grads[k] = grad * step_size; + _exp_avg_sq[k] = variance; + } + } + } +} + +template +void Adagrad_Optimizer::Step_4(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<4>(&rounded_size, _params, grads, _exp_avg_sq, _param_size); +#endif + if (_param_size > rounded_size) + Step_1((_params + rounded_size), + (grads + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size)); +} + +int create_adagrad_optimizer(int optimizer_id, + float alpha = 1e-2, + float eps = 1e-8, + float weight_decay = 0, + bool should_log = false) +{ + auto opt = std::make_shared(alpha, eps, weight_decay); + + s_optimizers[optimizer_id] = opt; + + if (should_log) { + std::string avx_type = ""; +#if defined(__AVX512__) + avx_type = "AVX512"; +#else +#if defined(__AVX256__) + avx_type = "AVX2"; +#else + avx_type = "scalar"; +#endif +#endif + + printf("Adagrad Optimizer #%d is created with %s arithmetic capability.\n", + optimizer_id, + avx_type.c_str()); + printf("Config: alpha=%f, weight_decay=%f\n", alpha, weight_decay); + } + + return 0; +} + +template +void Adagrad_Optimizer::Step_8(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<8>(&rounded_size, _params, grads, _exp_avg_sq, _param_size); +#endif + if (_param_size > rounded_size) + Step_4((_params + rounded_size), + (grads + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size)); +} + +template +void step_invoker(std::shared_ptr opt, + void* _params, + void* grads, + void* _exp_avg_sq, + size_t _param_size) +{ + opt->Step_8((ds_params_precision_t*)(_params), + (ds_params_precision_t*)(grads), + (ds_state_precision_t*)(_exp_avg_sq), + _param_size); +} + +std::map, + std::function, void*, void*, void*, size_t)>> + invokers; + +// Fill map with template functions for each type +template +void create_invoker() +{ + invokers[std::tuple(c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; +} +struct InvokerInitializer { + InvokerInitializer() + { + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + } +} _invoker_initializer; + +void invoke(std::shared_ptr opt, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg_sq, + size_t param_size) +{ + c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); + c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg_sq.options().dtype()); + + auto it = invokers.find(std::tuple(params_type, state_type)); + if (it == invokers.end()) { + throw std::runtime_error("Adagrad optimizer with param type "s + + c10::toString(params_type) + " and state type "s + + c10::toString(state_type) + + " is not supported on current hardware"s); + } + + it->second(opt, params.data_ptr(), grads.data_ptr(), exp_avg_sq.data_ptr(), param_size); +} + +int ds_adagrad_step(int optimizer_id, + size_t step, + float lr, + float epsilon, + float weight_decay, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg_sq) +{ + auto params_c = params.contiguous(); + auto grads_c = grads.contiguous(); + auto exp_avg_sq_c = exp_avg_sq.contiguous(); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + opt->IncrementStep(step); + opt->update_state(lr, epsilon, weight_decay); + + invoke(opt, params_c, grads_c, exp_avg_sq_c, params_c.numel()); + + return 0; +} + +int destroy_adagrad_optimizer(int optimizer_id) +{ + s_optimizers.erase(optimizer_id); + + return 0; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("adagrad_update", &ds_adagrad_step, "DeepSpeed CPU Adagrad update (C++)"); + m.def("create_adagrad", &create_adagrad_optimizer, "DeepSpeed CPU Adagrad (C++)"); + m.def("destroy_adagrad", &destroy_adagrad_optimizer, "DeepSpeed CPU Adagrad destroy (C++)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp new file mode 100644 index 0000000000000000000000000000000000000000..263c443cb4d4e5f8135f208a6fde9ced513ac8ae --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp @@ -0,0 +1,13 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_adam.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)"); + m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)"); + m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp new file mode 100644 index 0000000000000000000000000000000000000000..465aae7b9a343d5a3315c1531e11e8648c585b58 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp @@ -0,0 +1,244 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include +#include +#include +#include +#include "cpu_adam.h" + +using namespace std::string_literals; +static std::unordered_map> s_optimizers; + +// C++ interface + +template +void Adam_Optimizer::Step_1(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<1>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size); +#endif + if (_param_size > rounded_size) { + float betta1_minus1 = 1 - _betta1; + float betta2_minus1 = 1 - _betta2; + + float step_size = -1 * _alpha / _bias_correction1; + float w_decay = -1 * _alpha * _weight_decay; + + for (size_t t = rounded_size; t < _param_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > _param_size) copy_size = _param_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t k = t; k < offset; k++) { + float grad = (float)grads[k]; + float param = (float)_params[k]; + float momentum = _exp_avg[k]; + float variance = _exp_avg_sq[k]; + if (_weight_decay > 0 && !_adamw_mode) { grad = param * _weight_decay + grad; } + momentum = momentum * _betta1; + momentum = grad * betta1_minus1 + momentum; + + variance = variance * _betta2; + grad = grad * grad; + variance = grad * betta2_minus1 + variance; + + grad = sqrt(variance); + grad = grad * _bias_correction2 + _eps; + grad = momentum / grad; + if (_weight_decay > 0 && _adamw_mode) { param += w_decay * param; } + param = grad * step_size + param; + _params[k] = param; + _exp_avg[k] = momentum; + _exp_avg_sq[k] = variance; + } + } + } +} + +template +void Adam_Optimizer::Step_4(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<4>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size); +#endif + if (_param_size > rounded_size) + Step_1((_params + rounded_size), + (grads + rounded_size), + (_exp_avg + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size)); +} + +int create_adam_optimizer(int optimizer_id, + float alpha, + float betta1, + float betta2, + float eps, + float weight_decay, + bool adamw_mode, + bool should_log) +{ + auto opt = + std::make_shared(alpha, betta1, betta2, eps, weight_decay, adamw_mode); + + s_optimizers[optimizer_id] = opt; + + if (should_log) { + std::string avx_type = ""; +#if defined(__AVX512__) + avx_type = "AVX512"; +#else +#if defined(__AVX256__) + avx_type = "AVX2"; +#else + avx_type = "scalar"; +#endif +#endif + + printf("Adam Optimizer #%d is created with %s arithmetic capability.\n", + optimizer_id, + avx_type.c_str()); + printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n", + alpha, + betta1, + betta2, + weight_decay, + (int)adamw_mode); + } + + return 0; +} + +template +void Adam_Optimizer::Step_8(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<8>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size); +#endif + if (_param_size > rounded_size) + Step_4((_params + rounded_size), + (grads + rounded_size), + (_exp_avg + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size)); +} + +template +void step_invoker(std::shared_ptr opt, + void* _params, + void* grads, + void* _exp_avg, + void* _exp_avg_sq, + size_t _param_size) +{ + opt->Step_8((ds_params_precision_t*)(_params), + (ds_params_precision_t*)(grads), + (ds_state_precision_t*)(_exp_avg), + (ds_state_precision_t*)(_exp_avg_sq), + _param_size); +} + +std::map, + std::function, void*, void*, void*, void*, size_t)>> + invokers; + +// Fill map with template functions for each type +template +void create_invoker() +{ + invokers[std::tuple(c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; +} +struct InvokerInitializer { + InvokerInitializer() + { + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + } +} _invoker_initializer; + +void invoke(std::shared_ptr opt, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq, + size_t param_size) +{ + c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); + c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype()); + + auto it = invokers.find(std::tuple(params_type, state_type)); + if (it == invokers.end()) { + throw std::runtime_error("Adam optimizer with param type "s + c10::toString(params_type) + + " and state type "s + c10::toString(state_type) + + " is not supported on current hardware"s); + } + + it->second(opt, + params.data_ptr(), + grads.data_ptr(), + exp_avg.data_ptr(), + exp_avg_sq.data_ptr(), + param_size); +} + +int ds_adam_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + bool bias_correction, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq) +{ + auto params_c = params.contiguous(); + auto grads_c = grads.contiguous(); + auto exp_avg_c = exp_avg.contiguous(); + auto exp_avg_sq_c = exp_avg_sq.contiguous(); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + opt->IncrementStep(step, beta1, beta2); + opt->update_state(lr, epsilon, weight_decay, bias_correction); + + invoke(opt, params_c, grads_c, exp_avg_c, exp_avg_sq_c, params_c.numel()); + + return 0; +} + +int destroy_adam_optimizer(int optimizer_id) +{ + s_optimizers.erase(optimizer_id); + + return 0; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp new file mode 100644 index 0000000000000000000000000000000000000000..13b390248608b046dab443f85346b5446a47d722 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_adam", + &multi_tensor_adam_cuda, + "Compute and apply gradient update to parameters for Adam optimizer"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu new file mode 100644 index 0000000000000000000000000000000000000000..a1fc7d15aec9ad16ddb242fd62c99cbfdca55d0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu @@ -0,0 +1,203 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +// Another possibility: +// #include + +#include + +#include "multi_tensor_apply.cuh" +#include "type_shim.h" + +#define BLOCK_SIZE 512 +#define ILP 4 + +typedef enum : int { + ADAM_MODE_0 = 0, // L2 regularization mode + ADAM_MODE_1 = 1 // Decoupled weight decay mode(AdamW) +} adamMode_t; + +using MATH_T = float; + +template +struct AdamFunctor { + __device__ __forceinline__ void operator()(int chunk_size, + volatile int* noop_gmem, + TensorListMetadata<4>& tl, + const float beta1, + const float beta2, + const float beta1_correction, + const float beta2_correction, + const float epsilon, + const float lr, + adamMode_t mode, + const float decay) + { + // I'd like this kernel to propagate infs/nans. + // if(*noop_gmem == 1) + // return; + + index_t tensor_loc = tl.block_to_tensor[blockIdx.x]; + + // potentially use to pass in list of scalar + // int tensor_num = tl.start_tensor_this_launch + tensor_loc; + + index_t chunk_idx = tl.block_to_chunk[blockIdx.x]; + index_t n = tl.sizes[tensor_loc]; + + T* g = (T*)tl.addresses[0][tensor_loc]; + g += chunk_idx * chunk_size; + + T* p = (T*)tl.addresses[1][tensor_loc]; + p += chunk_idx * chunk_size; + + T* m = (T*)tl.addresses[2][tensor_loc]; + m += chunk_idx * chunk_size; + + T* v = (T*)tl.addresses[3][tensor_loc]; + v += chunk_idx * chunk_size; + + n -= chunk_idx * chunk_size; + + // see note in multi_tensor_scale_kernel.cu + for (index_t i_start = 0; i_start < n && i_start < chunk_size; + i_start += blockDim.x * ILP) { + MATH_T r_g[ILP]; + MATH_T r_p[ILP]; + MATH_T r_m[ILP]; + MATH_T r_v[ILP]; +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) { + r_g[ii] = g[i]; + r_p[ii] = p[i]; + r_m[ii] = m[i]; + r_v[ii] = v[i]; + } else { + r_g[ii] = MATH_T(0); + r_p[ii] = MATH_T(0); + r_m[ii] = MATH_T(0); + r_v[ii] = MATH_T(0); + } + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + if (mode == ADAM_MODE_0) { // L2 + r_g[ii] = r_g[ii] + (decay * r_p[ii]); + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sqrtf(next_v_unbiased) + epsilon; + MATH_T update = next_m_unbiased / denom; + r_p[ii] = r_p[ii] - (lr * update); + } else { // weight decay + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sqrtf(next_v_unbiased) + epsilon; + MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); + r_p[ii] = r_p[ii] - (lr * update); + } + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) { + p[i] = r_p[ii]; + m[i] = r_m[ii]; + v[i] = r_v[ii]; + } + } + } + } +}; + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay) +{ + using namespace at; + + // Handle bias correction mode + float bias_correction1 = 1.0f, bias_correction2 = 1.0f; + if (bias_correction == 1) { + bias_correction1 = 1 - std::pow(beta1, step); + bias_correction2 = 1 - std::pow(beta2, step); + } + + size_t max_size = 0; + bool requires_64bit_indexing = false; + for (auto it = tensor_lists.begin(); it != tensor_lists.end(); it++) { + for (auto it2 = it->begin(); it2 != it->end(); it2++) { + if (it2->numel() > max_size) { + max_size = it2->numel(); + if (max_size >= INT_MAX) { + requires_64bit_indexing = true; + break; + } + } + } + if (requires_64bit_indexing) { break; } + } + + // Assume single type across p,g,m1,m2 now + if (requires_64bit_indexing) { + DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), + 0, + "adam", + multi_tensor_apply<4>((int64_t)BLOCK_SIZE, + (int64_t)chunk_size, + noop_flag, + tensor_lists, + AdamFunctor(), + beta1, + beta2, + bias_correction1, + bias_correction2, + epsilon, + lr, + (adamMode_t)mode, + weight_decay);) + } else { + DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), + 0, + "adam", + multi_tensor_apply<4>(BLOCK_SIZE, + chunk_size, + noop_flag, + tensor_lists, + AdamFunctor(), + beta1, + beta2, + bias_correction1, + bias_correction2, + epsilon, + lr, + (adamMode_t)mode, + weight_decay);) + } + + AT_CUDA_CHECK(cudaGetLastError()); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh new file mode 100644 index 0000000000000000000000000000000000000000..342376c141be7578553910df43b003d29e0e4fc4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh @@ -0,0 +1,132 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +#include +#include "compat.h" + +#include + +// #include + +// This header is the one-stop shop for all your multi-tensor apply needs. + +// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) +constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; +constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; + +template +struct TensorListMetadata { + void* addresses[n][depth_to_max_tensors[n - 1]]; + int sizes[depth_to_max_tensors[n - 1]]; + unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. + int start_tensor_this_launch; +}; + +template +__global__ void multi_tensor_apply_kernel(int64_t chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) +{ + // Hand the chunk information to the user-supplied functor to process however it likes. + callable(chunk_size, noop_flag, tl, args...); +} + +template +void multi_tensor_apply(int64_t block_size, + int64_t chunk_size, + const at::Tensor& noop_flag, + const std::vector>& tensor_lists, + T callable, + ArgTypes... args) +{ + TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); + int len0 = tensor_lists[0].size(); + TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); + auto ref_device = tensor_lists[0][0].device(); + TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda"); + for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices + { + TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); + for (int t = 0; t < tensor_lists[l].size(); t++) { + // TODO: Print which tensor fails. + bool contiguous_memory = tensor_lists[l][t].is_contiguous(); +#ifdef VERSION_GE_1_5 + contiguous_memory = (contiguous_memory || + tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); +#endif + TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); + TORCH_CHECK(tensor_lists[l][t].device() == ref_device, + "A tensor was not on the same device as the first tensor"); + TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); + } + } + + int ntensors = tensor_lists[0].size(); + + TensorListMetadata tl; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); + auto stream = at::cuda::getCurrentCUDAStream(); + + tl.start_tensor_this_launch = 0; + int loc_block_info = 0; + int loc_tensor_info = 0; + for (int t = 0; t < ntensors; t++) { + tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); + for (int d = 0; d < depth; d++) + tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); + loc_tensor_info++; + + auto chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; + + for (auto chunk = 0; chunk < chunks_this_tensor; chunk++) { + // std::cout << chunks_this_tensor << std::endl; + tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; + tl.block_to_chunk[loc_block_info] = chunk; + loc_block_info++; + + bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && + chunk == chunks_this_tensor - 1); + bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); + bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); + if (tensors_full || blocks_full || last_chunk) { + // using accscalar_t = acc_type; + multi_tensor_apply_kernel<<>>( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); + + AT_CUDA_CHECK(cudaGetLastError()); + + // Reset. The control flow possibilities here make my brain hurt. + loc_block_info = 0; + if (chunk == chunks_this_tensor - 1) { + // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + loc_tensor_info = 0; + tl.start_tensor_this_launch = t + 1; + } else { + // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; + for (int d = 0; d < depth; d++) + tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; + loc_tensor_info = 1; + tl.start_tensor_this_launch = t; + } + } + } + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp new file mode 100644 index 0000000000000000000000000000000000000000..9d7ff5093017f60c06b96554f69630ffe9825918 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp @@ -0,0 +1,341 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "deepspeed_aio_common.h" + +using namespace std; +using namespace std::chrono; + +#define DEBUG_DS_AIO_PERF 0 +#define DEBUG_DS_AIO_SUBMIT_PERF 0 + +static const std::string c_library_name = "deepspeed_aio"; + +static void _report_aio_statistics(const char* tag, + const std::vector>& latencies) + __attribute__((unused)); + +static void _report_aio_statistics(const char* tag, + const std::vector>& latencies) +{ + std::vector lat_usec; + for (auto& lat : latencies) { lat_usec.push_back(lat.count() * 1e6); } + const auto min_lat = *(std::min_element(lat_usec.begin(), lat_usec.end())); + const auto max_lat = *(std::max_element(lat_usec.begin(), lat_usec.end())); + const auto avg_lat = std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size(); + + std::cout << c_library_name << ": latency statistics(usec) " << tag + << " min/max/avg = " << min_lat << " " << max_lat << " " << avg_lat << std::endl; +} + +static void _get_aio_latencies(std::vector>& raw_latencies, + struct deepspeed_aio_latency_t& summary_latencies) +{ + std::vector lat_usec; + for (auto& lat : raw_latencies) { lat_usec.push_back(lat.count() * 1e6); } + summary_latencies._min_usec = *(std::min_element(lat_usec.begin(), lat_usec.end())); + summary_latencies._max_usec = *(std::max_element(lat_usec.begin(), lat_usec.end())); + summary_latencies._avg_usec = + std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size(); +} + +static void _do_io_submit_singles(const int64_t n_iocbs, + const int64_t iocb_index, + std::unique_ptr& aio_ctxt, + std::vector>& submit_times) +{ + for (auto i = 0; i < n_iocbs; ++i) { + const auto st = std::chrono::high_resolution_clock::now(); + const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, 1, aio_ctxt->_iocbs.data() + i); + submit_times.push_back(std::chrono::high_resolution_clock::now() - st); +#if DEBUG_DS_AIO_SUBMIT_PERF + printf("submit(usec) %f io_index=%lld buf=%p len=%lu off=%llu \n", + submit_times.back().count() * 1e6, + iocb_index, + aio_ctxt->_iocbs[i]->u.c.buf, + aio_ctxt->_iocbs[i]->u.c.nbytes, + aio_ctxt->_iocbs[i]->u.c.offset); +#endif + assert(submit_ret > 0); + } +} + +static void _do_io_submit_block(const int64_t n_iocbs, + const int64_t iocb_index, + std::unique_ptr& aio_ctxt, + std::vector>& submit_times) +{ + const auto st = std::chrono::high_resolution_clock::now(); + const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, n_iocbs, aio_ctxt->_iocbs.data()); + submit_times.push_back(std::chrono::high_resolution_clock::now() - st); +#if DEBUG_DS_AIO_SUBMIT_PERF + printf("submit(usec) %f io_index=%lld nr=%lld buf=%p len=%lu off=%llu \n", + submit_times.back().count() * 1e6, + iocb_index, + n_iocbs, + aio_ctxt->_iocbs[0]->u.c.buf, + aio_ctxt->_iocbs[0]->u.c.nbytes, + aio_ctxt->_iocbs[0]->u.c.offset); +#endif + assert(submit_ret > 0); +} + +static int _do_io_complete(const int64_t min_completes, + const int64_t max_completes, + std::unique_ptr& aio_ctxt, + std::vector>& reap_times) +{ + const auto start_time = std::chrono::high_resolution_clock::now(); + int64_t n_completes = io_pgetevents(aio_ctxt->_io_ctxt, + min_completes, + max_completes, + aio_ctxt->_io_events.data(), + nullptr, + nullptr); + reap_times.push_back(std::chrono::high_resolution_clock::now() - start_time); + assert(n_completes >= min_completes); + return n_completes; +} + +void do_aio_operation_sequential(const bool read_op, + std::unique_ptr& aio_ctxt, + std::unique_ptr& xfer_ctxt, + deepspeed_aio_config_t* config, + deepspeed_aio_perf_t* perf) +{ + struct io_prep_context prep_ctxt(read_op, xfer_ctxt, aio_ctxt->_block_size, &aio_ctxt->_iocbs); + + const auto num_io_blocks = static_cast( + ceil(static_cast(xfer_ctxt->_num_bytes) / aio_ctxt->_block_size)); +#if DEBUG_DS_AIO_PERF + const auto io_op_name = std::string(read_op ? "read" : "write"); + std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes + << " bytes with " << num_io_blocks << " io blocks" << std::endl; +#endif + + std::vector> submit_times; + std::vector> reap_times; + const auto max_queue_bytes = + static_cast(aio_ctxt->_queue_depth * aio_ctxt->_block_size); + + auto start = std::chrono::high_resolution_clock::now(); + for (int64_t iocb_index = 0; iocb_index < num_io_blocks; iocb_index += aio_ctxt->_queue_depth) { + const auto start_offset = iocb_index * aio_ctxt->_block_size; + const auto start_buffer = (char*)xfer_ctxt->_mem_buffer + start_offset; + const auto n_iocbs = + min(static_cast(aio_ctxt->_queue_depth), (num_io_blocks - iocb_index)); + const auto num_bytes = min(max_queue_bytes, (xfer_ctxt->_num_bytes - start_offset)); + prep_ctxt.prep_iocbs(n_iocbs, num_bytes, start_buffer, start_offset); + + if (config->_single_submit) { + _do_io_submit_singles(n_iocbs, iocb_index, aio_ctxt, submit_times); + } else { + _do_io_submit_block(n_iocbs, iocb_index, aio_ctxt, submit_times); + } + + _do_io_complete(n_iocbs, n_iocbs, aio_ctxt, reap_times); + } + const std::chrono::duration elapsed = std::chrono::high_resolution_clock::now() - start; + + if (perf) { + _get_aio_latencies(submit_times, perf->_submit); + _get_aio_latencies(reap_times, perf->_complete); + perf->_e2e_usec = elapsed.count() * 1e6; + perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9); + } + +#if DEBUG_DS_AIO_PERF + _report_aio_statistics("submit", submit_times); + _report_aio_statistics("complete", reap_times); +#endif + +#if DEBUG_DS_AIO_PERF + std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6 + << " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl; +#endif + +#if DEBUG_DS_AIO_PERF + std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes + << " bytes " << std::endl; +#endif +} + +void do_aio_operation_overlap(const bool read_op, + std::unique_ptr& aio_ctxt, + std::unique_ptr& xfer_ctxt, + deepspeed_aio_config_t* config, + deepspeed_aio_perf_t* perf) +{ + struct io_prep_generator io_gen(read_op, xfer_ctxt, aio_ctxt->_block_size); + +#if DEBUG_DS_AIO_PERF + const auto io_op_name = std::string(read_op ? "read" : "write"); + std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes + << " bytes with " << io_gen._num_io_blocks << " io blocks" << std::endl; +#endif + + std::vector> submit_times; + std::vector> reap_times; + + auto request_iocbs = aio_ctxt->_queue_depth; + auto n_pending_iocbs = 0; + const auto min_completes = 1; + auto start = std::chrono::high_resolution_clock::now(); + while (true) { + const auto n_iocbs = io_gen.prep_iocbs(request_iocbs - n_pending_iocbs, &aio_ctxt->_iocbs); + if (n_iocbs > 0) { + if (config->_single_submit) { + _do_io_submit_singles( + n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times); + } else { + _do_io_submit_block( + n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times); + } + } + + n_pending_iocbs += n_iocbs; + assert(n_pending_iocbs <= aio_ctxt->_queue_depth); + + if (n_pending_iocbs == 0) { break; } + + const auto n_complete = + _do_io_complete(min_completes, n_pending_iocbs, aio_ctxt, reap_times); + n_pending_iocbs -= n_complete; + } + + const std::chrono::duration elapsed = std::chrono::high_resolution_clock::now() - start; + + if (perf) { + _get_aio_latencies(submit_times, perf->_submit); + _get_aio_latencies(reap_times, perf->_complete); + perf->_e2e_usec = elapsed.count() * 1e6; + perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9); + } + +#if DEBUG_DS_AIO_PERF + _report_aio_statistics("submit", submit_times); + _report_aio_statistics("complete", reap_times); +#endif + +#if DEBUG_DS_AIO_PERF + std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6 + << " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl; +#endif + +#if DEBUG_DS_AIO_PERF + std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes + << " bytes " << std::endl; +#endif +} + +void report_file_error(const char* filename, const std::string file_op, const int error_code) +{ + std::string err_msg = file_op + std::string(" failed on ") + std::string(filename) + + " error = " + std::to_string(error_code); + std::cerr << c_library_name << ": " << err_msg << std::endl; +} + +int open_file(const char* filename, const bool read_op) +{ + const int flags = read_op ? (O_RDONLY | O_DIRECT) : (O_WRONLY | O_CREAT | O_DIRECT); +#if defined(__ENABLE_CANN__) + int* flags_ptr = (int*)&flags; + *flags_ptr = read_op ? (O_RDONLY) : (O_WRONLY | O_CREAT); +#endif + const int mode = 0600; + const auto fd = open(filename, flags, mode); + if (fd == -1) { + const auto error_code = errno; + const auto error_msg = read_op ? " open for read " : " open for write "; + report_file_error(filename, error_msg, error_code); + return -1; + } + return fd; +} + +int regular_read(const char* filename, std::vector& buffer) +{ + const auto fd = open(filename, O_RDONLY, 0600); + assert(fd != -1); + struct stat fs; + const auto result = fstat(fd, &fs); + assert(result != -1); + int64_t num_bytes = fs.st_size; + buffer.resize(num_bytes); + int64_t read_bytes = 0; + auto r = 0; + do { + const auto buffer_ptr = buffer.data() + read_bytes; + const auto bytes_to_read = num_bytes - read_bytes; + r = read(fd, buffer_ptr, bytes_to_read); + read_bytes += r; + } while (r > 0); + + if (read_bytes != num_bytes) { + std::cerr << "read error " << " read_bytes (read) = " << read_bytes + << " num_bytes (fstat) = " << num_bytes << std::endl; + } + assert(read_bytes == num_bytes); + close(fd); + return 0; +} + +static bool _validate_buffer(const char* filename, void* aio_buffer, const int64_t num_bytes) +{ + std::vector regular_buffer; + const auto reg_ret = regular_read(filename, regular_buffer); + assert(0 == reg_ret); + std::cout << "regular read of " << filename << " returned " << regular_buffer.size() << " bytes" + << std::endl; + + if (static_cast(regular_buffer.size()) != num_bytes) { return false; } + + return (0 == memcmp(aio_buffer, regular_buffer.data(), regular_buffer.size())); +} + +bool validate_aio_operation(const bool read_op, + const char* filename, + void* aio_buffer, + const int64_t num_bytes) +{ + const auto msg_suffix = std::string("deepspeed_aio_") + + std::string(read_op ? "read()" : "write()") + + std::string("using read()"); + + if (false == _validate_buffer(filename, aio_buffer, num_bytes)) { + std::cout << "Fail: correctness of " << msg_suffix << std::endl; + return false; + } + + std::cout << "Pass: correctness of " << msg_suffix << std::endl; + return true; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h new file mode 100644 index 0000000000000000000000000000000000000000..aa4e49f4f4edee792e81bc0581e16ef3ab968811 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include +#include + +using namespace std; + +void do_aio_operation_sequential(const bool read_op, + std::unique_ptr& aio_ctxt, + std::unique_ptr& xfer_ctxt, + deepspeed_aio_config_t* config, + deepspeed_aio_perf_t* perf); + +void do_aio_operation_overlap(const bool read_op, + std::unique_ptr& aio_ctxt, + std::unique_ptr& xfer_ctxt, + deepspeed_aio_config_t* config, + deepspeed_aio_perf_t* perf); + +int open_file(const char* filename, const bool read_op); + +void report_file_error(const char* filename, const std::string file_op, const int error_code); + +int regular_read(const char* filename, std::vector& buffer); + +bool validate_aio_operation(const bool read_op, + const char* filename, + void* aio_buffer, + const int64_t num_bytes); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp new file mode 100644 index 0000000000000000000000000000000000000000..5e34a61065d4b0f16ba904d0e397eeccc149e621 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp @@ -0,0 +1,76 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include + +#include "deepspeed_aio_utils.h" + +using namespace std; + +const int c_block_size = 128 * 1024; +const int c_io_queue_depth = 8; + +deepspeed_aio_config_t::deepspeed_aio_config_t() + : _block_size(c_block_size), + _queue_depth(c_io_queue_depth), + _single_submit(false), + _overlap_events(false), + _lock_memory(false) +{ +} + +deepspeed_aio_config_t::deepspeed_aio_config_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const bool lock_memory) + : _block_size(block_size), + _queue_depth(queue_depth), + _single_submit(single_submit), + _overlap_events(overlap_events), + _lock_memory(lock_memory) +{ +} + +void deepspeed_aio_latency_t::dump(const std::string tag) +{ + std::cout << tag << _min_usec << " " << _max_usec << " " << _avg_usec << " " << std::endl; +} + +void deepspeed_aio_latency_t::accumulate(const struct deepspeed_aio_latency_t& other) +{ + _min_usec += other._min_usec; + _max_usec += other._max_usec; + _avg_usec += other._avg_usec; +} + +void deepspeed_aio_latency_t::scale(const float scaler) +{ + _min_usec *= scaler; + _max_usec *= scaler; + _avg_usec *= scaler; +} + +aio_context::aio_context(const int block_size, const int queue_depth) +{ + _block_size = block_size; + _queue_depth = queue_depth; + for (auto i = 0; i < queue_depth; ++i) { + _iocbs.push_back((struct iocb*)calloc(1, sizeof(struct iocb))); + } + _io_events.resize(queue_depth); + io_queue_init(queue_depth, &_io_ctxt); +} + +aio_context::~aio_context() +{ + for (auto& iocb : _iocbs) { free(iocb); } + _io_events.resize(0); + io_queue_release(_io_ctxt); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h new file mode 100644 index 0000000000000000000000000000000000000000..ce6a4e5cdfa75120c24dad599a8bc717765c220d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h @@ -0,0 +1,59 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include + +#include +#include + +using namespace std; + +struct deepspeed_aio_latency_t { + double _min_usec; + double _max_usec; + double _avg_usec; + + void dump(const std::string tag); + void accumulate(const deepspeed_aio_latency_t&); + void scale(const float value); +}; + +struct deepspeed_aio_perf_t { + deepspeed_aio_latency_t _submit; + deepspeed_aio_latency_t _complete; + double _e2e_usec; + double _e2e_rate_GB; +}; + +struct deepspeed_aio_config_t { + const int _block_size; + const int _queue_depth; + const bool _single_submit; + const bool _overlap_events; + const bool _lock_memory; + + deepspeed_aio_config_t(); + deepspeed_aio_config_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const bool lock_memory); +}; + +struct aio_context { + io_context_t _io_ctxt; + std::vector _io_events; + std::vector _iocbs; + int _block_size; + int _queue_depth; + + aio_context(const int block_size, const int queue_depth); + ~aio_context(); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c8e577f299ae9b8f1aae41d742d5b69395fc026b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp @@ -0,0 +1,140 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include + +#include "deepspeed_aio_utils.h" + +using namespace std; + +const int c_block_size = 128 * 1024; +const int c_io_queue_depth = 8; + +io_xfer_ctxt::io_xfer_ctxt(const int fd, + const int64_t file_offset, + const int64_t buffer_offset, + const int64_t num_bytes, + const void* buffer) + : _fd(fd), + _file_base_offset(file_offset), + _buffer_base_offset(buffer_offset), + _mem_buffer(buffer), + _num_bytes(num_bytes) +{ +} + +io_prep_context::io_prep_context(const bool read_op, + const std::unique_ptr& xfer_ctxt, + const size_t block_size, + const std::vector* iocbs) + : _read_op(read_op), _xfer_ctxt(xfer_ctxt), _block_size(block_size), _iocbs(iocbs) +{ +} + +void io_prep_context::prep_iocbs(const int n_iocbs, + const size_t num_bytes, + const void* start_buffer, + const int64_t start_offset) +{ + assert(static_cast(n_iocbs) <= _iocbs->size()); + for (auto i = 0; i < n_iocbs; ++i) { + const auto shift = i * _block_size; + const auto xfer_buffer = (char*)start_buffer + _xfer_ctxt->_buffer_base_offset + shift; + const auto xfer_offset = _xfer_ctxt->_file_base_offset + start_offset + shift; + auto byte_count = _block_size; + + if ((shift + _block_size) > num_bytes) { byte_count = num_bytes - shift; } + + if (_read_op) { + io_prep_pread(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset); + } else { + io_prep_pwrite(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset); + } + } +} + +io_prep_generator::io_prep_generator(const bool read_op, + const std::unique_ptr& xfer_ctxt, + const size_t block_size) + : _read_op(read_op), + _xfer_ctxt(xfer_ctxt), + _block_size(block_size), + _remaining_bytes(xfer_ctxt->_num_bytes), + _next_iocb_index(0) +{ + _num_io_blocks = + static_cast(ceil(static_cast(xfer_ctxt->_num_bytes) / block_size)); + _remaining_io_blocks = _num_io_blocks; +} + +int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* iocbs) +{ + if ((_remaining_bytes) == 0 || (_remaining_io_blocks == 0)) { + assert(static_cast(_remaining_bytes) == _remaining_io_blocks); + return 0; + } + + assert(static_cast(n_iocbs) <= iocbs->size()); + + auto actual_n_iocbs = min(static_cast(n_iocbs), _remaining_io_blocks); + for (auto i = 0; i < actual_n_iocbs; ++i, ++_next_iocb_index) { + const auto xfer_buffer = (char*)_xfer_ctxt->_mem_buffer + _xfer_ctxt->_buffer_base_offset + + (_next_iocb_index * _block_size); + const auto xfer_offset = _xfer_ctxt->_file_base_offset + (_next_iocb_index * _block_size); + const auto num_bytes = min(static_cast(_block_size), _remaining_bytes); + if (_read_op) { + io_prep_pread(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset); + } else { + io_prep_pwrite(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset); + } + _remaining_bytes -= num_bytes; + } + _remaining_io_blocks -= actual_n_iocbs; + + return actual_n_iocbs; +} + +int64_t get_file_size(const char* filename, int64_t& size) +{ + struct stat st; + if (stat(filename, &st) == -1) { return -1; } + size = st.st_size; + return 0; +} + +int64_t get_fd_file_size(const int fd, int64_t& size) +{ + struct stat st; + if (fstat(fd, &st) == -1) { return -1; } + size = st.st_size; + return 0; +} + +void* ds_page_aligned_alloc(const int64_t size, const bool lock) +{ + void* ptr; + int retval; + + retval = posix_memalign(&ptr, (size_t)sysconf(_SC_PAGESIZE), size); + if (retval) { return nullptr; } + + if (lock == false) { return ptr; } + + auto mlock_ret = mlock(ptr, size); + if (mlock_ret != 0) { + auto mlock_error = errno; + std::cerr << "mlock failed to allocate " << size << " bytes with error no " << mlock_error + << " msg " << strerror(mlock_error) << std::endl; + free(ptr); + return nullptr; + } + + return ptr; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..8742bf5bff543948032ce19de2fc4bea194a9ca9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +struct io_xfer_ctxt { + const int _fd; + const int64_t _file_base_offset; + const int64_t _buffer_base_offset; + const void* _mem_buffer; + const int64_t _num_bytes; + + io_xfer_ctxt(const int fd, + const int64_t file_offset, + const int64_t buffer_offset, + const int64_t num_bytes, + const void* buffer); +}; + +struct io_prep_context { + const bool _read_op; + const std::unique_ptr& _xfer_ctxt; + const size_t _block_size; + const std::vector* _iocbs; + + io_prep_context(const bool read_op, + const std::unique_ptr& xfer_ctxt, + const size_t block_size, + const std::vector* iocbs); + + void prep_iocbs(const int n_iocbs, + const size_t num_bytes, + const void* start_buffer, + const int64_t start_offset); +}; + +struct io_prep_generator { + const bool _read_op; + const std::unique_ptr& _xfer_ctxt; + const size_t _block_size; + + int64_t _remaining_bytes; + int64_t _num_io_blocks; + int64_t _remaining_io_blocks; + int64_t _next_iocb_index; + + io_prep_generator(const bool read_op, + const std::unique_ptr& xfer_ctxt, + const size_t block_size); + + int prep_iocbs(const int n_iocbs, std::vector* iocbs); +}; + +void* ds_page_aligned_alloc(const int64_t size, const bool lock = false); + +int64_t get_file_size(const char* filename, int64_t& size); +int64_t get_fd_file_size(const int fd, int64_t& size); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp new file mode 100644 index 0000000000000000000000000000000000000000..8387e667b33236e9ffcb7d070baec43cf498d20f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepspeed_aio_op_desc.h" + +using namespace std; + +io_op_desc_t::io_op_desc_t(const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const int intra_op_parallelism, + const bool validate, + const int64_t file_offset) + : _read_op(read_op), + _buffer(buffer), + _fd(fd), + _filename((filename == nullptr) ? std::string() : filename), + _file_offset(file_offset), + _intra_op_parallelism(intra_op_parallelism), + _num_bytes_per_thread(static_cast(buffer.nbytes()) / intra_op_parallelism), + _validate(validate) +{ + if (validate) { assert(nullptr != filename); } +} + +char* io_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } + +void io_op_desc_t::finish() {} + +void io_op_desc_t::validate() {} + +void io_op_desc_t::run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config) +{ +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h new file mode 100644 index 0000000000000000000000000000000000000000..cc7f15d7465830b32c273232c4e4353391c4a4d4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#ifndef _IO_OP_DESC_T_ +#define _IO_OP_DESC_T_ +#include +#include +#include "deepspeed_py_aio.h" + +struct io_op_desc_t { + const bool _read_op; + torch::Tensor _buffer; + int _fd; + std::string _filename; + const int _intra_op_parallelism; + const int64_t _num_bytes_per_thread; + torch::Tensor _contiguous_buffer; + const bool _validate; + const int64_t _file_offset; + + io_op_desc_t(const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const int intra_op_parallelism, + const bool validate, + const int64_t file_offset); + + virtual void run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config); + + virtual char* data_ptr() const; + + virtual void validate(); + + virtual void finish(); +}; +#endif // _IO_OP_DESC_T_ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp new file mode 100644 index 0000000000000000000000000000000000000000..30c3b49143979cd4e71f7e0898b0d454a8099a51 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_aio_thread.h" + +using namespace std; + +deepspeed_aio_thread_t::deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config) + : _tid(tid), + _aio_config(aio_config), + _aio_ctxt(new aio_context(aio_config._block_size, aio_config._queue_depth)), + _time_to_exit(false) +{ +} + +deepspeed_aio_thread_t::~deepspeed_aio_thread_t() {} + +void deepspeed_aio_thread_t::run() +{ + while (true) { + std::shared_ptr next_io_op = nullptr; + + { + std::unique_lock lock(_work_sync._mutex); + _work_sync._cond_var.wait(lock, + [this] { return (!_work_queue.empty() || _time_to_exit); }); + if (!_work_queue.empty()) { + next_io_op = _work_queue.front(); + _work_queue.pop(); + } + } + + if (next_io_op) { + next_io_op->run(_tid, _aio_ctxt, &_aio_config); + + { + std::lock_guard lock(_complete_sync._mutex); + _complete_queue.push(next_io_op); + } + _complete_sync._cond_var.notify_one(); + } + + if (_time_to_exit) { break; } + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h new file mode 100644 index 0000000000000000000000000000000000000000..a192804db13d8a9a1a160134ce5f90e0ca091cd7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include +#include "deepspeed_cpu_op.h" + +struct thread_sync_t { + std::mutex _mutex; + std::condition_variable _cond_var; +}; + +struct deepspeed_aio_thread_t { + const int _tid; + deepspeed_aio_config_t& _aio_config; + + std::unique_ptr _aio_ctxt; + std::queue> _work_queue; + std::queue> _complete_queue; + + bool _time_to_exit; + + struct thread_sync_t _work_sync; + struct thread_sync_t _complete_sync; + + deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config); + + ~deepspeed_aio_thread_t(); + + void run(); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..3ee1ecd3b4a8c9acd49c516fcb3eb69a088916d6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp @@ -0,0 +1,102 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepspeed_cpu_op.h" +#include "deepspeed_pin_tensor.h" + +using namespace std; + +cpu_op_desc_t::cpu_op_desc_t( + const std::unique_ptr& pinned_tensor_mgr, + const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const int intra_op_parallelism, + const bool validate, + const int64_t file_offset) + : io_op_desc_t(read_op, buffer, fd, filename, intra_op_parallelism, validate, file_offset), + _cpu_buffer(buffer), + _pinned_tensor_mgr(pinned_tensor_mgr), + _is_managed_bounce_buffer(false) +{ + // Need to use CPU bounce buffer if buffer is not a page-locked DRAM memory. + _use_bounce_buffer = + !(_buffer.is_cpu() && (_buffer.is_pinned() || _pinned_tensor_mgr->is_managed(_buffer))); + if (_use_bounce_buffer) { + _alloc_bounce_buffer(); + if (!_read_op) { _cpu_buffer.copy_(_buffer); } + } + _contiguous_buffer = _cpu_buffer.contiguous(); +} + +char* cpu_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } + +void cpu_op_desc_t::finish() +{ + if (_use_bounce_buffer) { + if (_read_op) { + if (_buffer.is_cuda()) { + _buffer.copy_(_cpu_buffer.to(torch::Device(torch::kCUDA, _buffer.get_device()), + /*non_blocking=*/true)); + } + if (_buffer.is_xpu()) { _buffer.copy_(_cpu_buffer.to(torch::kXPU)); } + if (_buffer.is_cpu()) { _buffer.copy_(_cpu_buffer); } +#if defined(__ENABLE_CANN__) + if (torch_npu::utils::is_npu(_buffer)) { + auto device = at::Device("npu:0"); + _buffer.copy_(_cpu_buffer.to(device)); + } +#endif + } + + _free_bounce_buffer(); + } +} + +void cpu_op_desc_t::validate() +{ + const auto num_io_bytes = static_cast(_contiguous_buffer.nbytes()); + validate_aio_operation(_read_op, _filename.c_str(), data_ptr(), num_io_bytes); +} + +void cpu_op_desc_t::run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config) +{ + assert(tid < _intra_op_parallelism); + const auto buffer_base_offset = _num_bytes_per_thread * tid; + const auto file_base_offset = _file_offset + (_num_bytes_per_thread * tid); + + std::unique_ptr xfer_ctxt(new io_xfer_ctxt( + _fd, file_base_offset, buffer_base_offset, _num_bytes_per_thread, data_ptr())); + + if (aio_config->_overlap_events) { + do_aio_operation_overlap(_read_op, aio_ctxt, xfer_ctxt, aio_config, nullptr); + } else { + do_aio_operation_sequential(_read_op, aio_ctxt, xfer_ctxt, aio_config, nullptr); + } +} + +void cpu_op_desc_t::_alloc_bounce_buffer() +{ + auto options = torch::TensorOptions() + .dtype(_buffer.dtype()) + .layout(_buffer.layout()) + .device(torch::kCPU) + .requires_grad(false); + +#if defined(__CUDA_ARCH__) + _cpu_buffer = torch::empty(_buffer.numel(), options).pin_memory(); +#else + _is_managed_bounce_buffer = true; + _cpu_buffer = _pinned_tensor_mgr->alloc(_buffer.numel(), options); +#endif +} + +void cpu_op_desc_t::_free_bounce_buffer() +{ + if (_is_managed_bounce_buffer) { _pinned_tensor_mgr->free(_cpu_buffer); } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h new file mode 100644 index 0000000000000000000000000000000000000000..7cc648bace8ed680e85c53652788c6d7349f293e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include "deepspeed_aio_op_desc.h" + +struct cpu_op_desc_t : io_op_desc_t { + torch::Tensor _cpu_buffer; + bool _use_bounce_buffer; + bool _is_managed_bounce_buffer; + const std::unique_ptr& _pinned_tensor_mgr; + + cpu_op_desc_t(const std::unique_ptr& pinned_tensor_mgr, + const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const int intra_op_parallelism, + const bool validate, + const int64_t file_offset); + + void run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config); + + char* data_ptr() const; + + void validate(); + + void finish(); + + void _alloc_bounce_buffer(); + void _free_bounce_buffer(); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp new file mode 100644 index 0000000000000000000000000000000000000000..a97a4ac18ba811dfaad713fdee15248caecf035b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for managing CPU tensors occupying page-locked memory. +*/ + +#include "deepspeed_pin_tensor.h" + +using namespace std; + +deepspeed_pin_tensor_t::~deepspeed_pin_tensor_t() +{ + for (auto iter = _locked_tensors.begin(); iter != _locked_tensors.end(); ++iter) { + munlock(iter->first, iter->second); + std::free((void*)iter->first); + } + _locked_tensors.clear(); +} + +torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, + const torch::TensorOptions& options) +{ + const auto scalar_dtype = torch::typeMetaToScalarType(options.dtype()); + const auto num_bytes = num_elem * torch::elementSize(scalar_dtype); + auto pinned_buffer = ds_page_aligned_alloc(num_bytes, true); + assert(nullptr != pinned_buffer); + + _locked_tensors[pinned_buffer] = num_bytes; + + return at::from_blob(pinned_buffer, static_cast(num_elem), options); +} + +torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, const at::ScalarType& elem_type) +{ + auto options = torch::TensorOptions().dtype(elem_type).device(torch::kCPU).requires_grad(false); + return alloc(num_elem, options); +} + +bool deepspeed_pin_tensor_t::free(torch::Tensor& locked_tensor) +{ + auto addr = locked_tensor.data_ptr(); + if (_locked_tensors.find(addr) != _locked_tensors.end()) { + munlock(addr, _locked_tensors[addr]); + std::free(addr); + _locked_tensors.erase(addr); + return true; + } + + return false; +} + +bool deepspeed_pin_tensor_t::is_managed(const torch::Tensor& buffer) +{ + if (!buffer.is_cpu()) { return false; } + auto addr = buffer.data_ptr(); + if (_locked_tensors.find(addr) != _locked_tensors.end()) { return true; } + return false; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h new file mode 100644 index 0000000000000000000000000000000000000000..4b8ad7e76085877818aa26e511968efb0dc31fd2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for managing CPU tensors occupying page-locked memory. +TODO: Implement a full-featured manager that +1. Avoid page-locked memory leaks +2. Minimize page-locked memory usage by reducing internal fragmentation +Functionality for managing CPU tensors occupying page-locked memory. +*/ + +#include +#include "deepspeed_py_aio.h" + +struct deepspeed_pin_tensor_t { + std::map _locked_tensors; + + deepspeed_pin_tensor_t() = default; + + ~deepspeed_pin_tensor_t(); + + torch::Tensor alloc(const int64_t num_elem, const at::ScalarType& elem_type); + torch::Tensor alloc(const int64_t num_elem, const torch::TensorOptions& options); + + bool free(torch::Tensor& locked_tensor); + + bool is_managed(const torch::Tensor& buffer); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1ff0397043fae2d904fb22d5d0b9eea5af6b6cd4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "deepspeed_py_aio.h" + +using namespace std; +using namespace std::chrono; + +#define DEBUG_DS_AIO_READ 0 +#define DEBUG_DS_AIO_WRITE 0 + +static const std::string c_library_name = "deepspeed_aio"; + +int deepspeed_py_aio_write(const torch::Tensor& buffer, + const char* filename, + const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const bool validate) +{ + const auto start_time = std::chrono::high_resolution_clock::now(); + deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false); + + const auto fd = open_file(filename, false); + if (fd == -1) { return -1; } + + auto write_buffer = (char*)buffer.data_ptr(); + const auto num_write_bytes = static_cast(buffer.nbytes()); + + std::unique_ptr xfer_ctxt( + new io_xfer_ctxt(fd, 0, 0, num_write_bytes, write_buffer)); + std::unique_ptr aio_ctxt(new aio_context(config._block_size, config._queue_depth)); + + if (config._overlap_events) { + do_aio_operation_overlap(false, aio_ctxt, xfer_ctxt, &config, nullptr); + } else { + do_aio_operation_sequential(false, aio_ctxt, xfer_ctxt, &config, nullptr); + } + const std::chrono::duration aio_time = + std::chrono::high_resolution_clock::now() - start_time; + + close(fd); + + if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); } + + const std::chrono::duration fn_time = + std::chrono::high_resolution_clock::now() - start_time; + std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 + << " call = " << fn_time.count() * 1e6 << std::endl; + return 0; +} + +int deepspeed_py_aio_read(torch::Tensor& buffer, + const char* filename, + const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const bool validate) +{ + const auto start_time = std::chrono::high_resolution_clock::now(); + int64_t num_file_bytes; + if (-1 == get_file_size(filename, num_file_bytes)) { + const auto error_code = errno; + report_file_error(filename, " fstat for read", error_code); + return -1; + } + + deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false); + const auto fd = open_file(filename, true); + if (fd == -1) { return -1; } + + auto read_buffer = (char*)buffer.data_ptr(); + assert(static_cast(buffer.nbytes()) == num_file_bytes); + + std::unique_ptr xfer_ctxt( + new io_xfer_ctxt(fd, 0, 0, num_file_bytes, read_buffer)); + std::unique_ptr aio_ctxt(new aio_context(config._block_size, config._queue_depth)); + + if (config._overlap_events) { + do_aio_operation_overlap(true, aio_ctxt, xfer_ctxt, &config, nullptr); + } else { + do_aio_operation_sequential(true, aio_ctxt, xfer_ctxt, &config, nullptr); + } + const std::chrono::duration aio_time = + std::chrono::high_resolution_clock::now() - start_time; + + close(fd); + + if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); } + + const std::chrono::duration fn_time = + std::chrono::high_resolution_clock::now() - start_time; + std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 + << " call = " << fn_time.count() * 1e6 << std::endl; + return 0; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h new file mode 100644 index 0000000000000000000000000000000000000000..ba794db5440d540967053fc2ae4b184cfb8825d8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include + +int deepspeed_py_aio_write(const torch::Tensor& buffer, + const char* filename, + const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const bool validate); + +int deepspeed_py_aio_read(torch::Tensor& buffer, + const char* filename, + const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const bool validate); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2b1093e9928659bdaac1c14323928c644061155e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_py_aio_handle.h" +#include + +using namespace std; + +deepspeed_aio_handle_t::deepspeed_aio_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism) + : deepspeed_io_handle_t(block_size, + queue_depth, + single_submit, + overlap_events, + intra_op_parallelism) +{ +} + +deepspeed_aio_handle_t::~deepspeed_aio_handle_t() {} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h new file mode 100644 index 0000000000000000000000000000000000000000..c9fcb6d2b46227251ec56a286c9144935daffe0b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ +#include +#include +#include "deepspeed_py_io_handle.h" + +struct deepspeed_aio_handle_t : deepspeed_io_handle_t { + deepspeed_aio_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism); + + ~deepspeed_aio_handle_t(); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp new file mode 100644 index 0000000000000000000000000000000000000000..f5480e9d9d836db587cc9ded27c8f437bc294452 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp @@ -0,0 +1,135 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_py_copy.h" +#include + +#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) + +#if defined(__AVX512__) or defined(__AVX256__) +union AVX_Data { +#if defined(__AVX512__) + __m512 data; +#else + __m256 data; +#endif +}; +#endif + +static void helper_memcpy_1(float* dest, float* src, size_t param_size) +{ + size_t rounded_size = 0; + +#if defined(__AVX512__) or defined(__AVX256__) + + rounded_size = ROUND_DOWN(param_size, SIMD_WIDTH); + + for (size_t t = 0; t < rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > rounded_size) copy_size = rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH) { + AVX_Data src_4; + src_4.data = SIMD_LOAD(src + i); + + SIMD_STORE(dest + i, src_4.data); + } + } + +#endif + + if (param_size > rounded_size) { +#pragma omp parallel for + for (size_t k = rounded_size; k < param_size; k++) { dest[k] = src[k]; } + } +} + +static void helper_memcpy_4(float* dest, float* src, size_t param_size) +{ + size_t rounded_size = 0; + +#if defined(__AVX512__) or defined(__AVX256__) + + rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2)); + + for (size_t t = 0; t < rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > rounded_size) copy_size = rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += (SIMD_WIDTH << 2)) { + AVX_Data src_4[4]; + src_4[0].data = SIMD_LOAD(src + i); + src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH); + src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1)); + src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3); + + SIMD_STORE(dest + i, src_4[0].data); + SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data); + SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data); + SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data); + } + } +#endif + if (param_size > rounded_size) + helper_memcpy_1((dest + rounded_size), (src + rounded_size), (param_size - rounded_size)); +} + +static void helper_mempcy_8(float* dest, float* src, size_t param_size) +{ + size_t rounded_size = 0; + +#if defined(__AVX512__) or defined(__AVX256__) + + rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2)); + + for (size_t t = 0; t < rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > rounded_size) copy_size = rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += (SIMD_WIDTH << 3)) { + AVX_Data src_4[8]; + src_4[0].data = SIMD_LOAD(src + i); + src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH); + src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1)); + src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3); + src_4[4].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 2)); + src_4[5].data = SIMD_LOAD(src + i + SIMD_WIDTH * 5); + src_4[6].data = SIMD_LOAD(src + i + SIMD_WIDTH * 6); + src_4[7].data = SIMD_LOAD(src + i + SIMD_WIDTH * 7); + + SIMD_STORE(dest + i, src_4[0].data); + SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data); + SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data); + SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data); + SIMD_STORE(dest + i + (SIMD_WIDTH << 2), src_4[4].data); + SIMD_STORE(dest + i + SIMD_WIDTH * 5, src_4[5].data); + SIMD_STORE(dest + i + SIMD_WIDTH * 6, src_4[6].data); + SIMD_STORE(dest + i + SIMD_WIDTH * 7, src_4[7].data); + } + } +#endif + if (param_size > rounded_size) + helper_memcpy_4((dest + rounded_size), (src + rounded_size), (param_size - rounded_size)); +} + +int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src) +{ + auto dest_c = dest.contiguous(); + auto src_c = src.contiguous(); + + float* dest_ptr = (float*)dest_c.data_ptr(); + float* src_ptr = (float*)src_c.data_ptr(); + + helper_mempcy_8(dest_ptr, src_ptr, dest_c.size(0)); + + return 0; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h new file mode 100644 index 0000000000000000000000000000000000000000..f443571a3e7b35f3bc70ec835152c34e5d979137 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#if (__x86_64__ || __i386__) +#include +#include +#endif + +#include +#include +#include + +#define TILE (1024 * 1024 * 1024) + +#if defined(__AVX512__) +#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d) +#define SIMD_LOAD(x) _mm512_loadu_ps(x) +#define SIMD_SET(x) _mm512_set1_ps(x) +#define SIMD_MUL(x, y) _mm512_mul_ps(x, y) +#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) +#define SIMD_SQRT(x) _mm512_sqrt_ps(x) +#define SIMD_DIV(x, y) _mm512_div_ps(x, y) +#define SIMD_WIDTH 16 +#else +#if defined(__AVX256__) +#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d) +#define SIMD_LOAD(x) _mm256_loadu_ps(x) +#define SIMD_SET(x) _mm256_set1_ps(x) +#define SIMD_MUL(x, y) _mm256_mul_ps(x, y) +#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) +#define SIMD_SQRT(x) _mm256_sqrt_ps(x) +#define SIMD_DIV(x, y) _mm256_div_ps(x, y) +#define SIMD_WIDTH 8 +#endif +#endif + +int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp new file mode 100644 index 0000000000000000000000000000000000000000..48668a842949a718a2e4ab1fdac2b33026568433 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp @@ -0,0 +1,366 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_py_io_handle.h" +#include + +#define O_DIRECT_ALIGNMENT 512 + +using namespace std; + +static void _start_aio_thread(std::shared_ptr ctxt) { ctxt->run(); } + +static bool is_valid_bytes_to_read(const char* filename, + const int64_t file_offset, + const int64_t num_bytes_to_read) +{ + int64_t num_file_bytes; + if (-1 == get_file_size(filename, num_file_bytes)) { + const auto error_code = errno; + report_file_error(filename, " fstat for read", error_code); + return false; + } + if ((file_offset + num_bytes_to_read) > num_file_bytes) { + std::cout << filename << ": file_offset + buffer nbytes > file bytes " + << (file_offset + num_bytes_to_read) << " > " << num_file_bytes << std::endl; + } + assert((file_offset + num_bytes_to_read) <= num_file_bytes); + return true; +} + +deepspeed_io_handle_t::deepspeed_io_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism) + : _aio_ctxt(new aio_context(block_size, queue_depth)), + _single_submit(single_submit), + _overlap_events(overlap_events), + _intra_op_parallelism(intra_op_parallelism), + _aio_config(block_size, queue_depth, single_submit, overlap_events, false), + _num_pending_ops(0), + _pinned_tensor_mgr(new deepspeed_pin_tensor_t()) +{ + for (auto i = 0; i < intra_op_parallelism; ++i) { + _thread_contexts.push_back(std::make_shared(i, _aio_config)); + } + + for (auto& ctxt : _thread_contexts) { + _threads.push_back(std::thread(_start_aio_thread, ctxt)); + } +} + +deepspeed_io_handle_t::~deepspeed_io_handle_t() +{ + _stop_threads(); + for (auto& thr : _threads) { thr.join(); } +} + +const int deepspeed_io_handle_t::get_block_size() const +{ + return _aio_ctxt ? _aio_ctxt->_block_size : -1; +} + +const int deepspeed_io_handle_t::get_queue_depth() const +{ + return _aio_ctxt ? _aio_ctxt->_queue_depth : -1; +} + +const bool deepspeed_io_handle_t::get_single_submit() const { return _single_submit; } + +const bool deepspeed_io_handle_t::get_overlap_events() const { return _overlap_events; } + +const int deepspeed_io_handle_t::get_intra_op_parallelism() const { return _intra_op_parallelism; } + +const int deepspeed_io_handle_t::get_alignment() const +{ + return _intra_op_parallelism * O_DIRECT_ALIGNMENT; +} + +int deepspeed_io_handle_t::read(torch::Tensor& buffer, + const char* filename, + const bool validate, + const int64_t file_offset) +{ + const auto start_time = std::chrono::high_resolution_clock::now(); + + assert(_aio_ctxt); + + int64_t num_file_bytes; + if (-1 == get_file_size(filename, num_file_bytes)) { + const auto error_code = errno; + report_file_error(filename, " fstat for read", error_code); + return -1; + } + assert(static_cast(buffer.nbytes()) == num_file_bytes); + + const auto fd = open_file(filename, true); + if (fd == -1) { return -1; } + + auto read_buffer = (char*)buffer.data_ptr(); + std::unique_ptr xfer_ctxt( + new io_xfer_ctxt(fd, file_offset, 0, num_file_bytes, read_buffer)); + + if (_aio_config._overlap_events) { + do_aio_operation_overlap(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); + } else { + do_aio_operation_sequential(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); + } + + close(fd); + const std::chrono::duration aio_time = + std::chrono::high_resolution_clock::now() - start_time; + + if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); } + const std::chrono::duration fn_time = + std::chrono::high_resolution_clock::now() - start_time; + std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 + << " call = " << fn_time.count() * 1e6 << std::endl; + return 0; +} + +int deepspeed_io_handle_t::write(const torch::Tensor& buffer, + const char* filename, + const bool validate, + const int64_t file_offset) +{ + assert(_aio_ctxt); + + const auto start_time = std::chrono::high_resolution_clock::now(); + + const auto fd = open_file(filename, false); + if (fd == -1) { return -1; } + + auto write_buffer = (char*)buffer.data_ptr(); + const auto num_write_bytes = static_cast(buffer.nbytes()); + std::unique_ptr xfer_ctxt( + new io_xfer_ctxt(fd, file_offset, 0, num_write_bytes, write_buffer)); + + if (_aio_config._overlap_events) { + do_aio_operation_overlap(false, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); + } else { + do_aio_operation_sequential(false, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); + } + const std::chrono::duration aio_time = + std::chrono::high_resolution_clock::now() - start_time; + + close(fd); + + if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); } + + const std::chrono::duration fn_time = + std::chrono::high_resolution_clock::now() - start_time; + std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 + << " call = " << fn_time.count() * 1e6 << std::endl; + return 0; +} + +void deepspeed_io_handle_t::_schedule_aio_work(std::shared_ptr scheduled_op) +{ + for (auto& ctxt : _thread_contexts) { + { + std::lock_guard lock(ctxt->_work_sync._mutex); + ctxt->_work_queue.push(scheduled_op); + } + ctxt->_work_sync._cond_var.notify_one(); + } + _num_pending_ops++; +} + +std::shared_ptr deepspeed_io_handle_t::_wait_for_aio_work() +{ + std::shared_ptr completed_op = nullptr; + for (auto& ctxt : _thread_contexts) { + std::unique_lock lock(ctxt->_complete_sync._mutex); + ctxt->_complete_sync._cond_var.wait(lock, + [ctxt] { return !ctxt->_complete_queue.empty(); }); + completed_op = ctxt->_complete_queue.front(); + ctxt->_complete_queue.pop(); + } + return completed_op; +} + +void deepspeed_io_handle_t::_stop_threads() +{ + assert(0 == _num_pending_ops); + for (auto& ctxt : _thread_contexts) { + { + std::lock_guard lock(ctxt->_work_sync._mutex); + ctxt->_time_to_exit = true; + } + ctxt->_work_sync._cond_var.notify_one(); + } +} + +int deepspeed_io_handle_t::wait() +{ + assert(_num_pending_ops > 0); + auto num_completed_ops = 0; + + while (_num_pending_ops > 0) { + auto completed_op = _wait_for_aio_work(); + + if (completed_op->_validate) { completed_op->validate(); } + + completed_op->finish(); + + if (!completed_op->_filename.empty()) { (completed_op->_fd); } + + --_num_pending_ops; + ++num_completed_ops; + } + + return num_completed_ops; +} + +bool deepspeed_io_handle_t::_is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes) +{ + const auto op_string = read_op ? "Read" : "Write"; + if (num_bytes % get_intra_op_parallelism()) { + std::cout << "deepspeed_aio failure: parallel " << op_string << " num_bytes = " << num_bytes + << " not divisible by intra op parallelism = " << get_intra_op_parallelism() + << std::endl; + return false; + } + + return true; +} + +std::shared_ptr deepspeed_io_handle_t::_create_io_op_desc( + const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const int64_t file_offset) +{ + return std::make_shared(_pinned_tensor_mgr, + read_op, + buffer, + fd, + filename, + _intra_op_parallelism, + validate, + file_offset); +} + +int deepspeed_io_handle_t::_pread(const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset) +{ + auto scheduled_op = _create_io_op_desc(true, buffer, fd, filename, validate, file_offset); + + _schedule_aio_work(scheduled_op); + + if (async) { return 0; } + + return wait(); +} + +int deepspeed_io_handle_t::pread(const torch::Tensor& buffer, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset) +{ + const auto buffer_bytes = static_cast(buffer.nbytes()); + + if (!is_valid_bytes_to_read(filename, file_offset, buffer_bytes)) { return -1; } + + if (!_is_valid_parallel_aio_op(true, buffer_bytes)) { return -1; } + + const auto fd = open_file(filename, true); + if (fd == -1) { return -1; } + + return _pread(buffer, fd, filename, validate, async, file_offset); +} + +int deepspeed_io_handle_t::_pwrite(const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset) +{ + auto scheduled_op = _create_io_op_desc(false, buffer, fd, filename, validate, file_offset); + + _schedule_aio_work(scheduled_op); + + if (async) { return 0; } + + return wait(); +} + +int deepspeed_io_handle_t::pwrite(const torch::Tensor& buffer, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset) +{ + const auto num_write_bytes = static_cast(buffer.nbytes()); + + if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; } + + const auto fd = open_file(filename, false); + if (fd == -1) { return -1; } + + return _pwrite(buffer, fd, filename, validate, async, file_offset); +} + +int deepspeed_io_handle_t::sync_pread(torch::Tensor& buffer, + const char* filename, + const int64_t file_offset) +{ + return pread(buffer, filename, false, false, file_offset); +} + +int deepspeed_io_handle_t::sync_pwrite(const torch::Tensor& buffer, + const char* filename, + const int64_t file_offset) +{ + return pwrite(buffer, filename, false, false, file_offset); +} + +int deepspeed_io_handle_t::async_pread(torch::Tensor& buffer, + const char* filename, + const int64_t file_offset) +{ + return pread(buffer, filename, false, true, file_offset); +} + +int deepspeed_io_handle_t::async_pwrite(const torch::Tensor& buffer, + const char* filename, + const int64_t file_offset) +{ + return pwrite(buffer, filename, false, true, file_offset); +} + +int deepspeed_io_handle_t::async_pwrite(const torch::Tensor& buffer, + const int fd, + const int64_t file_offset = 0) +{ + const auto num_write_bytes = static_cast(buffer.nbytes()); + if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; } + + return _pwrite(buffer, fd, nullptr, false, true, file_offset); +} + +at::Tensor deepspeed_io_handle_t::new_cpu_locked_tensor(const int64_t num_elem, + const torch::Tensor& example_tensor) +{ + return _pinned_tensor_mgr->alloc(num_elem, example_tensor.scalar_type()); +} + +bool deepspeed_io_handle_t::free_cpu_locked_tensor(torch::Tensor& locked_tensor) +{ + return _pinned_tensor_mgr->free(locked_tensor); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h new file mode 100644 index 0000000000000000000000000000000000000000..8cb43c5b38e5d8f1981176c6459db0135c5f896d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h @@ -0,0 +1,109 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include "deepspeed_aio_thread.h" +#include "deepspeed_pin_tensor.h" + +struct deepspeed_io_handle_t { + std::unique_ptr _aio_ctxt; + const bool _single_submit; + const bool _overlap_events; + const int _intra_op_parallelism; + deepspeed_aio_config_t _aio_config; + + std::vector> _thread_contexts; + std::vector _threads; + int _num_pending_ops; + std::unique_ptr _pinned_tensor_mgr; + + deepspeed_io_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism); + + virtual ~deepspeed_io_handle_t() = 0; + + const int get_block_size() const; + const int get_queue_depth() const; + const bool get_single_submit() const; + const bool get_overlap_events() const; + const int get_intra_op_parallelism() const; + const int get_alignment() const; + + int read(torch::Tensor& buffer, + const char* filename, + const bool validate, + const int64_t file_offset); + + int write(const torch::Tensor& buffer, + const char* filename, + const bool validate, + const int64_t file_offset); + + int pread(const torch::Tensor& buffer, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset); + + int pwrite(const torch::Tensor& buffer, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset); + + int sync_pread(torch::Tensor& buffer, const char* filename, const int64_t file_offset); + + int sync_pwrite(const torch::Tensor& buffer, const char* filename, const int64_t file_offset); + + int async_pread(torch::Tensor& buffer, const char* filename, const int64_t file_offset); + + int async_pwrite(const torch::Tensor& buffer, const char* filename, const int64_t file_offset); + int async_pwrite(const torch::Tensor& buffer, const int fd, const int64_t file_offset); + + // TODO: Make API's args to be shape and dtype. + torch::Tensor new_cpu_locked_tensor(const int64_t num_elem, + const torch::Tensor& example_tensor); + + bool free_cpu_locked_tensor(torch::Tensor&); + + int wait(); + + void _stop_threads(); + + void _schedule_aio_work(std::shared_ptr scheduled_op); + + std::shared_ptr _wait_for_aio_work(); + + bool _is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes); + + int _pread(const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset); + + int _pwrite(const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const bool async, + const int64_t file_offset); + + virtual std::shared_ptr _create_io_op_desc(const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const int64_t file_offset); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp new file mode 100644 index 0000000000000000000000000000000000000000..62500bf4a6e997d4444dc8f6dfdfdc10173199f5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ +#include +#include "deepspeed_py_aio_handle.h" +#include "deepspeed_py_copy.h" +using namespace pybind11::literals; + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("aio_read", &deepspeed_py_aio_read, "DeepSpeed Asynchronous I/O Read"); + + m.def("aio_write", &deepspeed_py_aio_write, "DeepSpeed Asynchronous I/O Write"); + + m.def("deepspeed_memcpy", &deepspeed_py_memcpy, "DeepSpeed Memory Copy"); + + py::class_(m, "aio_handle") + .def(py::init(), + "AIO handle constructor", + "block_size"_a = 1024 * 1024, + "queue_depth"_a = 128, + "single_submit"_a = false, + "overlap_events"_a = false, + "intra_op_parallelism"_a = 1) + + .def("get_block_size", &deepspeed_aio_handle_t::get_block_size) + .def("get_queue_depth", &deepspeed_aio_handle_t::get_queue_depth) + .def("get_single_submit", &deepspeed_aio_handle_t::get_single_submit) + .def("get_overlap_events", &deepspeed_aio_handle_t::get_overlap_events) + .def("get_intra_op_parallelism", &deepspeed_aio_handle_t::get_intra_op_parallelism) + .def("get_alignment", &deepspeed_aio_handle_t::get_alignment) + + .def("read", + &deepspeed_aio_handle_t::read, + "Synchronous and non-parallel file read. Returns count of completed read ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "file_offset"_a = 0) + + .def("write", + &deepspeed_aio_handle_t::write, + "Synchronous and non-parallel file write. Returns count of completed write ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "file_offset"_a = 0) + + .def("pread", + &deepspeed_aio_handle_t::pread, + "Parallel file read with option of asynchronous completion. If synchronous, returns " + "count of completed read ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "async"_a, + "file_offset"_a = 0) + + .def("pwrite", + &deepspeed_aio_handle_t::pwrite, + "Parallel file write with option of asynchronous completion. If synchronous, returns " + "count of completed write ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "async"_a, + "file_offset"_a = 0) + + .def("sync_pread", + &deepspeed_aio_handle_t::sync_pread, + "Synchronous parallel file read. Returns count of completed read ops", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def("sync_pwrite", + &deepspeed_aio_handle_t::sync_pwrite, + "Synchronous parallel file write. Returns count of completed write ops", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def("async_pread", + &deepspeed_aio_handle_t::async_pread, + "Asynchronous parallel file read. Returns 0 on success. Returns 0 on success, and " + "subsequent wait() returns count of completed ops.", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def( + "async_pwrite", + py::overload_cast( + &deepspeed_aio_handle_t::async_pwrite), + "Asynchronous parallel file write. Returns 0 on success, and subsequent wait() returns " + "count of completed ops.", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def("async_pwrite", + py::overload_cast( + &deepspeed_aio_handle_t::async_pwrite), + "Asynchronous parallel file write using opened python file object.", + "buffer"_a, + "fd"_a, + "file_offset"_a = 0) + + .def("new_cpu_locked_tensor", + &deepspeed_aio_handle_t::new_cpu_locked_tensor, + "Allocate pinned CPU tensor.", + "num_elem"_a, + "example_tenosr"_a) + + .def("free_cpu_locked_tensor", + &deepspeed_aio_handle_t::free_cpu_locked_tensor, + "Free pinned CPU tensor.", + "tensor"_a) + + .def("wait", + &deepspeed_aio_handle_t::wait, + "Wait for (ongoing) asynchronous operations to complete"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json new file mode 100644 index 0000000000000000000000000000000000000000..4a224711b5ff4330ca704af0aaa9a4a4b0ea8ca3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json @@ -0,0 +1,34 @@ +{ + "block_size": [ + "128K", + "256K", + "1M", + "2M", + "4M", + "8M", + "16M" + ], + "queue_depth": [ + 8, + 16, + 32, + 64 + ], + "io_parallel": [ + 1, + 2, + 4, + 8 + ], + "single_submit": [ + true, + false + ], + "sequential_requests": [ + true, + false + ], + "threads": [ + 1 + ] +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2eca0a33262e69cf8901c8293d6c7859b2619149 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp @@ -0,0 +1,188 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#define USE_C10D_NCCL + +namespace dc { + +std::shared_ptr param_registry; +std::unordered_map> executors; +std::shared_ptr reduce_buckets = nullptr; + +c10::intrusive_ptr process_group = nullptr; +c10::intrusive_ptr symm_mem = nullptr; +ncclComm_t nccl_comm; +bool use_symm_mem; +bool clone_custom_op_output; +bool profile = false; +bool pre_div_reduce = true; + +bool sync_before_reduce; // for debugging +bool sync_after_reduce; // for debugging +bool sync_before_allgather; // for debugging +bool sync_after_allgather; // for debugging + +std::vector sizes_to_int_vector(at::IntArrayRef sizes) +{ + std::vector result; + for (int i = 0; i < sizes.size(); i++) { result.push_back(sizes[i]); } + return result; +} + +void enable_profiling(bool enable) { profile = enable; } + +bool is_profiling() { return profile; } + +c10::intrusive_ptr getSymmMemWorkspace(int64_t size) +{ + c10::Device device = c10::Device(c10::kCUDA, c10::cuda::current_device()); + std::vector sizes = {size}; + std::vector strides = {1}; + at::Tensor sym_mem_ws = c10d::symmetric_memory::empty_strided_p2p( + {size}, {1}, c10::ScalarType::Byte, device, process_group->getGroupName(), std::nullopt); + return c10d::symmetric_memory::rendezvous(sym_mem_ws); +} + +void lazy_init_symm_memory() +{ + if (use_symm_mem && !symm_mem) { + int64_t max_param_size = 0; + for (const auto& it : param_registry->getParams()) { + int64_t size = it.second.getDSTensor().numel() * it.second.getDSTensor().element_size(); + if (size > max_param_size) { max_param_size = size; } + } + symm_mem = getSymmMemWorkspace(max_param_size); + } +} + +ncclDataType_t get_nccl_data_type(at::ScalarType scalar_type) +{ + switch (scalar_type) { + case at::kFloat: return ncclFloat; + case at::kHalf: return ncclHalf; + case at::kDouble: return ncclDouble; + case at::kBFloat16: return ncclBfloat16; + case at::kLong: return ncclInt64; + case at::kInt: return ncclInt; + case at::kChar: return ncclInt8; + default: throw std::runtime_error("Unsupported scalar type"); + } +} + +void reset() +{ + executors.clear(); + // We keep the buckets for memory estimation + // reduce_buckets->clear(); +} + +void cleanup() +{ + reset(); + + ncclCommDestroy(nccl_comm); + process_group = nullptr; + symm_mem = nullptr; +} + +at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id) +{ + if (sync_before_reduce) { c10::cuda::device_synchronize(); } + + assert(hasKey(executors, graph_id)); + if (!profile) { executors[graph_id]->reduceGrad(grad_tensor, ds_id); } + + if (sync_after_reduce) { c10::cuda::device_synchronize(); } + + return at::Tensor(); +} + +at::Tensor reduce_grad_meta(at::Tensor grad_tensor, long graph_id, long ds_id) +{ + return at::Tensor(); +} + +void free_tensors(std::vector tensors) +{ + int64_t THRESHOLD = 10 * 1024 * 1024; + + if (!profile) { + for (auto& tensor : tensors) { + if (tensor.is_cuda() && tensor.numel() > THRESHOLD) { + tensor.record_stream(at::cuda::getCurrentCUDAStream()); + tensor.set_data(torch::empty({0}, tensor.options())); + } + } + } +} + +void free_tensors_meta(std::vector tensors) {} + +void init(c10::intrusive_ptr pg, + int64_t initial_reduce_bucket_size, + bool enable_double_buffer, + bool _use_symm_mem, + bool _clone_custom_op_output, + bool _sync_before_reduce, + bool _sync_after_reduce, + bool _sync_before_allgather, + bool _sync_after_allgather) +{ + process_group = pg; + + ncclUniqueId ncclID; + ncclGetUniqueId(&ncclID); + + // ProcessGroup doesn't have an API to get the CUDA stream for comm calls. + // So we create a NCCL communicator and call NCCL APIs directly. + auto vec = std::vector(reinterpret_cast(&ncclID), + reinterpret_cast(&ncclID) + NCCL_UNIQUE_ID_BYTES); + auto device = torch::Device(torch::kCUDA); + at::Tensor tensor = torch::from_blob(vec.data(), {static_cast(vec.size())}, torch::kUInt8) + .to(torch::Device(torch::kCUDA)); + std::vector bcast_input = {tensor}; + + process_group->broadcast(bcast_input, c10d::BroadcastOptions())->wait(); + + // create a new nccl communicator + std::memcpy(&ncclID, tensor.to(torch::Device(torch::kCPU)).data_ptr(), NCCL_UNIQUE_ID_BYTES); + ncclCommInitRank(&nccl_comm, process_group->getSize(), ncclID, process_group->getRank()); + + param_registry = std::make_shared(); + reduce_buckets = std::make_shared(initial_reduce_bucket_size, + enable_double_buffer); + use_symm_mem = _use_symm_mem; + clone_custom_op_output = _clone_custom_op_output; + + sync_before_reduce = _sync_before_reduce; + sync_after_reduce = _sync_after_reduce; + sync_before_allgather = _sync_before_allgather; + sync_after_allgather = _sync_after_allgather; +} + +void start_forward() +{ + lazy_init_symm_memory(); + for (auto& it : executors) { it.second->startForward(); } +} + +void end_forward() +{ + for (auto& it : executors) { it.second->endForward(); } +} + +void start_backward(bool update) +{ + for (auto& it : executors) { it.second->startBackward(update); } +} + +// We don't call this +// void end_backward(bool update) +// { +// } + +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/init.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/init.cpp new file mode 100644 index 0000000000000000000000000000000000000000..9b80fff4e6a7fab16ed8dacaf2af30a55ea71567 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/init.cpp @@ -0,0 +1,103 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" +#include "z1.h" +#include "z2.h" +#include "z3.h" + +TORCH_LIBRARY(dc, m) +{ + m.def("allgather_param(Tensor a, int graph_id, int id) -> Tensor"); + m.def("prefetch_params_fused(int graph_id, Tensor[] params, int[] ids) -> ()"); + m.def("wait_allgather(Tensor a, int graph_id, int id) -> Tensor"); + m.def("release_param(Tensor a, int graph_id, int id, int n_users) -> Tensor"); + m.def("reduce_grad(Tensor a, int graph_id, int id) -> Tensor"); + m.def("free_tensors(Tensor[] a) -> ()"); + m.def("offload_tensor(Tensor a, int id, int id) -> Tensor"); + m.def("reload_tensor(Tensor a, int id, int id) -> Tensor"); + m.def("wait_offload(Tensor a, int id, int id) -> Tensor"); + m.def("wait_reload(Tensor a, int id, int id) -> Tensor"); + m.def("offload_parameter(Tensor a, int id, int id) -> ()"); + m.def("reload_parameter(Tensor a, int id, int id) -> ()"); + + m.def("test_call(Tensor a) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(dc, CPU, m) +{ + m.impl("allgather_param", &dc::allgather_param); + m.impl("prefetch_params_fused", &dc::prefetch_params_fused); + m.impl("wait_allgather", &dc::wait_allgather); + m.impl("release_param", &dc::release_param); + m.impl("reduce_grad", &dc::reduce_grad); + m.impl("free_tensors", &dc::free_tensors); + m.impl("offload_tensor", &dc::offload_tensor); + m.impl("reload_tensor", &dc::reload_tensor); + m.impl("wait_offload", &dc::wait_offload); + m.impl("wait_reload", &dc::wait_reload); + m.impl("offload_parameter", &dc::offload_parameter); + m.impl("reload_parameter", &dc::reload_parameter); + + m.impl("test_call", &dc::test_call); +} + +TORCH_LIBRARY_IMPL(dc, CUDA, m) +{ + m.impl("allgather_param", &dc::allgather_param); + m.impl("prefetch_params_fused", &dc::prefetch_params_fused); + m.impl("wait_allgather", &dc::wait_allgather); + m.impl("release_param", &dc::release_param); + m.impl("reduce_grad", &dc::reduce_grad); + m.impl("free_tensors", &dc::free_tensors); + m.impl("offload_tensor", &dc::offload_tensor); + m.impl("reload_tensor", &dc::reload_tensor); + m.impl("wait_offload", &dc::wait_offload); + m.impl("wait_reload", &dc::wait_reload); + m.impl("offload_parameter", &dc::offload_parameter); + m.impl("reload_parameter", &dc::reload_parameter); + + m.impl("test_call", &dc::test_call); +} + +TORCH_LIBRARY_IMPL(dc, Meta, m) +{ + m.impl("allgather_param", &dc::allgather_param_meta); + m.impl("prefetch_params_fused", &dc::prefetch_params_fused_meta); + m.impl("release_param", &dc::release_param_meta); + m.impl("wait_allgather", &dc::wait_allgather_meta); + m.impl("reduce_grad", &dc::reduce_grad_meta); + m.impl("free_tensors", &dc::free_tensors_meta); + m.impl("reload_parameter", &dc::reload_parameter_meta); + m.impl("offload_parameter", &dc::offload_parameter_meta); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("set_persistent", &dc::set_persistent, "Set persistent flag for a parameter"); + m.def("enable_profiling", &dc::enable_profiling, "Enable profiling"); + m.def("is_profiling", &dc::is_profiling, "Check if profiling is enabled"); + m.def("init", &dc::init, "Set the process group"); + m.def("cleanup", &dc::cleanup, "Cleanup the process group"); + m.def("register_param", &dc::register_param, "Register a parameter"); + m.def("register_graph_z1", + &dc::register_graph_z1, + "Register graph with a list of ds parameter ids"); + m.def("register_graph_z2", + &dc::register_graph_z2, + "Register graph with a list of ds parameter ids"); + m.def("register_z3_param", &dc::register_z3_param, "Register a parameter"); + m.def("register_graph_z3", + &dc::register_graph_z3, + "Register graph with a list of ds parameter ids"); + m.def("start_forward", &dc::start_forward, "Start forward pass"); + m.def("end_forward", &dc::end_forward, "End forward pass"); + m.def("start_backward", &dc::start_backward, "Start backward pass"); + // m.def("end_backward", &dc::end_backward, "End backward pass"); + m.def("cleanup", &dc::cleanup, "Clean up DeepCompile"); + m.def("reset", &dc::reset, "Reset the state"); + m.def("invalidate_gathered_param", &dc::invalidate_gathered_param, "Invalidate gathered param"); + m.def("clear_all_gathered_params", &dc::clear_all_gathered_params, "Clear all gathered params"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/util.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/util.cpp new file mode 100644 index 0000000000000000000000000000000000000000..948338028059bc1eff42e00d1d7390c09b58fbb0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/util.cpp @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#include + +namespace dc { + +std::string tensorToString(const at::Tensor& t, size_t max_elem, size_t max_str_len) +{ + auto t_cpu = t.flatten() + .slice(0, 0, std::min((int64_t)max_elem, t.numel())) + .to(c10::Device(c10::kCPU), false, true); + + size_t size = std::min(max_elem, productDim(t.sizes())); + + if (t.scalar_type() == c10::ScalarType::Half || t.scalar_type() == c10::ScalarType::BFloat16) { + auto float_ten = t_cpu.to(c10::ScalarType::Float, false, true).contiguous(); + return tensorPtrToString((float*)float_ten.data_ptr(), size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Float) { + return tensorPtrToString((float*)t_cpu.data_ptr(), size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Double) { + return tensorPtrToString((double*)t_cpu.data_ptr(), size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Int) { + int* ptr = static_cast(t_cpu.data_ptr()); + return tensorPtrToString(ptr, size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Long) { + long* ptr = static_cast(t_cpu.data_ptr()); + return tensorPtrToString(ptr, size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Byte) { + unsigned char* ptr = static_cast(t_cpu.data_ptr()); + std::vector vec; + vec.reserve(size); + for (size_t i = 0; i < size; i++) { + vec.push_back(*ptr); + ptr++; + } + return tensorPtrToString(&vec[0], size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Bool) { + bool* ptr = static_cast(t_cpu.data_ptr()); + std::vector vec; + vec.reserve(size); + for (size_t i = 0; i < size; i++) { + vec.push_back(*ptr); + ptr++; + } + return tensorPtrToString(&vec[0], size, max_str_len); + } + std::stringstream ss; + ss << "Failed to convert tensor to string. Invalid type of tensor: " + << toString(t.scalar_type()); + throw std::invalid_argument(ss.str()); +} + +std::string tensorPtrToString(void* ptr, + size_t size, + c10::ScalarType datatype, + size_t max_elem, + size_t max_str_len) +{ + int64_t elem_size = std::min((size_t)max_elem, size); + + if (datatype == c10::ScalarType::Long) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Int) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Double) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Float) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Half || datatype == c10::ScalarType::BFloat16) { + const auto ten = torch::from_blob(ptr, {(int64_t)elem_size}, datatype); + auto float_ten = ten.to(c10::ScalarType::Float, false, true).contiguous(); + return tensorPtrToString((float*)float_ten.data_ptr(), elem_size, max_str_len); + } + std::stringstream ss; + ss << "Failed to convert tensor ptr to string. Invalid type of tensor: " << toString(datatype); + throw std::invalid_argument(ss.str()); +} + +std::string tensorDimToString(const at::Tensor& t) +{ + const auto dim = t.sizes(); + return join_as_str(dim); +} +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.cpp new file mode 100644 index 0000000000000000000000000000000000000000..cbec2dec82ab13e055a433a4f2f96c704f3709cc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.cpp @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "z1.h" +#include "deepcompile.h" + +namespace dc { + +class Z1CustomOpExecutor : public CustomOpExecutor { +public: + Z1CustomOpExecutor(c10::intrusive_ptr process_group, + std::shared_ptr param_registry, + std::shared_ptr reduce_buckets, + std::vector ds_ids, + ncclComm_t nccl_comm, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream, + bool pre_div_reduce) + : CustomOpExecutor(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce) + { + } + ~Z1CustomOpExecutor() {} + + at::Tensor reduceGrad(at::Tensor grad_tensor, long ds_id) override + { + if (!hasKey(grad_tensors_, ds_id)) { + grad_tensors_[ds_id] = grad_tensor; + } else { + grad_tensors_[ds_id].add_(grad_tensor); + } + + if (param_updated_) { + CustomOpExecutor::reduceGrad(grad_tensors_[ds_id], ds_id); + grad_tensors_.erase(ds_id); + } + + return at::Tensor(); + } + + void flushReduceBucket(at::ScalarType scalar_type) override + { + if (!hasKey(reduce_tasks_, scalar_type)) { return; } + + blockCopyEvents(scalar_type); + applyPreDivision(scalar_type); + + // NCCL AllReduce operation + ncclGroupStart(); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(), + t.getSendBuf().data_ptr(), + t.getSendBuf().numel(), + get_nccl_data_type(scalar_type), + getReductionOp(), + nccl_comm_, + rs_stream_); + if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); } + } + ncclGroupEnd(); + + // Copy results to gradient buffers + { + at::cuda::CUDAStreamGuard guard(rs_stream_); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto param = param_registry_->getParam(t.getDSId()); + auto grad_buf = param.getGradBuffer().flatten(); + + if (grad_buf.numel() == 0) { continue; } + + int64_t offset = param.getOffset(); + auto recv_buf = t.getSendBuf().flatten().index( + {torch::indexing::Slice(offset, offset + grad_buf.numel())}); + grad_buf.copy_(recv_buf); + } + } + + performCleanup(scalar_type); + } + +protected: + std::unordered_map grad_tensors_; +}; + +static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); + +void register_graph_z1(long graph_id, const std::vector& ds_ids) +{ + executors[graph_id] = std::make_shared(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce); +} + +void register_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + int64_t offset) +{ + param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, false, offset, false); +} + +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.h new file mode 100644 index 0000000000000000000000000000000000000000..1d3607a59b066559ab73b11d23d292e70afd432f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.h @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#pragma once + +namespace dc { + +void register_graph_z1(long graph_id, const std::vector& ds_ids); +void register_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + int64_t offset); +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2effa83f02e516b2aeb8f60b6ba1d2f1dd97cb61 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.cpp @@ -0,0 +1,102 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "z2.h" +#include "deepcompile.h" + +namespace dc { + +class Z2CustomOpExecutor : public CustomOpExecutor { +public: + Z2CustomOpExecutor(c10::intrusive_ptr process_group, + std::shared_ptr param_registry, + std::shared_ptr reduce_buckets, + std::vector ds_ids, + ncclComm_t nccl_comm, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream, + bool pre_div_reduce) + : CustomOpExecutor(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce) + { + } + ~Z2CustomOpExecutor() {} + + void endBackward() override + { + if (param_updated_) { + for (auto& it : has_acc_grad_) { it.second = false; } + } + } + + void flushReduceBucket(at::ScalarType scalar_type) override + { + if (!hasKey(reduce_tasks_, scalar_type)) { return; } + + blockCopyEvents(scalar_type); + applyPreDivision(scalar_type); + + // NCCL AllReduce operation + ncclGroupStart(); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(), + t.getSendBuf().data_ptr(), + t.getSendBuf().numel(), + get_nccl_data_type(scalar_type), + getReductionOp(), + nccl_comm_, + rs_stream_); + if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); } + } + ncclGroupEnd(); + + // Copy or accumulate results to gradient buffers + { + at::cuda::CUDAStreamGuard guard(rs_stream_); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + bool acc_grad = has_acc_grad_.at(t.getDSId()); + auto param = param_registry_->getParam(t.getDSId()); + auto grad_buf = param.getGradBuffer().flatten(); + + if (grad_buf.numel() == 0) { continue; } + + int64_t offset = param.getOffset(); + auto recv_buf = t.getSendBuf().flatten().index( + {torch::indexing::Slice(offset, offset + grad_buf.numel())}); + if (acc_grad) { + grad_buf.add_(recv_buf); + } else { + grad_buf.copy_(recv_buf); + } + has_acc_grad_[t.getDSId()] = true; + } + } + + performCleanup(scalar_type); + } +}; + +static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); + +void register_graph_z2(long graph_id, const std::vector& ds_ids) +{ + executors[graph_id] = std::make_shared(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce); +} + +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.h new file mode 100644 index 0000000000000000000000000000000000000000..cc6c3136c20c9c69d253caafde31c743d2243dda --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.h @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#pragma once + +namespace dc { + +void register_graph_z2(long graph_id, const std::vector& ds_ids); + +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.cpp new file mode 100644 index 0000000000000000000000000000000000000000..75352c2bafdeb548b08d0e7a2e3fbe84beaafa63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.cpp @@ -0,0 +1,524 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "z3.h" +#include "deepcompile.h" + +namespace dc { + +const size_t TIMEOUT_SYMMETRIC_MEMORY_BARRIER = 60000; + +class Z3CustomOpExecutor : public CustomOpExecutor { +public: + Z3CustomOpExecutor(c10::intrusive_ptr process_group, + std::shared_ptr param_registry, + std::shared_ptr reduce_buckets, + std::vector ds_ids, + ncclComm_t nccl_comm, + at::cuda::CUDAStream ag_stream, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream, + at::cuda::CUDAStream offload_stream, + at::cuda::CUDAStream reload_stream, + bool pre_div_reduce) + : CustomOpExecutor(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce), + ag_stream_(ag_stream), + offload_stream_(offload_stream), + reload_stream_(reload_stream) + { + for (long ds_id : ds_ids_) { + ag_comm_done_events_[ds_id] = + std::make_shared(cudaEventDisableTiming); + ag_comp_done_events_[ds_id] = + std::make_shared(cudaEventDisableTiming); + + param_use_count_[ds_id] = 0; + } + } + ~Z3CustomOpExecutor() {} + + void endBackward() override + { + if (param_updated_) { + for (auto& it : has_acc_grad_) { + it.second = false; + param_registry_->setValid(it.first, false); + } + } + + for (auto& it : reload_buffers_) { + it.second.record_stream(at::cuda::getCurrentCUDAStream()); + } + reload_buffers_.clear(); + } + + void launchAllGather(at::Tensor output_buf, + long ds_id, + c10::intrusive_ptr symm_mem) + { + const DSParam& param = param_registry_->getParam(ds_id); + const at::Tensor& ds_tensor = param.getDSTensor(); + + if (symm_mem == nullptr) { + ncclResult_t result = ncclAllGather(ds_tensor.contiguous().data_ptr(), + output_buf.data_ptr(), + ds_tensor.numel(), + get_nccl_data_type(ds_tensor.scalar_type()), + nccl_comm_, + ag_stream_); + + if (result != ncclSuccess) { throw std::runtime_error("NCCL AllGather failed"); } + } else { + at::cuda::CUDAStreamGuard guard(ag_stream_); + int world_size = process_group_->getSize(); + int rank = process_group_->getRank(); + + at::Tensor local_buf = + symm_mem->get_buffer(rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0); + local_buf.copy_(ds_tensor, true); + + symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER); + auto chunks = output_buf.flatten().chunk(world_size); + for (int step = 0; step < world_size; step++) { + int remote_rank = (rank - step + world_size) % world_size; + auto src_buf = symm_mem->get_buffer( + remote_rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0); + chunks[remote_rank].copy_(src_buf.flatten(), true); + } + symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER); + } + + param_registry_->registerGatheredParam(ds_id, output_buf); + param_registry_->setValid(ds_id, true); + } + + at::Tensor allgatherParam(long ds_id, + c10::intrusive_ptr symm_mem) + { + if (param_registry_->isValid(ds_id)) { return param_registry_->getGatheredParam(ds_id); } + + const DSParam& param = param_registry_->getParam(ds_id); + const at::Tensor& ds_tensor = param.getDSTensor(); + at::Tensor output_buf = param_registry_->hasGatheredParam(ds_id) + ? param_registry_->getGatheredParam(ds_id) + : torch::empty(param.getShape(), ds_tensor.options()); + + assert(hasKey(ag_comp_done_events_, ds_id)); + ag_comp_done_events_[ds_id]->record(); + ag_comp_done_events_[ds_id]->block(ag_stream_); + + launchAllGather(output_buf, ds_id, symm_mem); + + ag_comm_done_events_[ds_id]->record(ag_stream_); + return output_buf; + } + + void prefetchParamsFused(std::vector ds_ids, + c10::intrusive_ptr symm_mem) + { + std::vector invalid_ds_ids; + for (const auto& ds_id : ds_ids) { + if (!param_registry_->isValid(ds_id)) { invalid_ds_ids.push_back(ds_id); } + } + + std::unordered_map output_bufs; + for (long ds_id : invalid_ds_ids) { + const DSParam& param = param_registry_->getParam(ds_id); + if (param_registry_->hasGatheredParam(ds_id)) { + output_bufs[ds_id] = param_registry_->getGatheredParam(ds_id); + } else { + output_bufs[ds_id] = torch::empty(param.getShape(), param.getDSTensor().options()); + } + } + + for (long ds_id : invalid_ds_ids) { + ag_comp_done_events_[ds_id]->record(); + ag_comp_done_events_[ds_id]->block(ag_stream_); + } + + ncclGroupStart(); + for (long ds_id : invalid_ds_ids) { + assert(hasKey(output_bufs, ds_id)); + launchAllGather(output_bufs.at(ds_id), ds_id, symm_mem); + } + ncclGroupEnd(); + + for (long ds_id : invalid_ds_ids) { ag_comm_done_events_[ds_id]->record(ag_stream_); } + } + + void releaseParam(long ds_id, long n_users) + { + const DSParam& param = param_registry_->getParam(ds_id); + + assert(hasKey(param_use_count_, ds_id)); + if (param_use_count_[ds_id] == 0) { param_use_count_[ds_id] = n_users; } + param_use_count_[ds_id]--; + + if (param_use_count_[ds_id] == 0 && !param.isPersistent()) { + at::Tensor gathered_param = param_registry_->getGatheredParam(ds_id); + + if (gathered_param.defined()) { // gathered param is undefined while profiling + const auto options = gathered_param.options(); + at::Tensor empty_buffer = torch::empty({0}, options); + gathered_param.set_data(empty_buffer); + } + + param_registry_->unregisterGatheredParam(ds_id); + } + } + + at::Tensor waitAllgather(at::Tensor v, long ds_id) + { + assert(hasKey(ag_comm_done_events_, ds_id)); + ag_comm_done_events_[ds_id]->block(at::cuda::getCurrentCUDAStream()); + return v; + } + + void flushReduceBucket(at::ScalarType scalar_type) override + { + if (!hasKey(reduce_tasks_, scalar_type)) { return; } + + blockCopyEvents(scalar_type); + + // Calculate temporary buffer size for accumulated gradients + int64_t tmp_recv_numel = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + if (has_acc_grad_.at(t.getDSId())) { + tmp_recv_numel += param_registry_->getParam(t.getDSId()).getGradBuffer().numel(); + } + } + + // Allocate temporary buffer if needed + at::Tensor tmp_recv_buf = at::Tensor(); + if (tmp_recv_numel > 0) { + at::cuda::CUDAStreamGuard guard(rs_stream_); + tmp_recv_buf = torch::empty({tmp_recv_numel}, + at::TensorOptions().dtype(scalar_type).device(at::kCUDA)); + } + + applyPreDivision(scalar_type); + + // NCCL ReduceScatter operation + ncclGroupStart(); + int64_t offset = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto recv_buf = param_registry_->getParam(t.getDSId()).getGradBuffer(); + bool acc_grad = has_acc_grad_.at(t.getDSId()); + + if (acc_grad) { + recv_buf = + tmp_recv_buf.index({torch::indexing::Slice(offset, offset + recv_buf.numel())}); + } + + ncclResult_t result = ncclReduceScatter(t.getSendBuf().data_ptr(), + recv_buf.data_ptr(), + recv_buf.numel(), + get_nccl_data_type(scalar_type), + getReductionOp(), + nccl_comm_, + rs_stream_); + if (result != ncclSuccess) { throw std::runtime_error("NCCL ReduceScatter failed"); } + + if (acc_grad) { offset += recv_buf.numel(); } + } + ncclGroupEnd(); + + // Handle gradient accumulation with temporary buffer + { + at::cuda::CUDAStreamGuard guard(rs_stream_); + int64_t offset = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + bool acc_grad = has_acc_grad_.at(t.getDSId()); + + if (acc_grad) { + auto recv_buf = param_registry_->getParam(t.getDSId()).getGradBuffer(); + recv_buf.add_(tmp_recv_buf.index( + {torch::indexing::Slice(offset, offset + recv_buf.numel())})); + offset += recv_buf.numel(); + } + has_acc_grad_[t.getDSId()] = true; + } + } + + performCleanup(scalar_type); + + // Record stream for temporary buffer to prevent early deallocation + if (tmp_recv_numel > 0) { tmp_recv_buf.record_stream(rs_stream_); } + } + + at::Tensor offloadTensor(at::Tensor tensor, long id) + { + if (!hasKey(offload_events_, id)) { + offload_events_[id] = std::make_shared(cudaEventDisableTiming); + offload_comp_done_events_[id] = + std::make_shared(cudaEventDisableTiming); + + const auto options = at::TensorOptions().pinned_memory(true).device(torch::kCPU); + offload_buffers_[id] = at::empty_like(tensor, options); + } + + offload_comp_done_events_[id]->record(); + offload_comp_done_events_[id]->block(offload_stream_); + { + at::cuda::CUDAStreamGuard guard(offload_stream_); + offload_buffers_.at(id).copy_(tensor, true); + } + + tensor.record_stream(offload_stream_); + + offload_events_[id]->record(offload_stream_); + assert(hasKey(offload_buffers_, id)); + return offload_buffers_.at(id); + } + + at::Tensor reloadTensor(at::Tensor tensor, long id) + { + if (!hasKey(reload_events_, id)) { + reload_events_[id] = std::make_shared(cudaEventDisableTiming); + } + + assert(hasKey(offload_buffers_, id)); + offload_events_[id]->block(reload_stream_); + + at::Tensor ten; + { + at::cuda::CUDAStreamGuard guard(reload_stream_); + + assert(hasKey(offload_buffers_, id)); + at::Tensor buf = offload_buffers_.at(id); + const auto options = at::TensorOptions().device(torch::kCUDA); + ten = at::empty_like(buf, options); + ten.copy_(buf, true); + + reload_buffers_[id] = ten; + } + + reload_events_[id]->record(reload_stream_); + return ten; + } + + at::Tensor waitOffload(at::Tensor tensor, long id) + { + assert(hasKey(offload_events_, id)); + offload_events_[id]->block(at::cuda::getCurrentCUDAStream()); + + assert(hasKey(offload_buffers_, id)); + return offload_buffers_.at(id); + } + + at::Tensor waitReload(at::Tensor tensor, long id) + { + assert(hasKey(reload_events_, id)); + reload_events_[id]->block(at::cuda::getCurrentCUDAStream()); + + assert(hasKey(reload_buffers_, id)); + auto ten = reload_buffers_.at(id); + + // We can't release here because the tensor is still being used + // We will need "freeReloadedTensor" after the last user of the tensor to call + // ".record_stream". As it is a bit complicated, we clear the buffer and do at the end of + // the backward pass for now. reload_buffers_.erase(id); + return ten; + } + + void offloadParameter(at::Tensor tensor, long ds_id) { param_registry_->offload(ds_id); } + void reloadParameter(at::Tensor tensor, long ds_id) { param_registry_->reload(ds_id); } + + bool hasReloadBuffer(long id) { return hasKey(reload_buffers_, id); } + + bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); } + +private: + at::cuda::CUDAStream ag_stream_; + at::cuda::CUDAStream offload_stream_; + at::cuda::CUDAStream reload_stream_; + + std::unordered_map> ag_comp_done_events_; + std::unordered_map> ag_comm_done_events_; + + std::unordered_map> offload_events_; + std::unordered_map> offload_comp_done_events_; + std::unordered_map> reload_events_; + std::unordered_map offload_buffers_; + std::unordered_map reload_buffers_; + + std::unordered_map param_use_count_; +}; + +static at::cuda::CUDAStream ag_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream offload_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream reload_stream = at::cuda::getStreamFromPool(true); + +void register_graph_z3(long graph_id, const std::vector& ds_ids) +{ + executors[graph_id] = std::make_shared(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + ag_stream, + rs_stream, + copy_stream, + offload_stream, + reload_stream, + pre_div_reduce); +} + +void register_z3_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + bool persistent) +{ + param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, true, 0, persistent); + if (persistent) { param_registry->registerGatheredParam(ds_id, ds_tensor); } +} + +at::Tensor allgather_param(at::Tensor param_tensor, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + + if (sync_before_allgather) { c10::cuda::device_synchronize(); } + auto ret = executor->allgatherParam(ds_id, symm_mem); + if (sync_after_allgather) { c10::cuda::device_synchronize(); } + return ret; +} + +void set_persistent(long ds_id) +{ + param_registry->setPersistent(ds_id, true); + + // Allocate buffer here + // Memory fragmentation will be more severe if we allocate in forward/backward + for (auto& it : executors) { + if (it.second->hasParam(ds_id)) { + auto executor = getExecutor(it.first, executors); + executor->allgatherParam(ds_id, symm_mem); + } + } +} + +void prefetch_params_fused(long graph_id, + const std::vector params, + const std::vector& ds_ids) +{ + auto executor = getExecutor(graph_id, executors); + executor->prefetchParamsFused(ds_ids, symm_mem); +} + +void prefetch_params_fused_meta(long graph_id, + const std::vector params, + const std::vector& ds_ids) +{ +} + +// for profiling +void invalidate_gathered_param(long ds_id) +{ + const DSParam& param = param_registry->getParam(ds_id); + if (param.isPersistent()) { return; } + + param_registry->unregisterGatheredParam(ds_id); + param_registry->registerGatheredParam(ds_id, at::Tensor()); +} + +void clear_all_gathered_params() +{ + for (const auto& it : param_registry->getParams()) { + long ds_id = it.first; + const DSParam& param = param_registry->getParam(ds_id); + if (param.isPersistent()) { continue; } + if (param_registry->hasGatheredParam(ds_id)) { + param_registry->unregisterGatheredParam(ds_id); + } + } +} + +at::Tensor allgather_param_meta(at::Tensor param_tensor, long graph_id, long ds_id) +{ + const DSParam& param = param_registry->getParam(ds_id); + auto options = param.getDSTensor().options().device(c10::kMeta); + at::Tensor output_buf = torch::empty(param.getShape(), options); + return output_buf; +} + +at::Tensor release_param(at::Tensor dummy, long graph_id, long ds_id, long n_users) +{ + auto executor = getExecutor(graph_id, executors); + executor->releaseParam(ds_id, n_users); + + if (clone_custom_op_output) { return dummy.clone(); } + return dummy; +} + +at::Tensor release_param_meta(at::Tensor dummy, long graph_id, long ds_id, long n_users) +{ + return dummy; +} + +at::Tensor wait_allgather(at::Tensor v, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + executor->waitAllgather(v, ds_id); + return v; +} + +at::Tensor wait_allgather_meta(at::Tensor v, long graph_id, long ds_id) { return v; } + +at::Tensor offload_tensor(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + return executor->offloadTensor(tensor, id); +} + +at::Tensor reload_tensor(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + return executor->reloadTensor(tensor, id); +} + +at::Tensor wait_offload(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + return executor->waitOffload(tensor, id); +} + +at::Tensor wait_reload(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + if (profile && !executor->hasReloadBuffer(id)) { return tensor; } + return executor->waitReload(tensor, id); +} + +at::Tensor test_call(at::Tensor a) +{ + std::cout << "test_call" << std::endl; + return a; +} + +void reload_parameter(at::Tensor tensor, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + executor->reloadParameter(tensor, ds_id); +} + +void offload_parameter(at::Tensor tensor, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + executor->offloadParameter(tensor, ds_id); +} +void reload_parameter_meta(at::Tensor param_tensor, long graph_id, long ds_id) {} +void offload_parameter_meta(at::Tensor tensor, long graph_id, long ds_id) {} + +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.h new file mode 100644 index 0000000000000000000000000000000000000000..1031f0c84f7c1447d662c1df4629f50c7c47fce1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.h @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#pragma once + +namespace dc { + +void register_graph_z3(long graph_id, const std::vector& ds_ids); +void register_graph_ops_z3(long graph_id, + const std::vector& op_names, + const std::vector& n_args); +void register_bwd_graph_ops_z3(long graph_id, + const std::vector& op_names, + const std::vector& n_args); +void register_z3_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + bool persistent); +at::Tensor allgather_param(at::Tensor param_tensor, long graph_id, long ds_id); +void set_persistent(long ds_id); +void prefetch_params_fused(long graph_id, + const std::vector params, + const std::vector& ds_ids); +void prefetch_params_fused_meta(long graph_id, + const std::vector params, + const std::vector& ds_ids); +// for profiling +void invalidate_gathered_param(long ds_id); +void clear_all_gathered_params(); +at::Tensor allgather_param_meta(at::Tensor param_tensor, long graph_id, long ds_id); +at::Tensor release_param(at::Tensor dummy, long graph_id, long ds_id, long n_users); +at::Tensor release_param_meta(at::Tensor dummy, long graph_id, long ds_id, long n_users); +at::Tensor wait_allgather(at::Tensor v, long graph_id, const long ds_id); +at::Tensor wait_allgather_meta(at::Tensor v, long graph_id, long ds_id); +at::Tensor offload_tensor(at::Tensor tensor, long graph_id, long id); +at::Tensor reload_tensor(at::Tensor tensor, long graph_id, long id); +at::Tensor wait_offload(at::Tensor tensor, long graph_id, long id); +at::Tensor wait_reload(at::Tensor tensor, long graph_id, long id); +void reload_parameter(at::Tensor tensor, long graph_id, long id); +void offload_parameter(at::Tensor tensor, long graph_id, long id); +void reload_parameter_meta(at::Tensor tensor, long graph_id, long id); +void offload_parameter_meta(at::Tensor tensor, long graph_id, long id); +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d25578f410da278dace61fb5b488b14ed1257c9f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_adam.h" + +// C++ interface + +void multi_tensor_adam(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, /*gpmv*/ + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay) +{ + static bool initialized = false; + if (!initialized) { + create_adam_optimizer(0); + initialized = true; + } + for (int i = 0; i < tensor_lists[0].size(); i++) { + ds_adam_step(0, + step, + lr, + beta1, + beta2, + epsilon, + weight_decay, + bias_correction, + tensor_lists[1][i], + tensor_lists[0][i], + tensor_lists[2][i], + tensor_lists[3][i]); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_adam", + &multi_tensor_adam, + "Compute and apply gradient update to parameters for Adam optimizer"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d28509e592668aca100555e18ae47e3aadacb43a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp @@ -0,0 +1,344 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include +#include "shm.h" + +// #define DO_PROFILE +#ifdef DO_PROFILE +#include +#include +#endif + +// Communication settings +static int world_rank = -1; +static int world_size = -1; + +static std::set _comm_ids; +static std::set _colors; +static std::vector _ccl_comms; +static ccl::shared_ptr_class sub_kvs; +static std::map, int> group_to_comm_id; + +ccl::communicator& _get_comm_from_group() { return _ccl_comms[0]; } +ccl::communicator& _get_comm_from_group(py::object group) { return _ccl_comms[0]; } +ccl::communicator& _get_comm_from_group(std::vector ranks) +{ + if (group_to_comm_id.find(ranks) != group_to_comm_id.end()) { + auto id = group_to_comm_id.find(ranks); + return _ccl_comms[id->second]; + } + return _ccl_comms[0]; +} + +#define CCLCHECK(cmd) \ + do { \ + cmd; \ + } while (0) + +#define KVS_CREATE_SUCCESS 0 +#define KVS_CREATE_FAILURE -1 + +static bool is_initialized = 0; + +static ccl::shared_ptr_class kvs; + +static bool all_ranks_local_p = false; + +void initialize(int size, int rank, torch::Tensor& kvs_data) +{ + if (is_initialized) return; + + // Check whether all ranks is on the same physical machine. + // If true, we will use an SHM based low latency allreduce + + auto ls_string = std::getenv("LOCAL_SIZE"); + int ls = 0; + if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } + + if (size >= 1 && size == ls) { all_ranks_local_p = true; } + + world_size = size; + world_rank = rank; + is_initialized = 1; + + ccl::kvs::address_type main_addr; + + if (rank != 0) { + memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); + kvs = ccl::create_kvs(main_addr); + } + + _ccl_comms.emplace_back(ccl::create_communicator(size, rank, kvs)); + + auto addr_string = std::getenv("MASTER_ADDR"); + if (addr_string == NULL) { addr_string = ""; } + auto port_string = std::getenv("MASTER_PORT"); + if (port_string == NULL) { port_string = ""; } + + if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } +} + +/* + rank == 0: create main kvs and return its address + rank == else: return an empty address +*/ +std::vector get_kvs_addr(int rank) +{ + if (rank == 0) { + kvs = ccl::create_main_kvs(); + ccl::kvs::address_type main_addr = kvs->get_address(); + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } else { + ccl::kvs::address_type main_addr; + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } +} + +int get_rank(int group = 0) { return world_rank; } + +int get_world_size(int group = 0) { return world_size; } + +// Find the next ordered, unique value to a set. E.g. <0,1,2,7> --> 3 +int next_unique_val(std::set s) +{ + std::set::iterator itr; + // Base case. Add 0 to start of set. + if (s.empty() || *s.begin() != 0) { + return 0; + // second base case where s = {0} (the case of s = {n != 0} is caught above) + } else if (s.size() == 1) { + return 1; + } else { + int prev_val = *s.begin(); + for (itr = std::next(s.begin()); itr != s.end(); itr++) { + if (*itr != prev_val + 1) { return prev_val + 1; } + prev_val = *itr; + } + return *(s.end()) + 1; + } +} + +std::vector get_sub_kvs_addr(bool first) +{ + if (first) { + sub_kvs = ccl::create_main_kvs(); + ccl::kvs::address_type main_addr = sub_kvs->get_address(); + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } else { + ccl::kvs::address_type main_addr; + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } +} + +void initialize_sub_comm(int size, int rank, torch::Tensor& kvs_data, std::vector ranks) +{ + ccl::kvs::address_type main_addr; + if (rank != 0) { + memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); + sub_kvs = ccl::create_kvs(main_addr); + } + _ccl_comms.push_back(ccl::create_communicator(size, rank, sub_kvs)); + group_to_comm_id[ranks] = _ccl_comms.size() - 1; +} + +ccl::datatype get_ccl_datatype(c10::ScalarType type) +{ + ccl::datatype ccl_type; + switch (type) { + case c10::ScalarType::Int: ccl_type = ccl::datatype::int32; break; + case c10::ScalarType::Long: ccl_type = ccl::datatype::int64; break; + case c10::ScalarType::Float: ccl_type = ccl::datatype::float32; break; + case c10::ScalarType::Double: ccl_type = ccl::datatype::float64; break; + case c10::ScalarType::BFloat16: ccl_type = ccl::datatype::bfloat16; break; + case c10::ScalarType::Half: ccl_type = ccl::datatype::float16; break; + default: ccl_type = ccl::datatype::int8; + } + return ccl_type; +} + +ccl::reduction get_ccl_reduce_op(py::object op, at::Tensor& input) +{ + py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); + if (!py::isinstance(op, ReduceOp)) { + throw std::runtime_error("Error: Op must be of type ReduceOp"); + } + + int op_val = py::int_(op.attr("value")); + ccl::reduction ccl_op; + + if (input.scalar_type() == at::kBool) { + if (op_val == (int)py::int_(ReduceOp.attr("SUM").attr("value"))) { + // For bool tensors, map sum to max, which both represent a bitwise or. + // This is to prevent overflow issues with sum, since we use uint8 to + // represent a bool (see cclDataType mapping). + ccl_op = ccl::reduction::max; + } else if (op_val == (int)py::int_(ReduceOp.attr("AVG").attr("value"))) { + throw std::runtime_error("Error: For bool tensors, op must be of type ReduceOp"); + } + } + + if (op_val == (int)py::int_(ReduceOp.attr("SUM").attr("value"))) { + ccl_op = ccl::reduction::sum; + } else if (op_val == (int)py::int_(ReduceOp.attr("MIN").attr("value"))) { + ccl_op = ccl::reduction::min; + } else if (op_val == (int)py::int_(ReduceOp.attr("MAX").attr("value"))) { + ccl_op = ccl::reduction::max; + } else if (op_val == (int)py::int_(ReduceOp.attr("PRODUCT").attr("value"))) { + ccl_op = ccl::reduction::prod; + } else { + throw std::runtime_error("Error: Unrecognized ReduceOp type"); + } + return ccl_op; +} + +void broadcast(torch::Tensor& data, int src, std::vector group, bool async_op) +{ + CCLCHECK(ccl::broadcast(data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + src, + _get_comm_from_group(group)) + .wait()); +} + +// TODO: implement torch's async_op behavior, document it. +void all_reduce(torch::Tensor& data, py::object op, std::vector group, bool async_op) +{ + CCLCHECK(ccl::allreduce(data.data_ptr(), + data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + get_ccl_reduce_op(op, data), + _get_comm_from_group(group)) + .wait()); +} + +void all_reduce_caching(torch::Tensor& data, + py::object op, + std::string match_id, + std::vector group, + bool async_op) +{ + ccl::allreduce_attr attr = ccl::default_allreduce_attr; + auto match_str = ccl::v1::string(match_id); + attr.template set(true); + attr.template set(match_str); + // To control this, use operation attribute and set true value for to_cache field and unique + // string (for example, tensor name) for match_id field. Note that: + // match_id should be the same for a specific communication operation across all ranks. + // If the same tensor is a part of different communication operations, match_id should have + // different values for each of these operations. + CCLCHECK(ccl::allreduce(data.data_ptr(), + data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + get_ccl_reduce_op(op, data), + _get_comm_from_group(group), + attr) + .wait()); +} + +void inference_all_reduce(torch::Tensor& data, py::object op) +{ +#ifdef DO_PROFILE + static double total_time = 0.0; + static double total_time_sq = 0.0; + static int count = -16; // warmup + static double max_time = 0.0; + static double min_time = DBL_MAX; + // make sure all rank reach this point before measuring time + // turn on this if you suspect each rank didn't reach here at the same time (stragger) + // if (all_ranks_local_p) { + // barrier_wait(0, world_size); + //} + auto start = std::chrono::system_clock::now(); +#endif + + static py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); + static auto ReduceOpSum = (int)py::int_(ReduceOp.attr("SUM").attr("value")); + + assert(py::int_(op.attr("value")) == ReduceOpSum); + + auto numel = data.numel(); + + int data_size = 0; + bool data_type_fallback = false; + + switch (data.scalar_type()) { + case c10::ScalarType::BFloat16: data_size = numel * 2; break; + case c10::ScalarType::Float: data_size = numel * 4; break; + default: data_type_fallback = true; + } + + if (data_type_fallback || !all_ranks_local_p) { + // fallback to oneccl allreduce + CCLCHECK(ccl::allreduce(data.data_ptr(), + data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + get_ccl_reduce_op(op, data), + _get_comm_from_group()) + .wait()); + } else { + all_reduce_outer_loop(data, numel, data_size); + } + +#ifdef DO_PROFILE + auto end = std::chrono::system_clock::now(); + count++; + if (count > 0) { + double elapsed = std::chrono::duration_cast(end - start).count(); + if (elapsed > max_time) { max_time = elapsed; } + if (elapsed < min_time) { min_time = elapsed; } + total_time += elapsed; + total_time_sq += elapsed * elapsed; + if (world_rank == 0 && count == 1000) { + auto avg = total_time / count; + auto sd = + sqrt(total_time_sq / count - total_time * total_time / (count * count)) / avg * 100; + printf(" C++ kernel\t\t %.2f\t %.2f\t%.2f\t %.2f\n", + min_time, + max_time, + total_time / count, + sd); + } + } +#endif +} + +void barrier(std::vector group, bool async_op) +{ + CCLCHECK(ccl::barrier(_get_comm_from_group(group)).wait()); +} + +std::vector get_available_coll() +{ + std::vector colls{ + "broadcast", "all_reduce", "inference_all_reduce", "all_reduce_caching", "barrier"}; + return colls; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("get_kvs_addr", &get_kvs_addr, "create and get main kvs addr"); + m.def("initialize", &initialize, "ccl initialize"); + m.def("get_rank", &get_rank, "get rank"); + m.def("get_world_size", &get_world_size, "get world size"); + m.def("broadcast", &broadcast, "ccl broadcast"); + m.def("all_reduce", &all_reduce, "ccl all_reduce"); + m.def("inference_all_reduce", &inference_all_reduce, "low latency all_reduce implementation"); + m.def("all_reduce_caching", &all_reduce_caching, "ccl all_reduce with caching"); + m.def("barrier", &barrier, "barrier"); + m.def("initialize_sub_comm", &initialize_sub_comm, "initialize_sub_comm"); + m.def("get_sub_kvs_addr", &get_sub_kvs_addr, "get_sub_kvs_addr"); + m.def("get_available_coll", &get_available_coll, "get_available_coll"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp new file mode 100644 index 0000000000000000000000000000000000000000..be44681ca0626a9dcef9ccf5e7ea09b38bf27d6c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp @@ -0,0 +1,692 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include +#include +#include +#include +#include +#include "shm.h" + +// #define DO_PROFILE +#ifdef DO_PROFILE +#include +#include +#endif + +// states for collectives +enum coll_state { + coll_begin = 0, + coll_allreduce_naive__copy_in_done, + coll_allreduce_naive__reduce_done, + // alternative state when allreduce is working on alternative buffer + // of the double buffer. + coll_alt1_allreduce_naive__copy_in_done, + coll_alt2_allreduce_naive__copy_in_done, + coll_alt1_allreduce_naive__reduce_done, +}; + +// SHM building blocks +struct SharedData { + const char* name; + int descriptor; + void* bytes; + size_t nbytes; +}; + +void shared_open(SharedData* data, const char* name, size_t nbytes) +{ + int d = shm_open(name, O_RDWR, S_IRUSR | S_IWUSR); + if (d != -1) { + void* bytes = mmap(NULL, nbytes, PROT_READ | PROT_WRITE, MAP_SHARED, d, 0); + data->name = name; + data->descriptor = d; + data->bytes = bytes; + data->nbytes = nbytes; + } else { + if (errno != ENOENT) { + // don't print if shm can not be found because we want to loop over from + // caller again until the other ranks created the shm + printf("shared_open %s failed, errno=%d\n", name, errno); + } + data->descriptor = -1; + } +} + +void shared_create(SharedData* data, const char* name, void* bytes, size_t nbytes) +{ + int d = shm_open(name, O_CREAT | O_RDWR, S_IRUSR | S_IWUSR); + if (d != -1) { + if (nbytes = write(d, bytes, nbytes)) { shared_open(data, name, nbytes); } + } else { + printf("shared_create %s failed\n", name); + } +} + +void shared_close(SharedData* data) +{ + if (data->descriptor != -1) { + munmap(data->bytes, data->nbytes); + shm_unlink(data->name); + } +} + +static int world_size; + +// SHM based allreduce helper functions +// buffer that holds shm name +#define NAME_BUF_SIZE 1000 +#define MAX_BUF_SIZE 1048576 * 32 +#define NAIVE_ALLREDUCE_THRESHOLD 1048576 +#define SHM_BUFFER_NAME "deepspeed_allreduce_buffer" +struct allreduce_workspace { + enum coll_state states[2]; // idx=0 -- state for symmetric_naive_all_reduce + // idx=1 -- state for distributed_naive_all_reduce + // double buffer to avoid syncing between rounds + // offset=0 -- 2*NAIVE_ALLREDUCE_THRESHOLD : buffer for symmetric_naive_all_reduce + // after that : buffer for distributed_naive_all_reduce + char buffer[2 * NAIVE_ALLREDUCE_THRESHOLD + 2 * MAX_BUF_SIZE]; +}; + +#define BUFFER0_OFFSET(current_buffer) current_buffer* NAIVE_ALLREDUCE_THRESHOLD +#define BUFFER1_OFFSET(current_buffer) 2 * NAIVE_ALLREDUCE_THRESHOLD + current_buffer* MAX_BUF_SIZE + +struct allreduce_workspace** workspace; + +// buffer for small messages, double buffer +char** symmetric_buffer[2]; +// buffer for large messages, double buffer +char** distributed_buffer[2]; + +void wait_buffer_state_until_2(int index, + enum coll_state state0, + enum coll_state state1, + int state_group) +{ + volatile enum coll_state* state_ptr = &(workspace[index]->states[state_group]); + + while (1) { + volatile enum coll_state cur_state = *state_ptr; + if (cur_state == state0 || cur_state == state1) break; + } +} + +__m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_bf16_to_fp32(const __m256i src) +{ + auto y = _mm512_cvtepu16_epi32(src); + return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2)); +} + +inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_bf16(const __m512 src) +{ + __m512i value = _mm512_castps_si512(src); + __m512i nan = _mm512_set1_epi32(0xffff); + auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q); + __m512i ones = _mm512_set1_epi32(0x1); + __m512i vec_bias = _mm512_set1_epi32(0x7fff); + // uint32_t lsb = (input >> 16) & 1; + auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones); + // uint32_t rounding_bias = 0x7fff + lsb; + t_value = _mm512_add_epi32(t_value, vec_bias); + // input += rounding_bias; + t_value = _mm512_add_epi32(t_value, value); + // input = input >> 16; + t_value = _mm512_srli_epi32(t_value, 16); + // Check NaN before converting back to bf16 + t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value); + return _mm512_cvtusepi32_epi16(t_value); +} + +__m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_fp16_to_fp32(const __m256i src) { return _mm512_cvtph_ps(src); } + +inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_fp16(const __m512 src) +{ + return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); +} + +void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) + __attribute__((target("avx512bw"))); + +void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) + __attribute__((target("avx512bw"))); + +void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) + __attribute__((target("avx512bw"))); + +void reduce_all_buffers(int start_elements, + int num_elements, + c10::ScalarType scalar_type, + int to_buffer_idx, + char* to_buffer, + char** buffers) +{ + switch (scalar_type) { + case c10::ScalarType::BFloat16: + reduce_bf16_buffers(start_elements, num_elements, to_buffer, buffers); + break; + case c10::ScalarType::Half: + reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers); + break; + case c10::ScalarType::Float: + reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers); + break; + default: assert(!"Should not get here"); + } +} + +#define CVT_ADD_BF16(x) \ + do { \ + auto in##x##_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ + inout_val = _mm512_add_ps(inout_val, in##x##_val); \ + } while (0) + +// Reduce functions down below use vectorized algorithm, the number of bytes processed each +// iteration depends on vector length. 256bit vector ==> 32 bytes, 512bit vector ==> 64 bytes +// If you change implementation of reduce_bf16_buffers, etc. , check whether this number needs +// to be changed +#define VECTOR_LENGTH_IN_BYTES 32 + +void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) +{ + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); + switch (world_size) { + case 16: CVT_ADD_BF16(15); + case 15: CVT_ADD_BF16(14); + case 14: CVT_ADD_BF16(13); + case 13: CVT_ADD_BF16(12); + case 12: CVT_ADD_BF16(11); + case 11: CVT_ADD_BF16(10); + case 10: CVT_ADD_BF16(9); + case 9: CVT_ADD_BF16(8); + case 8: CVT_ADD_BF16(7); + case 7: CVT_ADD_BF16(6); + case 6: CVT_ADD_BF16(5); + case 5: CVT_ADD_BF16(4); + case 4: CVT_ADD_BF16(3); + case 3: CVT_ADD_BF16(2); + case 2: CVT_ADD_BF16(1); + case 1: break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); + inout_val = _mm512_add_ps(inout_val, in_val); + } + } + _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_bf16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { val += *(at::BFloat16*)(buffers[j] + i); } + *(at::BFloat16*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +#define CVT_ADD_FP16(x) \ + do { \ + auto in##x##_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ + inout_val = _mm512_add_ps(inout_val, in##x##_val); \ + } while (0) + +void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) +{ + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); + switch (world_size) { + case 16: CVT_ADD_FP16(15); + case 15: CVT_ADD_FP16(14); + case 14: CVT_ADD_FP16(13); + case 13: CVT_ADD_FP16(12); + case 12: CVT_ADD_FP16(11); + case 11: CVT_ADD_FP16(10); + case 10: CVT_ADD_FP16(9); + case 9: CVT_ADD_FP16(8); + case 8: CVT_ADD_FP16(7); + case 7: CVT_ADD_FP16(6); + case 6: CVT_ADD_FP16(5); + case 5: CVT_ADD_FP16(4); + case 4: CVT_ADD_FP16(3); + case 3: CVT_ADD_FP16(2); + case 2: CVT_ADD_FP16(1); + case 1: break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); + inout_val = _mm512_add_ps(inout_val, in_val); + } + } + _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_fp16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { val += *(at::Half*)(buffers[j] + i); } + *(at::Half*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +#define CVT_ADD_F32(x) \ + do { \ + auto in##x##_val = _mm256_loadu_ps((float*)(buffers[x] + i)); \ + inout_val = _mm256_add_ps(inout_val, in##x##_val); \ + } while (0) + +void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) +{ + const int element_size = 4; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = _mm256_loadu_ps((float*)(buffers[0] + i)); + switch (world_size) { + case 16: CVT_ADD_F32(15); + case 15: CVT_ADD_F32(14); + case 14: CVT_ADD_F32(13); + case 13: CVT_ADD_F32(12); + case 12: CVT_ADD_F32(11); + case 11: CVT_ADD_F32(10); + case 10: CVT_ADD_F32(9); + case 9: CVT_ADD_F32(8); + case 8: CVT_ADD_F32(7); + case 7: CVT_ADD_F32(6); + case 6: CVT_ADD_F32(5); + case 5: CVT_ADD_F32(4); + case 4: CVT_ADD_F32(3); + case 3: CVT_ADD_F32(2); + case 2: CVT_ADD_F32(1); + case 1: break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = _mm256_loadu_ps((float*)(buffers[j] + i)); + inout_val = _mm256_add_ps(inout_val, in_val); + } + } + _mm256_storeu_ps((float*)(to_buffer + i), inout_val); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { val += *(float*)(buffers[j] + i); } + *(float*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +static bool is_initialized = 0; +static int world_rank; + +void shm_initialize(int size, int rank, char* addr_string, char* port_string) +{ + if (is_initialized) return; + is_initialized = 1; + + world_size = size; + world_rank = rank; + + char shm_name_prefix[NAME_BUF_SIZE]; + char shm_name[NAME_BUF_SIZE]; + snprintf(shm_name_prefix, + NAME_BUF_SIZE, + "%s_%d_%s_%s", + SHM_BUFFER_NAME, + getuid(), + addr_string, + port_string); + // create shared workspace for SHM based allreduce + SharedData allreduce_buffer; + // allocate workspace_buf for current rank + struct allreduce_workspace* workspace_buf; + struct allreduce_workspace* workspace_buf_other; + workspace_buf = (struct allreduce_workspace*)malloc(sizeof(struct allreduce_workspace)); + snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, rank); + shared_create(&allreduce_buffer, shm_name, workspace_buf, sizeof(struct allreduce_workspace)); + workspace_buf = (struct allreduce_workspace*)allreduce_buffer.bytes; + workspace_buf->states[0] = coll_alt2_allreduce_naive__copy_in_done; + workspace_buf->states[1] = coll_begin; + + // create the workspace pointer list + workspace = (struct allreduce_workspace**)malloc(size * sizeof(struct allreduce_workspace*)); + symmetric_buffer[0] = (char**)malloc(size * sizeof(char**)); + symmetric_buffer[1] = (char**)malloc(size * sizeof(char**)); + distributed_buffer[0] = (char**)malloc(size * sizeof(char**)); + distributed_buffer[1] = (char**)malloc(size * sizeof(char**)); + + // map shm of all ranks + for (int i = 0; i < size; i++) { + if (i != rank) { + snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, i); + // printf("open %s, %d\n", shm_name, rank); + do { + shared_open(&allreduce_buffer, shm_name, sizeof(struct allreduce_workspace)); + } while (allreduce_buffer.descriptor == -1 && errno == ENOENT); + workspace_buf_other = (struct allreduce_workspace*)allreduce_buffer.bytes; + workspace[i] = workspace_buf_other; + } else { + workspace[i] = workspace_buf; + } + symmetric_buffer[0][i] = workspace[i]->buffer + BUFFER0_OFFSET(0); + symmetric_buffer[1][i] = workspace[i]->buffer + BUFFER0_OFFSET(1); + distributed_buffer[0][i] = workspace[i]->buffer + BUFFER1_OFFSET(0); + distributed_buffer[1][i] = workspace[i]->buffer + BUFFER1_OFFSET(1); + } +} + +static void parallel_memcpy(void* to, void* from, size_t n_bytes) + __attribute__((target("avx512bw"))); +static void parallel_memcpy(void* to, void* from, size_t n_bytes) +{ + auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES); + // process aligned part +#pragma omp parallel for + for (int i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) { + auto val = _mm256_loadu_si256((__m256i*)((char*)from + i)); + _mm256_storeu_si256((__m256i*)((char*)to + i), val); + } + + // process remaining part + for (int i = aligned_bytes; i < n_bytes; i++) { *((char*)to + i) = *((char*)from + i); } +} + +#define positive_mod(num, mod) ((((num) % (mod)) + (mod)) % (mod)) +#define rank_mod(rank) positive_mod(rank, world_size) +size_t slice_size(size_t chunk_el, int slice_idx) +{ + size_t slice_size = chunk_el / world_size; + return slice_idx == world_size - 1 ? slice_size + (chunk_el % world_size) : slice_size; +} + +char* slice_data(char* data_ptr, size_t chunk_el, int el_size, int slice_idx) +{ + size_t slice_size = chunk_el / world_size; + size_t el_offset = slice_size * slice_idx; + return data_ptr + el_offset * el_size; +} + +size_t slice_el_start(size_t chunk_el, int slice_idx) +{ + size_t slice_size = chunk_el / world_size; + return slice_size * slice_idx; +} + +/* + Symmetrical naive all_reduce + step 0: before enter the function ith times, state is copy(i-1) + step 1: each rank copy data from input (data_ptr) to SHM buffer[i] + step 2: set own state to copy(i) + step 3: wait each other rank's state equal or later than copy(i) + step 4: reduce across SHM buffer(ith) directly into output (data_ptr) +*/ +void symmetric_naive_all_reduce(char* data_ptr, + c10::ScalarType scalar_type, + size_t chunk_size, + size_t chunk_el) +{ +#ifdef DO_PROFILE + static double total_t1_t0 = 0.0; + static double total_t2_t1 = 0.0; + static double total_t3_t2 = 0.0; + static int count = -16; // warmup + auto t0 = std::chrono::system_clock::now(); +#endif + + /* + We can't have infinite number of buffers and states. 2 sets of buffer + and 3 sets of states is just enough. Consider current rank is in step 3, + with it's own state set to copy(i), the other rank will them have the + following situations: + ------------------------------------------------ + my state | can I proceed? | the other rank state + ================================================ + | N | copy(i-1) + |----------------|--------------------- + copy(i) | Y | copy(i) + |----------------|--------------------- + | Y | copy(i+1) + ------------------------------------------------ + * When I have state as copy(i), the other rank cannot have state + copy(i-2) or before. In that case I'll be in state copy(i-1) and cannot + proceed to copy(i). + * The other rank cannot have state copy(i+2) or beyond because my + state is still copy(i), copy(i+1) is as far as the other rank could go. + * From a rank's POV, all the other ranks can be divided into three sets: + - Lagging ranks: ranks that are still working on previous iteration + - Syncing ranks: ranks that are working on current iteration + - Leading ranks: ranks that are working on next iteration + * We can have 3 sets of states, one set for syncing ranks; one set for + lagging ranks; one set of leading ranks. With 3 sets of states, we can + distinguish between lagging and leading ranks. + * Note from any rank's POV, leading ranks and lagging ranks does not + appear at the same time. Either all other ranks are syncing or + lagging, or all other ranks are syncing or leading. Otherwise leading + and lagging ranks will be 2 iterations apart and this should not happen. + * So we have 2 sets of buffers, one buffer is used by current iter; + one buffer used by either lagging ranks or leading ranks. + */ + const int state_group = 0; + static int current_buffer = 0; + static int state_idx = 0; + + enum coll_state copy_current, copy_next; + + switch (state_idx) { + case 0: + copy_current = coll_allreduce_naive__copy_in_done; + copy_next = coll_alt1_allreduce_naive__copy_in_done; + break; + case 1: + copy_current = coll_alt1_allreduce_naive__copy_in_done; + copy_next = coll_alt2_allreduce_naive__copy_in_done; + break; + case 2: + copy_current = coll_alt2_allreduce_naive__copy_in_done; + copy_next = coll_allreduce_naive__copy_in_done; + break; + default: assert(!"Should not get here."); + } + state_idx = (state_idx + 1) % 3; + + parallel_memcpy(symmetric_buffer[current_buffer][world_rank], data_ptr, chunk_size); + std::atomic_thread_fence(std::memory_order_release); + workspace[world_rank]->states[state_group] = copy_current; + +#ifdef DO_PROFILE + auto t1 = std::chrono::system_clock::now(); +#endif + + for (int i = 0; i < world_size; i++) { + // wait until the other rank copy the buffer + if (i != world_rank) { wait_buffer_state_until_2(i, copy_current, copy_next, state_group); } + } +#ifdef DO_PROFILE + auto t2 = std::chrono::system_clock::now(); +#endif + + // each rank reduce the buffer independently so therre is no need for synchronization afterward + reduce_all_buffers( + 0, chunk_el, scalar_type, world_rank, data_ptr, symmetric_buffer[current_buffer]); + + // switch buffer + current_buffer = 1 - current_buffer; + +#ifdef DO_PROFILE + auto t3 = std::chrono::system_clock::now(); + + count++; + if (count > 0) { + total_t1_t0 += std::chrono::duration_cast(t1 - t0).count(); + total_t2_t1 += std::chrono::duration_cast(t2 - t1).count(); + total_t3_t2 += std::chrono::duration_cast(t3 - t2).count(); + if (world_rank == 0 && count == 1000) { + printf("symmetric_naive_all_reduce time breakdown:\n"); + printf("\tcopy input buffer: %.2f\n", total_t1_t0 / count); + printf("\twait for copy: %.2f\n", total_t2_t1 / count); + printf("\treduce: %.2f\n", total_t3_t2 / count); + } + } +#endif +} + +// naive allreduce distributed, each rank do naive reduce on its slice +void distributed_naive_reduce(char* data_ptr, + c10::ScalarType scalar_type, + size_t chunk_size, + size_t chunk_el) +{ +#ifdef DO_PROFILE + static double total_t1_t0 = 0.0; + static double total_t2_t1 = 0.0; + static double total_t3_t2 = 0.0; + static double total_t4_t3 = 0.0; + static double total_t5_t4 = 0.0; + static int count = -16; // warmup + auto t0 = std::chrono::system_clock::now(); +#endif + + const int state_group = 1; + static int current_buffer = 0; + static int state_idx = 0; + + enum coll_state copy_current, copy_next, reduce_current; + + // similar to symmetric_naive_allreduce, but here we only need two sets of + // states, because distributed naive reduce has two barriers in the algorithm + switch (state_idx) { + case 0: + copy_current = coll_allreduce_naive__copy_in_done; + reduce_current = coll_allreduce_naive__reduce_done; + copy_next = coll_alt1_allreduce_naive__copy_in_done; + break; + case 1: + copy_current = coll_alt1_allreduce_naive__copy_in_done; + reduce_current = coll_alt1_allreduce_naive__reduce_done; + copy_next = coll_allreduce_naive__copy_in_done; + break; + default: assert(!"Should not get here."); + } + state_idx = (state_idx + 1) % 2; + + int data_size = chunk_size / chunk_el; + parallel_memcpy(distributed_buffer[current_buffer][world_rank], data_ptr, chunk_size); + std::atomic_thread_fence(std::memory_order_release); + workspace[world_rank]->states[state_group] = copy_current; + +#ifdef DO_PROFILE + auto t1 = std::chrono::system_clock::now(); +#endif + + for (int i = 0; i < world_size; i++) { + // wait until all the other ranks copy the buffer + if (i != world_rank) + wait_buffer_state_until_2(i, copy_current, reduce_current, state_group); + } + +#ifdef DO_PROFILE + auto t2 = std::chrono::system_clock::now(); +#endif + + // reduce scatter + reduce_all_buffers(slice_el_start(chunk_el, world_rank), + slice_size(chunk_el, world_rank), + scalar_type, + world_rank, + distributed_buffer[current_buffer][world_rank], + distributed_buffer[current_buffer]); + std::atomic_thread_fence(std::memory_order_release); + workspace[world_rank]->states[state_group] = reduce_current; + +#ifdef DO_PROFILE + auto t3 = std::chrono::system_clock::now(); +#endif + + for (int i = 0; i < world_size; i++) { + // wait until all the other ranks reduce the buffer + if (i != world_rank) wait_buffer_state_until_2(i, reduce_current, copy_next, state_group); + } + + auto t4 = std::chrono::system_clock::now(); + + for (int i = 0; i < world_size; i++) { + int rank = (i + world_rank) % world_size; + parallel_memcpy( + slice_data(data_ptr, chunk_el, data_size, rank), + slice_data( + distributed_buffer[current_buffer][rank], chunk_el, chunk_size / chunk_el, rank), + slice_size(chunk_el, rank) * data_size); + } + + current_buffer = 1 - current_buffer; + +#ifdef DO_PROFILE + auto t5 = std::chrono::system_clock::now(); + count++; + if (count > 0) { + total_t1_t0 += std::chrono::duration_cast(t1 - t0).count(); + total_t2_t1 += std::chrono::duration_cast(t2 - t1).count(); + total_t3_t2 += std::chrono::duration_cast(t3 - t2).count(); + total_t4_t3 += std::chrono::duration_cast(t4 - t3).count(); + total_t5_t4 += std::chrono::duration_cast(t5 - t4).count(); + if (world_rank == 0 && count == 1000) { + printf("distributed_naive_reduce time breakdown:\n"); + printf("\tcopy input buffer: %.2f\n", total_t1_t0 / count); + printf("\twait for copy: %.2f\n", total_t2_t1 / count); + printf("\treduce: %.2f\n", total_t3_t2 / count); + printf("\twait for reduce finish: %.2f\n", total_t4_t3 / count); + printf("\tcopy out: %.2f\n", total_t5_t4 / count); + } + } +#endif +} + +void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size) +{ + for (int offset = 0; offset < data_size; offset += MAX_BUF_SIZE) { + auto data_ptr = ((char*)(data.data_ptr()) + offset); + size_t chunk_size = data_size - offset > MAX_BUF_SIZE ? MAX_BUF_SIZE : data_size - offset; + size_t chunk_el = chunk_size / (data_size / numel); + if (chunk_size < NAIVE_ALLREDUCE_THRESHOLD) + symmetric_naive_all_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el); + else + distributed_naive_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el); + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h new file mode 100644 index 0000000000000000000000000000000000000000..7f73197a8caa5c4487c79c0c30ef926797bae97b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#ifndef __SHM_COLLECTIVES__ +#define __SHM_COLLECTIVES__ +#define VECTOR_LENGTH_IN_BYTES 32 +void shm_initialize(int size, int rank, char* addr_string, char* port_string); +void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size); +void barrier_wait(int root_idx, int num_ranks); +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp new file mode 100644 index 0000000000000000000000000000000000000000..5be5cb799a7bb47f624f5b4952f2947738864d62 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp @@ -0,0 +1,186 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include "shm.h" + +// #define DO_PROFILE +#ifdef DO_PROFILE +#include +#include +#endif + +// Communication settings +static int world_rank = -1; +static int world_size = -1; + +static bool is_initialized = 0; + +static bool all_ranks_local_p = false; + +void initialize(int size, int rank) +{ + if (is_initialized) return; + + // Check whether all ranks is on the same physical machine. + // If true, we will use an SHM based low latency allreduce + + auto ls_string = std::getenv("LOCAL_SIZE"); + int ls = 0; + if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } + + if (size >= 1 && size == ls) { all_ranks_local_p = true; } + + world_size = size; + world_rank = rank; + is_initialized = 1; + + auto addr_string = std::getenv("MASTER_ADDR"); + if (addr_string == NULL) { addr_string = ""; } + auto port_string = std::getenv("MASTER_PORT"); + if (port_string == NULL) { port_string = ""; } + + if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } +} + +void inference_all_reduce_(torch::Tensor& data, int op); + +// Success - return 0 +// Fail (cannot hornor the request and need to fall back) - return -1 +void inference_all_reduce_(torch::Tensor& data, int op) +{ + assert(op == 0); +#ifdef DO_PROFILE + static double total_time = 0.0; + static double total_time_sq = 0.0; + static int count = -16; // warmup + static double max_time = 0.0; + static double min_time = DBL_MAX; + // make sure all rank reach this point before measuring time + // turn on this if you suspect each rank didn't reach here at the same time (stragger) + // if (all_ranks_local_p) { barrier_wait(0, world_size); } + auto start = std::chrono::system_clock::now(); +#endif + + auto numel = data.numel(); + + int data_size = 0; + bool data_type_fallback = false; + + switch (data.scalar_type()) { + case c10::ScalarType::BFloat16: data_size = numel * 2; break; + case c10::ScalarType::Half: data_size = numel * 2; break; + case c10::ScalarType::Float: data_size = numel * 4; break; + default: data_type_fallback = true; + } + + if (data_type_fallback) return; + + all_reduce_outer_loop(data, numel, data_size); + +#ifdef DO_PROFILE + auto end = std::chrono::system_clock::now(); + count++; + if (count > 0) { + double elapsed = std::chrono::duration_cast(end - start).count(); + if (elapsed > max_time) { max_time = elapsed; } + if (elapsed < min_time) { min_time = elapsed; } + total_time += elapsed; + total_time_sq += elapsed * elapsed; + if (world_rank == 0 && count == 1000) { + auto avg = total_time / count; + auto sd = + sqrt(total_time_sq / count - total_time * total_time / (count * count)) / avg * 100; + printf(" C++ kernel\t\t %.2f\t %.2f\t%.2f\t %.2f\n", + min_time, + max_time, + total_time / count, + sd); + } + } +#endif + return; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("initialize", &initialize, "shm initialize"); } + +TORCH_LIBRARY(deepspeed, m) +{ + m.def("inference_all_reduce(Tensor self) -> Tensor"); + m.def("inference_all_reduce_(Tensor(a!) self) -> Tensor(a!)"); +} + +torch::Tensor inference_all_reduce_meta(const torch::Tensor& self_) +{ + torch::Tensor result_ = torch::empty_like(self_); + return result_; +} + +torch::Tensor& inference_all_reduce__meta(torch::Tensor& self_) { return self_; } + +torch::Tensor& inference_all_reduce__cpu(torch::Tensor& self_) +{ + TORCH_INTERNAL_ASSERT(self_.device().type() == torch::DeviceType::CPU); + torch::Tensor self_tensor = self_.contiguous(); + inference_all_reduce_(self_tensor, 0); + return self_; +} + +torch::Tensor inference_all_reduce_cpu(const torch::Tensor& self_) +{ + torch::Tensor result = self_.clone(); + inference_all_reduce__cpu(result); + return result; +} + +#include +// The boilerplate functionalization logic, that teaches functionalization +// how to map x_() calls into x() calls. +// Long term, we'd like to not require users to write this logic. +// HOWEVER, if you have a custom op that is mutable, +// You will still need to write an out-of-place version of that op! +at::Tensor& inference_all_reduce__functionalization_glue(at::Tensor& x) +{ + // We expect all tensor inputs to our op to be "functional tensors" + TORCH_INTERNAL_ASSERT(at::functionalization::impl::isFunctionalTensor(x)); + // First, sync and unwrap and functional tensors + at::functionalization::impl::sync(x); + auto x_ = at::functionalization::impl::from_functional_tensor(x); + // Grab the dispatcher entry corresponding to the out-of-place op, "x" + static auto op_handle = c10::Dispatcher::singleton() + // specify namespace::op_name, op_overload_name + .findSchemaOrThrow("deepspeed::inference_all_reduce", "") + // Specify the C++ schema of the out-of-place op. + .typed(); + // Next, redispatch to the out-of-place op, x() (user called x_, we call x) + at::Tensor tmp_output; + { + at::AutoDispatchSkipFunctionalize guard; + tmp_output = op_handle.call(x_); + } + // Finally, tell functionalization about this mutation. + at::functionalization::impl::replace_(x, tmp_output); + at::functionalization::impl::commit_update(x); + at::functionalization::impl::sync(x); + return x; +} + +TORCH_LIBRARY_IMPL(deepspeed, CPU, m) +{ + m.impl("inference_all_reduce", inference_all_reduce_cpu); + m.impl("inference_all_reduce_", inference_all_reduce__cpu); +} + +TORCH_LIBRARY_IMPL(deepspeed, Meta, m) +{ + m.impl("inference_all_reduce", inference_all_reduce_meta); + m.impl("inference_all_reduce_", inference_all_reduce__meta); +} + +TORCH_LIBRARY_IMPL(deepspeed, Functionalize, m) +{ + m.impl("inference_all_reduce_", inference_all_reduce__functionalization_glue); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp new file mode 100644 index 0000000000000000000000000000000000000000..708df7f0146aa996e1652ed938c038331738e149 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_lion.h" + +// C++ interface + +void multi_tensor_lion(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, /*gpmv*/ + const float lr, + const float beta1, + const float beta2, + const int step, + const int mode, + const float weight_decay) +{ + static bool initialized = false; + if (!initialized) { + create_lion_optimizer(0); + initialized = true; + } + for (int i = 0; i < tensor_lists[0].size(); i++) { + ds_lion_step(0, + step, + lr, + beta1, + beta2, + weight_decay, + tensor_lists[1][i], + tensor_lists[0][i], + tensor_lists[2][i]); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_lion", + &multi_tensor_lion, + "Compute and apply gradient update to parameters for Lion optimizer"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ac3364539ff1e47bceffcf79298f2bbc91845ac8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +void attention_impl(torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& o, + torch::Tensor& lse); +void attention(torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& o, + torch::Tensor& lse) +{ + attention_impl(q, k, v, bias1, bias2, o, lse); +} + +void attention_back_impl(torch::Tensor& go, + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& o, + torch::Tensor& lse, + torch::Tensor& delta, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& gq, + torch::Tensor& gk, + torch::Tensor& gv, + torch::Tensor& gb1, + torch::Tensor& gb2); +void attention_bwd(torch::Tensor& go, + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& o, + torch::Tensor& lse, + torch::Tensor& delta, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& gq, + torch::Tensor& gk, + torch::Tensor& gv, + torch::Tensor& gb1, + torch::Tensor& gb2) +{ + attention_back_impl(go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("attention", &attention, ""); + m.def("attention_bwd", &attention_bwd, ""); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu new file mode 100644 index 0000000000000000000000000000000000000000..a82c4ec68a13fccfe63c684a481afeebfe378031 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu @@ -0,0 +1,218 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include "gemm_kernel_utils.h" +#include "kernel_backward.h" +#include "transform/bias_broadcast.h" + +constexpr auto kBlockSizeI = 64; +constexpr auto kBlockSizeJ = 64; + +template + class Broadcast1_, + template + class Broadcast2_> +typename std::enable_if::value>::type attention_back_impl_template( + torch::Tensor& go, + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& o, + torch::Tensor& lse, + torch::Tensor& delta, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& gq, + torch::Tensor& gk, + torch::Tensor& gv, + torch::Tensor& gb1, + torch::Tensor& gb2) +{ + EVOFORMER_CHECK(false, "Unsupported GPU and data type combination") +} + +template + class Broadcast1_, + template + class Broadcast2_> +typename std::enable_if::value>::type attention_back_impl_template( + torch::Tensor& go, + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& o, + torch::Tensor& lse, + torch::Tensor& delta, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& gq, + torch::Tensor& gk, + torch::Tensor& gv, + torch::Tensor& gb1, + torch::Tensor& gb2) +{ + constexpr bool kPreload_ = arch::kMinComputeCapability >= 80; + using Kernel = AttentionBackwardKernel; + int head_size = q.size(-1); + int head_number = q.size(-2); + int seq_length = q.size(-3); + auto q_view = q.view({-1, seq_length, head_number, head_size}); + auto k_view = k.view({-1, seq_length, head_number, head_size}); + auto v_view = v.view({-1, seq_length, head_number, head_size}); + auto o_view = o.view({-1, seq_length, head_number, head_size}); + auto do_view = go.view({-1, seq_length, head_number, head_size}); + auto dk_view = gk.view({-1, seq_length, head_number, head_size}); + auto dv_view = gv.view({-1, seq_length, head_number, head_size}); + auto dq_view = gq.view({-1, seq_length, head_number, head_size}); + auto q_ptr = reinterpret_cast(q.data_ptr()); + auto k_ptr = reinterpret_cast(k.data_ptr()); + auto v_ptr = reinterpret_cast(v.data_ptr()); + auto o_ptr = reinterpret_cast(o.data_ptr()); + auto do_ptr = reinterpret_cast(go.data_ptr()); + auto dk_ptr = reinterpret_cast(gk.data_ptr()); + auto dv_ptr = reinterpret_cast(gv.data_ptr()); + auto dq_ptr = reinterpret_cast(gq.data_ptr()); + auto db1_ptr = gb1.size(0) > 0 ? reinterpret_cast(gb1.data_ptr()) : nullptr; + auto db2_ptr = gb2.size(0) > 0 ? reinterpret_cast(gb2.data_ptr()) : nullptr; + auto lse_ptr = reinterpret_cast(lse.data_ptr()); + auto delta_ptr = reinterpret_cast(delta.data_ptr()); + auto bias1_ptr = reinterpret_cast(bias1.data_ptr()); + auto bias2_ptr = reinterpret_cast(bias2.data_ptr()); + static_assert(Kernel::kKernelComputesDelta, "Kernel must compute delta"); + + typename Kernel::Params p; + p.query_ptr = q_ptr; + p.key_ptr = k_ptr; + p.value_ptr = v_ptr; + p.logsumexp_ptr = lse_ptr; + p.output_ptr = o_ptr; + p.grad_output_ptr = do_ptr; + p.delta_ptr = delta_ptr; + p.grad_query_ptr = dq_ptr; + p.grad_key_ptr = dk_ptr; + p.grad_value_ptr = dv_ptr; + + p.grad_bias1_ptr = db1_ptr; + p.grad_bias2_ptr = db2_ptr; + p.B = q.size(0); + p.N = q.size(1); + p.bias1_ptr = bias1.size(0) ? bias1_ptr : nullptr; + p.bias2_ptr = bias2.size(0) ? bias2_ptr : nullptr; + + p.scale = 1.0f / sqrtf(head_size); + + p.head_dim = head_size; + p.head_dim_value = head_size; + p.num_queries = seq_length; + p.num_keys = seq_length; + p.num_heads = head_number; + + p.q_strideM = q_view.stride(-3); + p.k_strideM = k_view.stride(-3); + p.v_strideM = v_view.stride(-3); + p.gO_strideM = do_view.stride(-3); + p.o_strideH = o_view.stride(-2); + p.q_strideH = q_view.stride(-2); + p.k_strideH = k_view.stride(-2); + p.v_strideH = v_view.stride(-2); + p.o_strideB = o_view.stride(-4); + p.q_strideB = q_view.stride(-4); + p.k_strideB = k_view.stride(-4); + p.v_strideB = v_view.stride(-4); + p.lse_strideB = lse.stride(-3); + p.lse_strideH = lse.stride(-2); + p.delta_strideB = delta.stride(-3); + p.delta_strideH = delta.stride(-2); + p.num_batches = q_view.size(-4); + + p.gO_strideB = do_view.stride(-4); + p.gQ_strideB = dq_view.stride(-4); + p.gK_strideB = dk_view.stride(-4); + p.gV_strideB = dv_view.stride(-4); + p.gO_strideH = do_view.stride(-2); + p.gQ_strideH = dq_view.stride(-2); + p.gK_strideH = dk_view.stride(-2); + p.gV_strideH = dv_view.stride(-2); + + torch::Tensor workspace = torch::empty(p.workspace_size() / 4, lse.options()); + p.workspace = workspace.data_ptr(); + + auto kernel_fn = attention_kernel_backward_batched_impl; + size_t smem_bytes = sizeof(typename Kernel::SharedStorage); + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, int(smem_bytes)); + if (!Kernel::check_supported(p)) { throw std::runtime_error("Unsupported parameters"); } + kernel_fn<<>>(p); +} + +#define CODE(scalar_t, torch_scalar_t) \ + do { \ + if (bias1.size(0) == 0 && bias2.size(0) == 0) { \ + attention_back_impl_template( \ + go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ + } else if (bias1.size(0) > 0 && bias2.size(0) > 0) { \ + attention_back_impl_template( \ + go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ + } else if (bias1.size(0) > 0) { \ + attention_back_impl_template( \ + go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ + } else { \ + attention_back_impl_template( \ + go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ + } \ + } while (0) + +void attention_back_impl(torch::Tensor& go, + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& o, + torch::Tensor& lse, + torch::Tensor& delta, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& gq, + torch::Tensor& gk, + torch::Tensor& gv, + torch::Tensor& gb1, + torch::Tensor& gb2) +{ + cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties(); + DISPATCH_ARCHTAG(prop->major * 10 + prop->minor, + DISPATCH_TYPES(q, { CODE(scalar_t, torch_scalar_t); })); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu new file mode 100644 index 0000000000000000000000000000000000000000..37636c4bf98842543fe1bd055a848cb52d0a6449 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu @@ -0,0 +1,160 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include "gemm_kernel_utils.h" +#include "kernel_forward.h" +#include "transform/bias_broadcast.h" + +template + class Broadcast1_, + template + class Broadcast2_> +typename std::enable_if::value>::type attention_impl_template( + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& o, + float* lse_ptr) +{ + EVOFORMER_CHECK(false, "Unsupported GPU and data type combination") +} + +template + class Broadcast1_, + template + class Broadcast2_> +typename std::enable_if::value>::type attention_impl_template( + torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& o, + float* lse_ptr) +{ + // Attention definition goes here, replaced with BroadcastType1 and + // BroadcastType2 + using Attention = AttentionKernel; + + static_assert(!Attention::kNeedsOutputAccumulatorBuffer, + "This test does not support output accumulator buffer"); + int head_size = q.size(-1); + int head_number = q.size(-2); + int seq_length = q.size(-3); + auto q_view = q.view({-1, seq_length, head_number, head_size}); + auto k_view = k.view({-1, seq_length, head_number, head_size}); + auto v_view = v.view({-1, seq_length, head_number, head_size}); + auto o_view = o.view({-1, seq_length, head_number, head_size}); + int batch_size = q_view.size(0); + auto q_ptr = reinterpret_cast(q.data_ptr()); + auto k_ptr = reinterpret_cast(k.data_ptr()); + auto v_ptr = reinterpret_cast(v.data_ptr()); + auto o_ptr = reinterpret_cast(o.data_ptr()); + + auto bias1_ptr = reinterpret_cast(bias1.data_ptr()); + auto bias2_ptr = reinterpret_cast(bias2.data_ptr()); + + typename Attention::Params p; + { // set parameters + p.query_ptr = q_ptr; + p.key_ptr = k_ptr; + p.value_ptr = v_ptr; + p.logsumexp_ptr = lse_ptr; // Only needed for bw + p.output_accum_ptr = nullptr; + p.output_ptr = o_ptr; + p.scale = 1.0f / sqrt(float(head_size)); + + p.bias1_ptr = bias1_ptr; + p.bias2_ptr = bias2_ptr; + p.B = q.size(0); + p.N = q.size(1); + + p.num_heads = head_number; + p.num_batches = batch_size; + p.head_dim = head_size; + p.head_dim_value = head_size; + p.num_queries = seq_length; + p.num_keys = seq_length; + + // All tensors are in BMHK shapes + p.q_strideH = q_view.stride(-2); + p.k_strideH = k_view.stride(-2); + p.v_strideH = v_view.stride(-2); + p.q_strideM = q_view.stride(-3); + p.k_strideM = k_view.stride(-3); + p.v_strideM = v_view.stride(-3); + p.o_strideM = o_view.stride(-3); + p.q_strideB = q_view.stride(-4); + p.k_strideB = k_view.stride(-4); + p.v_strideB = v_view.stride(-4); + } + + constexpr auto kernel_fn = attention_kernel_batched_impl; + int smem_bytes = sizeof(typename Attention::SharedStorage); + if (smem_bytes > 0xc000) { + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + } + if (!Attention::check_supported(p)) { throw std::runtime_error("Parameters not supported"); } + kernel_fn<<>>(p); +} + +#define CODE(scalar_t, torch_scalar_t) \ + do { \ + if (bias1.size(0) == 0 && bias2.size(0) == 0) { \ + attention_impl_template(q, k, v, bias1, bias2, o, lse_ptr); \ + } else if (bias1.size(0) == 0) { \ + attention_impl_template(q, k, v, bias1, bias2, o, lse_ptr); \ + } else if (bias2.size(0) == 0) { \ + attention_impl_template(q, k, v, bias1, bias2, o, lse_ptr); \ + } else { \ + attention_impl_template( \ + q, k, v, bias1, bias2, o, lse_ptr); \ + } \ + } while (0) + +// Function to select and call the correct template based on biases sizes +void attention_impl(torch::Tensor& q, + torch::Tensor& k, + torch::Tensor& v, + torch::Tensor& bias1, + torch::Tensor& bias2, + torch::Tensor& o, + torch::Tensor& lse) +{ + auto lse_ptr = lse.size(0) == 0 ? nullptr : reinterpret_cast(lse.data_ptr()); + cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties(); + DISPATCH_ARCHTAG(prop->major * 10 + prop->minor, + DISPATCH_TYPES(q, { CODE(scalar_t, torch_scalar_t); })); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h new file mode 100644 index 0000000000000000000000000000000000000000..17b6479ed8c53711a4b0c5fd7e23bee44ff9e142 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h @@ -0,0 +1,250 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once +#include +#include +#include "../iterators/predicated_tile_iterator_atomic.h" +#include "cutlass/epilogue/threadblock/epilogue.h" + +namespace cutlass { +namespace epilogue { +namespace threadblock { +template +struct EpilogueTensorOpAffineRankN : public DefaultEpilogueTensorOpAffineRankN { + using Base = DefaultEpilogueTensorOpAffineRankN; + using OutputTileIterator = + cutlass::epilogue::threadblock::PredicatedTileIteratorAffineRankNAtomic< + typename Base::OutputTileThreadMap, + typename Base::ElementOutput, + Rank>; + + using Epilogue = + cutlass::epilogue::threadblock::Epilogue; +}; + +template +struct EpilogueVoltaTensorOpAffineRankN + : public DefaultEpilogueVoltaTensorOpAffineRankN { + using Base = DefaultEpilogueVoltaTensorOpAffineRankN; + using OutputTileIterator = + cutlass::epilogue::threadblock::PredicatedTileIteratorAffineRankNAtomic< + typename Base::OutputTileThreadMap, + typename Base::ElementOutput, + Rank>; + + using Epilogue = + cutlass::epilogue::threadblock::Epilogue; +}; + +template +struct EpilogueTensorOp : public DefaultEpilogueTensorOp { + using Base = DefaultEpilogueTensorOp; + using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIteratorAtomic< + typename Base::OutputTileThreadMap, + typename Base::ElementOutput, + ScatterD, + PermuteDLayout>; + using Epilogue = + cutlass::epilogue::threadblock::Epilogue; +}; + +template +struct EpilogueVoltaTensorOp : public DefaultEpilogueVoltaTensorOp { + using Base = DefaultEpilogueVoltaTensorOp; + using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIteratorAtomic< + typename Base::OutputTileThreadMap, + typename Base::ElementOutput, + ScatterD, + PermuteDLayout>; + using Epilogue = + cutlass::epilogue::threadblock::Epilogue; +}; +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +template +struct BiasGradEpilogue { + using Epilogue = + typename cutlass::epilogue::threadblock::EpilogueTensorOp::Epilogue; +}; + +template +struct BiasGradEpilogue { + using Epilogue = + typename cutlass::epilogue::threadblock::EpilogueVoltaTensorOp::Epilogue; +}; + +template +struct BiasGradEpilogueAffineRankN { + using Epilogue = typename cutlass::epilogue::threadblock::EpilogueTensorOpAffineRankN< + Rank, + Shape_, + WarpMmaTensorOp_, + PartitionsK, + OutputOp_, + ElementsPerAccess>::Epilogue; +}; + +template +struct BiasGradEpilogueAffineRankN { + using Epilogue = typename cutlass::epilogue::threadblock::EpilogueVoltaTensorOpAffineRankN< + Rank, + Shape_, + WarpMmaTensorOp_, + PartitionsK, + OutputOp_, + ElementsPerAccess>::Epilogue; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h new file mode 100644 index 0000000000000000000000000000000000000000..3b7b32d6145202ea59d707b025669cbe75a8bc51 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h @@ -0,0 +1,592 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + File copied from "cutlass/epilogue/threadblock/epilogue.h" + then modified to: + (1) load 2 source fragments at the same time (pipelining) + (2) support reading from a different dtype + (3) pass the row id to the OutputOp if it takes it + (see MemoryEfficientAttentionNormalize) + Note that in general the fragment passed to the OutputOp could + span multiple rows but it does not happen with the configurations we have +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include "cutlass/aligned_buffer.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/functional.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/layout/vector.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_coord.h" + +#include "cutlass/gemm/gemm.h" + +#include "cutlass/transform/pitch_linear_thread_map.h" +#include "cutlass/transform/threadblock/regular_tile_iterator.h" + +#include "cutlass/epilogue/threadblock/epilogue_base.h" +#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "cutlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +template +struct ApplyEpilogueOp { + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum, + typename Op::FragmentOutput const& source) + { + return output_op(accum, source); + } + static CUTLASS_DEVICE typename Op::FragmentOutput + apply(Op const& output_op, int row_id, typename Op::FragmentAccumulator const& accum) + { + return output_op(accum); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template ::value), + typename OutputTileSourceIterator_ = + OutputTileIterator_ ///< Tile iterator reading tensors + > +class EpiloguePipelined : public EpilogueBase { +public: + using Base = EpilogueBase; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using OutputTileSourceIterator = OutputTileSourceIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + using ElementSource = typename OutputTileSourceIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename cutlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = + Array; + using SourceAccessType = Array; + + /// Array type used by output functor + using AccumulatorAccessType = + Array; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + static int constexpr kSmemTiles = + Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK; + static int constexpr kSmemPointerOffset = + Base::SharedStorage::StorageShape::kCount / kSmemTiles; + +public: + static_assert(OutputTileSourceIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between input tile and output tile iterator (kElements)"); + static_assert(OutputTileSourceIterator::kIterations == OutputTileIterator::kIterations, + "Mismatch between input tile and output tile iterator (kIterations)"); + static_assert(SharedLoadIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert(OutputTileIterator::kElementsPerAccess, + "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % + OutputTileIterator::kElementsPerAccess), + "Divisibility"); + +private: + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + +public: + /// Constructor + CUTLASS_DEVICE + EpiloguePipelined(typename Base::SharedStorage& shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.reference(), thread_idx) + { + } + + /// Streams the result to global memory + CUTLASS_DEVICE + void operator()(OutputOp const& output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& accumulators, ///< Complete warp-level accumulator tile + OutputTileSourceIterator source_iterator) + { ///< Threadblock tile coordinate in GEMM (in units + ///< of threadblock tiles) + + if (!output_op.is_source_needed()) { + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } else { + compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator); + } + } + CUTLASS_DEVICE + void operator()(OutputOp const& output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& accumulators) + { ///< Complete warp-level accumulator tile + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } + +private: + template + struct acc2smem_source_not_needed; + + template + struct acc2smem_source_not_needed> { + template + CUTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator& warp_tile_iterator) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { ++accum_fragment_iterator; } + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + warp_tile_iterator.store(accum_fragment); + if (p < Base::kFragmentsPerIteration - 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); + } + } + + if (Base::kFragmentsPerIteration > 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset * + (1 - Base::kFragmentsPerIteration)); + } + } + + CUTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const& iterator_begin, + WarpTileIterator& warp_tile_iterator) + { + int dummy[] = { + (pos == (Seq * Base::kFragmentsPerIteration)) && + (helper(iterator_begin, warp_tile_iterator), + 0)...}; + + CUTLASS_UNUSED(dummy[0]); + } + }; + + static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1, + "One of these must be exactly 1."); + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const& output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& accumulators ///< Complete warp-level accumulator tile + ) + { + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + +#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration \ + : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; + iter += Base::kFragmentsPerIteration) { + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem_source_not_needed>:: + push(iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + if (p < Base::kFragmentsPerIteration - 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + } else if (kPartitionsK > 1) { + plus add_fragments; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = + add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * + kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_source_not_needed_(destination_iterator.thread_start_row(), + output_fragment, + output_op, + aligned_accum_fragment[0]); + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + + if (Base::kFragmentsPerIteration > 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset * + (1 - Base::kFragmentsPerIteration)); + } + } + } + + template + struct acc2smem_source_needed; + + template + struct acc2smem_source_needed> { + template + CUTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator& warp_tile_iterator) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { ++accum_fragment_iterator; } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + CUTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const& iterator_begin, + WarpTileIterator& warp_tile_iterator) + { + int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_needed_( + OutputOp const& output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& accumulators, ///< Complete warp-level accumulator tile + OutputTileSourceIterator source_iterator ///< Threadblock tile coordinate in GEMM (in units + ///< of threadblock tiles) + ) + { + typename OutputTileSourceIterator::Fragment source_fragment[2]; + + source_fragment[0].clear(); + source_iterator.load(source_fragment[0]); + ++source_iterator; + source_fragment[1].clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + +#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + if (iter > 0) { __syncthreads(); } + // + // Load the source for next iteration (pipelining) + // + + if (iter + 1 < OutputTileIterator::kIterations) { + source_iterator.load(source_fragment[(iter + 1) % 2]); + } + ++source_iterator; + acc2smem_source_needed>:: + push(iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the + // k-slices + if (kPartitionsK > 1) { + plus add_fragments; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = + add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_(destination_iterator.thread_start_row(), + output_fragment, + output_op, + aligned_accum_fragment[0], + source_fragment[iter % 2]); + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_(int begin_row, + typename OutputTileIterator::Fragment& output_fragment, + OutputOp const& output_op, ///< Output operator + typename SharedLoadIterator::Fragment const& aligned_accum_fragment, + typename OutputTileSourceIterator::Fragment const& source_fragment) + { + OutputAccessType* output_frag_ptr = reinterpret_cast(&output_fragment); + + AccumulatorAccessType const* compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + SourceAccessType const* source_frag_ptr = + reinterpret_cast(&source_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = ApplyEpilogueOp::apply( + output_op, + begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), + compute_frag_ptr[i], + source_frag_ptr[i]); + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_source_not_needed_( + int begin_row, + typename OutputTileIterator::Fragment& output_fragment, + OutputOp const& output_op, ///< Output operator + typename SharedLoadIterator::Fragment const& aligned_accum_fragment) + { + OutputAccessType* output_frag_ptr = reinterpret_cast(&output_fragment); + + AccumulatorAccessType const* compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = ApplyEpilogueOp::apply( + output_op, + begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), + compute_frag_ptr[i]); + } + } + + // This should be constexpr, but it's only supported on c++14 + static int CUTLASS_HOST_DEVICE getRowOffset(int i) + { + using ThreadMap = typename OutputTileIterator::ThreadMap; + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + int frag_idx = ThreadMap::kElementsPerAccess * + (frag_row_idx * ThreadMap::Iterations::kColumn + column); + if (i < frag_idx + ThreadMap::kElementsPerAccess) { return row_offset; } + } + } + } + } + return -1; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h new file mode 100644 index 0000000000000000000000000000000000000000..f81a09f74f1ebed8e3d46fee52ac220efd98214b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h @@ -0,0 +1,251 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory + to match canonical tensor layouts in global memory. Epilogues support + conversion and reduction operations. + + This is a copy of cutlass/epilogue/threadblock/epilogue.h that can + handle "row_id" as a first argument, as uses it to get the corresponding + `m_prime` / `s_prime` to rescale the output. +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include "cutlass/aligned_buffer.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/functional.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/layout/vector.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_coord.h" + +#include "cutlass/gemm/gemm.h" + +#include "cutlass/transform/pitch_linear_thread_map.h" +#include "cutlass/transform/threadblock/regular_tile_iterator.h" + +#include "cutlass/epilogue/threadblock/epilogue_base.h" +#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "cutlass/numeric_types.h" + +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/functional.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "epilogue_pipelined.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +// output <- alpha * accumulator + beta * source +// with: +// alpha = 1 / s_prime (to normalize when isLast=True, 1 otherwise) +// beta = alpha / m_prime (renormalize the output when the max changes) +// source is the current output +template , + ///< but we use 64 or 32 sometimes when there are not enough data + ///< to store + typename ElementAccumulator_, ///< Accumulator data type + typename ElementCompute_, ///< Data type used to compute linear combination + bool isFirst, + bool isLast, + typename FragmentAlphaBeta_, + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest> +class MemoryEfficientAttentionNormalize { +public: + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentSource = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + using FragmentAlphaBeta = FragmentAlphaBeta_; + + static FloatRoundStyle const kRound = Round; + +private: + // + // Data members + // + + FragmentAlphaBeta const& s_prime_; + FragmentAlphaBeta const& m_prime_; + +public: + /// Constructs the function object, possibly loading from pointers in host + /// memory + CUTLASS_HOST_DEVICE + MemoryEfficientAttentionNormalize(FragmentAlphaBeta const& s_prime, + FragmentAlphaBeta const& m_prime) + : s_prime_(s_prime), m_prime_(m_prime) + { + } + + /// Returns true if source is needed + CUTLASS_HOST_DEVICE + bool is_source_needed() const { return !isFirst; } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) {} + + /// Computes linear scaling: D = alpha * accumulator + beta * source + CUTLASS_HOST_DEVICE + FragmentOutput operator()(int row, + FragmentAccumulator const& accumulator, + FragmentSource const& source) const + { + assert(!isFirst); + + // Convert source to internal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter + accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; + ElementCompute beta = alpha * m_prime_[row]; + + intermediate = mul_add_source(beta, converted_source); // X = beta * C + + intermediate = mul_add_accumulator( + alpha, converted_accumulator, intermediate); // D = alpha * Accum + X + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + CUTLASS_HOST_DEVICE + FragmentOutput operator()(int row, FragmentAccumulator const& accumulator) const + { + assert(isFirst); + + // Convert source to internal compute numeric type + NumericArrayConverter + accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + ComputeFragment intermediate; + multiplies mul_accumulator; + + ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; + + intermediate = mul_accumulator(alpha, converted_accumulator); // X = alpha * C + uniform + + return destination_converter(intermediate); + } +}; + +} // namespace thread + +namespace threadblock { +template +struct ApplyEpilogueOp< + thread::MemoryEfficientAttentionNormalize> { + using Op = thread::MemoryEfficientAttentionNormalize; + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum, + typename Op::FragmentSource const& source) + { + return output_op(row_id, accum, source); + } + static CUTLASS_DEVICE typename Op::FragmentOutput + apply(Op const& output_op, int row_id, typename Op::FragmentAccumulator const& accum) + { + return output_op(row_id, accum); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h new file mode 100644 index 0000000000000000000000000000000000000000..46fb2bf17c1c1c20400244593dd8a98a91973d0d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h @@ -0,0 +1,168 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include + +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/functional.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template +struct ArrayExponential { + CUTLASS_HOST_DEVICE + Array operator()( + Array const& input) const + { + Array result; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { result[i] = expf(input[i]); } + + return result; + } +}; + +template +struct ArrayExponential { + CUTLASS_DEVICE + Array operator()(Array const& input) const + { + Array result; + + int const kVectorCount = ElementsPerAccess / 2; + + __half2 const* input_ptr = reinterpret_cast<__half2 const*>(input.raw_data()); + __half2* res_ptr = reinterpret_cast<__half2*>(result.raw_data()); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kVectorCount; ++i) { res_ptr[i] = h2exp(input_ptr[i]); } + + return result; + } +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies: +/// output <- (input - lse).exp() +template +class ApplyLogSumExp { +public: + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementLSE = ElementLSE_; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + static const ScaleType::Kind kScale = cutlass::epilogue::thread::ScaleType::NoBetaScaling; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentLSE = Array; + using FragmentScaleBias = FragmentLSE; // Used by epilogue_smem_accumulator.h + +public: + // + // Methods + // + + CUTLASS_HOST_DEVICE + ApplyLogSumExp() {} + + /// Returns true if source is needed + CUTLASS_HOST_DEVICE + bool is_source_needed() const { return true; } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) {} + + CUTLASS_HOST_DEVICE + FragmentOutput operator()(FragmentAccumulator const& AB, + FragmentLSE const& scale_unused, + // bias used as LSE + FragmentLSE const& bias) const + { + FragmentCompute frag_AB = + NumericArrayConverter()(AB); + FragmentCompute frag_lse_compute = + NumericArrayConverter()(bias); + FragmentCompute frag_compute; + + minus minus_lse; + detail::ArrayExponential apply_exp; + frag_compute = minus_lse(frag_AB, frag_lse_compute); + frag_compute = apply_exp(frag_compute); + + return NumericArrayConverter()( + frag_compute); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h new file mode 100644 index 0000000000000000000000000000000000000000..75833bbfe7d2f3c4a91db1ddc7f4453b270db968 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h @@ -0,0 +1,119 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "custom_mma_multistage.h" +#include "custom_mma_pipelined.h" +#include "cutlass/gemm/threadblock/mma_multistage.h" +#include "cutlass/gemm/threadblock/mma_pipelined.h" + +template +struct MakeCustomMma; + +template +struct MakeCustomMma, + kMaxK> { + // Reduce the number of stages if we don't need that many + static int constexpr kStages = + kMaxK == cutlass::platform::numeric_limits::max() + ? Stages + : cutlass::const_min(Stages, (kMaxK + int(Shape::kK) - 1) / int(Shape::kK)); + using Mma = cutlass::gemm::threadblock::CustomMmaMultistage; +}; + +template +struct MakeCustomMma, + kMaxK> { + using Mma = cutlass::gemm::threadblock::CustomMmaPipelined; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h new file mode 100644 index 0000000000000000000000000000000000000000..bbf91240b90075da6f711dba55cf27e19506f898 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h @@ -0,0 +1,181 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/threadblock/mma_base.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class CustomMmaBase { +public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = + GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + template + struct OperandSharedStorage { + AlignedBuffer buffer; + using TensorRef = TensorRef; + + CUTLASS_DEVICE + static OperandLayout Layout() + { + return OperandLayout::packed({OperandShape::kRow, OperandShape::kColumn}); + } + + /// Returns a TensorRef to the operand + CUTLASS_HOST_DEVICE + TensorRef ref() { return TensorRef{buffer.data(), Layout()}; } + }; + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape; + + using SharedStorageA = + OperandSharedStorage; + using SharedStorageB = + OperandSharedStorage; + using TensorRefA = typename SharedStorageA::TensorRef; + using TensorRefB = typename SharedStorageB::TensorRef; + + struct SharedStorage { + /// Buffer for A operand + SharedStorageA operand_A; + + /// Buffer for B operand + SharedStorageB operand_B; + }; + +protected: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorageA& shared_storageA, + SharedStorageB& shared_storageB, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storageA.ref(), lane_idx), + warp_tile_iterator_B_(shared_storageB.ref(), lane_idx) + { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h new file mode 100644 index 0000000000000000000000000000000000000000..3760ccab852a5ceec65fa6693070a84fec284d8a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h @@ -0,0 +1,714 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/cache_operation.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" + +#include "custom_mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Upper boundon the K dimension + int kMaxK = cutlass::platform::numeric_limits::max(), + /// Used for partial specialization + typename Enable = bool> +class CustomMmaMultistage : public CustomMmaBase { +public: + ///< Base class + using Base = CustomMmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / + Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / + Base::kWarpGemmIterations; + }; + + static bool const kSmemContainsEntireMat = kMaxK <= Shape::kK * Stages; + static constexpr int kNumStagesConcurrentLoad = kSmemContainsEntireMat ? Stages : Stages - 1; + +private: + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + +private: + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + bool prologue_done_; + + // Set to `True` to ensure the accumulator will be zero outside the GEMM + // footprint + bool zero_outside_bounds_; + +public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storageA.ref(), thread_idx), + smem_iterator_B_(shared_storageB.ref(), thread_idx), + prologue_done_(false), + zero_outside_bounds_(false) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + CUTLASS_DEVICE + CustomMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& st, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : CustomMmaMultistage(st.operand_A, st.operand_B, thread_idx, warp_idx, lane_idx) + { + } + + CUTLASS_DEVICE + bool set_prologue_done(bool value) + { + prologue_done_ = value; + return true; + } + + CUTLASS_DEVICE + bool set_zero_outside_bounds(bool value) + { + zero_outside_bounds_ = value; + return true; + } + + template + CUTLASS_DEVICE static void prologue(typename Base::SharedStorage& shared_storage, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) + { + prologue(shared_storage.operand_A, + shared_storage.operand_B, + iterator_A, + iterator_B, + thread_idx, + problem_size_k); + } + + template + CUTLASS_DEVICE static void prologue(typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) + { + SmemIteratorA smem_iterator_A(shared_storageA.ref(), thread_idx); + SmemIteratorB smem_iterator_B(shared_storageB.ref(), thread_idx); + int32_t iter = (problem_size_k + Base::Shape::kK - 1) / Base::Shape::kK; + _prologue(iterator_A, iterator_B, iter, smem_iterator_A, smem_iterator_B); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance(IteratorA& iterator_A, + IteratorB& iterator_B, + int group_start_A = 0, + int group_start_B = 0) + { + iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + if (zero_outside_bounds_ || + SharedMemoryClear == SharedMemoryClearOption::kZfill) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + cutlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast(this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + if (zero_outside_bounds_ || + SharedMemoryClear == SharedMemoryClearOption::kZfill) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + cutlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + template + CUTLASS_DEVICE static void _prologue(IteratorA& iterator_A, + IteratorB& iterator_B, + int32_t& gemm_k_iterations, + SmemIteratorA& smem_iterator_A_, + SmemIteratorB& smem_iterator_B_) + { + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; ++stage, --gemm_k_iterations) { + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast(smem_iterator_A_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + if (kLoadA) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + } + + ++iterator_A; + } + + ++smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast(smem_iterator_B_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + if (kLoadB) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + } + + ++iterator_B; + } + + ++smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + smem_iterator_A_.add_tile_offset({0, 1}); + smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + cutlass::arch::cp_async_fence(); + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC& accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const& src_accum) + { + // + // Prologue + // + + if (!prologue_done_) { + _prologue( + iterator_A, iterator_B, gemm_k_iterations, smem_iterator_A_, smem_iterator_B_); + } else if (!kSmemContainsEntireMat) { + _prologue( + iterator_A, iterator_B, gemm_k_iterations, smem_iterator_A_, smem_iterator_B_); + } else { + gemm_k_iterations -= kNumStagesConcurrentLoad; + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // + // Clear the remaining tiles of SMEM. This is a functional requirement for + // some kernels so that all accumulator elements outside the GEMM footprint + // are zero. + // + + if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) { + /// Iterator to write threadblock-scoped tile of A operand to shared + /// memory + SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_); + + typename IteratorA::AccessType zero_A; + zero_A.clear(); + + last_smem_iterator_A.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast(last_smem_iterator_A.get()); + + *dst_ptr = zero_A; + + ++last_smem_iterator_A; + } + + /// Iterator to write threadblock-scoped tile of B operand to shared + /// memory + SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_); + typename IteratorB::AccessType zero_B; + + zero_B.clear(); + last_smem_iterator_B.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast(last_smem_iterator_B.get()); + + *dst_ptr = zero_B; + + ++last_smem_iterator_B; + } + } + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], + warp_transformed_frag_B[0], + warp_loaded_frag_A[0], + warp_loaded_frag_B[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same::value || + platform::is_same::value) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-kNumStagesConcurrentLoad);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); + + // In case of a non-circular buffer ("kSmemContainsEntireMat") + // make sure we don't load out of bounds data. + if (!kSmemContainsEntireMat || gemm_k_iterations > (-kNumStagesConcurrentLoad) || + warp_mma_k < Base::kWarpGemmIterations - 1) { + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + if (platform::is_same::value || + platform::is_same::value) { + warp_mma(tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma(accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum); + } + + // Issue global->shared copies for the this stage + if (!kSmemContainsEntireMat && warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + if (!kSmemContainsEntireMat) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (!kSmemContainsEntireMat && smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + } + + if (platform::is_same::value || + platform::is_same::value) { + accum = plus_accum(accum, tmp_accum); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h new file mode 100644 index 0000000000000000000000000000000000000000..07b26ca3129968f8d78810df5187dd92c3c166c9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h @@ -0,0 +1,388 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/numeric_conversion.h" + +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" + +#include "custom_mma_base.h" +#include "cutlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter, + /// + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter, + /// Used for partial specialization + typename Enable = bool> +class CustomMmaPipelined : public CustomMmaBase { +public: + ///< Base class + using Base = CustomMmaBase; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages == 2), "MmaPipelined requires kStages set to value 2"); + + static bool const kSmemContainsEntireMat = false; + +private: + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + +protected: + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaPipelined(typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ) + : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storageA.ref(), thread_idx), + smem_iterator_B_(shared_storageB.ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + CUTLASS_DEVICE + CustomMmaPipelined( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& st, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : CustomMmaPipelined(st.operand_A, st.operand_B, thread_idx, warp_idx, lane_idx) + { + } + + CUTLASS_DEVICE + bool set_prologue_done(bool value) + { + // NOT IMPLEMENTED FOR PIPELINED + } + + CUTLASS_DEVICE + bool set_zero_outside_bounds(bool value) + { + // NOT NEEDED FOR PIPELINED + // shared memory will always be zero-filled + } + + template + CUTLASS_DEVICE static void prologue(typename Base::SharedStorage& shared_storage, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) + { + prologue(shared_storage.operand_A, + shared_storage.operand_B, + iterator_A, + iterator_B, + thread_idx, + problem_size_k); + } + + template + CUTLASS_DEVICE static void prologue(typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) + { + // NOT IMPLEMENTED FOR PIPELINED + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC& accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const& src_accum, ///< source accumulator tile + TransformA transform_A = TransformA(), ///< transformation applied to A fragment + TransformB transform_B = TransformB()) + { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_A.clear_mask(gemm_k_iterations <= 1); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* + // issuing shared memory loads (which have the tightest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A.clear_mask(gemm_k_iterations <= 2); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + + warp_mma(accum, warp_frag_A[warp_mma_k % 2], warp_frag_B[warp_mma_k % 2], accum); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h new file mode 100644 index 0000000000000000000000000000000000000000..163dcbf852595c2f97fe950b5be6576d938d4efb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h @@ -0,0 +1,191 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Cutlass provides helper template functions to figure out the right + datastructures to instantiate to run a GEMM with various parameters (see + `cutlass/gemm/threadblock/default_mma.h`). However, due to template + instantiation priority rules, it will only create an MmaMultiStage with + kStages=3 (otherwise creates an MmePipelined - which is not compatible with + FastF32). kStages=3 uses too much shared memory and we want to use kStages=2, + so we just copy-pasted some code from `default_mma.h` and + `default_mma_core.h` files and wrapped this template to allow our usecase. + + This is really only for the FastF32 case - aka using TensorCores with fp32. +*/ + +#pragma once + +#include "cutlass/gemm/threadblock/default_mma.h" +#include "cutlass/gemm/threadblock/default_mma_core_simt.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" + +namespace cutlass { +namespace gemm { +namespace threadblock { + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + typename Enable_ = void> +struct FindDefaultMma { + static constexpr bool AccumulatorsInRowMajor = false; + static constexpr SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone; + using DefaultMma = cutlass::gemm::threadblock::DefaultMma; +}; + +/// Specialization for sm80 / FastF32 / multistage with kStages=2 +template +struct FindDefaultMma 1)>::type> { + using LayoutC = layout::RowMajor; + using OperatorClass = arch::OpClassTensorOp; + using ArchTag = arch::Sm80; + + using DefaultMma_ = cutlass::gemm::threadblock::DefaultMma; + struct DefaultMma : DefaultMma_ { + using MmaCore_ = typename DefaultMma_::MmaCore; + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = + cutlass::gemm::threadblock::MmaMultistage; + }; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h new file mode 100644 index 0000000000000000000000000000000000000000..5e2f0cf681bfe10fdc30c2dfcbe390c5f711dae0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h @@ -0,0 +1,347 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "cutlass/functional.h" +#include "cutlass/gemm/warp/mma_simt_tile_iterator.h" +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h" +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" +#include "cutlass/matrix_shape.h" + +/* +TensorCores have different accumulator layouts. +This file provides a class to easily map the accumulator +i-th element with the corresponding matrix row/col. +*/ + +template +struct AccumLambdaIteratorSm80 { + static_assert(cutlass::platform::is_same::value, + "only RowMajor is supported"); + + using Policy = typename T::Policy; + using InstructionShape = typename T::InstructionShape; + using OpDelta = typename T::OpDelta; + using Shape = typename T::Shape; + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + + static cutlass::MatrixCoord CUTLASS_DEVICE + get_lane_offset(int8_t lane_id, int8_t warp_id, typename T::TensorCoord const& tile_offset) + { + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + return cutlass::MatrixCoord( + quad + tile_offset.row() * Shape::kRow, + lane_in_quad * kElementsPerAccess + tile_offset.column() * Shape::kColumn); + } + + template + CUTLASS_DEVICE static void iterateRows(cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) + { + // See cutlass/gemm/warp/mma_tensor_op_tile_iterator.h + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kRowsPerTile + + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + CUTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col + + lane_offset.column(); + int idx = mma_accum_start + row * kElementsPerAccess + col; + op(accum_m, accum_n, idx); + } + } + + endRow(accum_m); + } + } + } + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) + { + // In each warp, 4 threads will work on the same row + // - the ones with the same `quad` + auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1); + myValue = fn(myValue, otherV); + otherV = __shfl_xor_sync(0xffffffff, myValue, 2); + myValue = fn(myValue, otherV); + int lane_in_quad = (lane_id & 3); + return lane_in_quad == 0; + } +}; + +template +struct AccumLambdaIteratorSm70 { + static_assert(cutlass::platform::is_same::value, + "only RowMajor is supported"); + + using Policy = typename T::Policy; + using InstructionShape = typename T::InstructionShape; + using OpDelta = typename T::OpDelta; + using Shape = typename T::Shape; + using Element = accum_t; + + static int const kElementsPerPartial = 4; + using EleShapePerPatial = + typename cutlass::platform::conditional::value, + cutlass::MatrixShape<2, 2>, + cutlass::MatrixShape<1, 4>>::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; + + static cutlass::MatrixCoord CUTLASS_DEVICE + get_lane_offset(int8_t lane_id, int8_t warp_id, typename T::TensorCoord const& tile_offset) + { + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (cutlass::platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + return cutlass::MatrixCoord(accum_m + tile_offset.row() * Shape::kRow, + accum_n + tile_offset.column() * Shape::kColumn); + } + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) + { + static_assert(cutlass::platform::is_same::value, + "update to support non-float accum"); + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-884-f16 + // T0 & T2 share same line within a quad + auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 1); + myValue = fn(myValue, otherV); + // quad 0 and quad 2 are on the same lines + otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 3); + myValue = fn(myValue, otherV); + return (lane_id & ((1 << 1) | (1 << 3))) == 0; + } + + template + CUTLASS_DEVICE static void iterateRows(cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) + { + CUTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2 + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + + mma_n) * + Policy::MmaIterations::kRow + + mma_m) * + kElementsPerMma; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn / 2 + n + + lane_offset.column(); + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + op(accum_m, accum_n, idx); + } + } + } + } + endRow(accum_m); + } + } + } + } +}; + +template +struct AccumLambdaIteratorSimt { + using Policy = typename T::Policy; + using Iterations = typename T::Iterations; + using Element = typename T::Element; + using Delta = typename T::Delta; + using Shape = typename T::Shape; + static_assert(cutlass::platform::is_same::value, + "only RowMajor is supported"); + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) + { + CUTLASS_PRAGMA_UNROLL + for (int bit = 1; bit < Policy::WarpShape::kColumn; bit *= 2) { + auto otherV = __shfl_xor_sync(0xffffffff, myValue, bit); + myValue = fn(myValue, otherV); + } + return (lane_id & (Policy::WarpShape::kColumn - 1)) == 0; + } + + template + CUTLASS_DEVICE static void iterateRows(cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) + { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + int accum_m = mma_m * Delta::kRow + m + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + int accum_n = mma_n * Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + + lane_offset.column(); + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + int idx = n + Policy::LaneMmaShape::kN * + (mma_n + Iterations::kColumn * + (m + mma_m * Policy::LaneMmaShape::kM)); + op(accum_m, accum_n + n, idx); + } + } + endRow(accum_m); + } + } + } + + static cutlass::MatrixCoord CUTLASS_DEVICE + get_lane_offset(int8_t lane_id, int8_t warp_id, typename T::TensorCoord const& tile_offset) + { + static_assert(cutlass::platform::is_same>::value, + ""); + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + cutlass::MatrixCoord lane_offset = + lane_layout.inverse(lane_id) * + cutlass::MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); + return lane_offset + tile_offset * cutlass::MatrixCoord(Shape::kRow, Shape::kColumn); + } +}; + +template +struct DefaultMmaAccumLambdaIterator; + +// Simt +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaSimtTileIterator, + accum_t, + kWarpSize> { + using WarpIterator = + typename cutlass::gemm::warp::MmaSimtTileIterator; + using Iterator = AccumLambdaIteratorSimt; +}; + +// TensorOp - Volta +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator>, + accum_t, + kWarpSize> { + using WarpIterator = typename cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + cutlass::MatrixShape<1, 1>>; + using Iterator = AccumLambdaIteratorSm70; +}; + +// TensorOp - Sm75+ +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp:: + MmaTensorOpAccumulatorTileIterator, + accum_t, + kWarpSize> { + using WarpIterator = typename cutlass::gemm::warp:: + MmaTensorOpAccumulatorTileIterator; + using Iterator = AccumLambdaIteratorSm80; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h new file mode 100644 index 0000000000000000000000000000000000000000..40d3265c7a638396ca5c12282f7da56fa63c925b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h @@ -0,0 +1,1939 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/threadblock/default_epilogue_simt.h" +#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "cutlass/functional.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/platform/platform.h" +#include "cutlass/transform/threadblock/vector_iterator.h" + +#include "../epilogue/epilogue_thread_apply_logsumexp.h" +#include "../gemm/mma_accum_lambda_iterator.h" +#include "../gemm_kernel_utils.h" +#include "../iterators/make_residual_last.h" +#include "../iterators/transpose_warp_iterator.h" +#include "../iterators/warp_iterator_from_smem.h" +#include "cutlass/epilogue/threadblock/epilogue_smem_accumulator.h" +#include "cutlass/gemm/threadblock/mma_base.h" +#include "cutlass/gemm/threadblock/mma_multistage.h" +#include "cutlass/gemm/threadblock/mma_pipelined.h" +#include "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" + +namespace cutlass { +namespace gemm { +namespace threadblock { + +/// Shared storage object needed by accumulator +/// From 13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h +template +class AccumulatorSharedStorage { +public: + // + // Type definitions + // + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using Padding = Padding_; + + /// Tensor reference to the accumulator + using TensorRefAccum = cutlass::TensorRef; + + /// Shape of the accumulator matrix in shared memory + using ShapeAccum = + cutlass::MatrixShape; + +public: + // + // Data members + // + + /// Buffer for accumulator + cutlass::AlignedBuffer accum; + +public: + // + // Methods + // + + /// Returns a layout object for the Accum matrix + CUTLASS_DEVICE + static Layout LayoutAccum() { return Layout::packed({ShapeAccum::kRow, ShapeAccum::kColumn}); } + + /// Returns a TensorRef to the Accumulator + CUTLASS_HOST_DEVICE + TensorRefAccum accum_ref() { return TensorRefAccum{accum.data(), LayoutAccum()}; } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + // Maximum value for K + int kMaxK, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaBaseFromSharedMemory { +public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = + GemmShape; + using WarpCount1 = WarpCount; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK); + static int const kWarpGemmIterations1 = kWarpGemmIterations; + + /// Number of stages + static int const kStages = Stages; + + /// If this is true, we fill the entire shmem buffer at start + /// and don't need to iterate through it in a circular fashion + static bool const kSmemContainsEntireB = kMaxK <= Shape::kK * kStages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape; + + public: + // + // Data members + // + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + // + // Methods + // + + /// Returns a layout object for the B matrix + CUTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() + { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the B operand + CUTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { return TensorRefB{operand_B.data(), LayoutB()}; } + }; + +protected: + // + // Data members + // + + // /// Iterator to load a warp-scoped tile of A operand from shared memory + // typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + /// Construct from tensor references + CUTLASS_DEVICE + MmaBaseFromSharedMemory( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage& shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) + { + } +}; + +namespace { + +// has necessary trait compliance with WarpIteratorFromSmem but doesn't do +// anything, can be default initialized, and uses fragment that takes up +// (almost) no space. this warp iterator is selected at compile time when +// elementwise on-the-fly scaling for operand A is disabled, in which case +// operations related to loading scale factors for operand A get wiped out by +// the compiler. +template +class NoOpWarpIteratorScale { +public: + // in pipelined+multistage MMA implementations we keep an array of fragments. + // if we aren't using scaling we don't want to waste registers on fragments + // of scale elements, so ideally this would be sized 0. + // Since arrays of zero-sized objects are not allowed, using size as 1. + // The compiler will most likely wipe it out anyways. + using Fragment = cutlass::Array; + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale() {} + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale(TensorRef const&, int) {} + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale& add_tile_offset(typename TensorRef::TensorCoord const&) { return *this; } + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale& operator++() { return *this; } + + CUTLASS_DEVICE + void load(Fragment&) const {} +}; + +// if scaling is enabled, performs fragment elementwise multiplication between +// fragment and its scaling factor. +template +class FragmentElementwiseScaler; + +// specialization for scaling being enabled. +template +class FragmentElementwiseScaler { +public: + // cast scale_frag to correct type then apply elementwise to fragment + CUTLASS_DEVICE + static Fragment apply(Fragment frag, FragmentScale const& scale_frag) + { + Fragment converted_scale_frag = + cutlass::NumericArrayConverter()(scale_frag); + return cutlass::multiplies()(frag, converted_scale_frag); + } +}; + +// specialization for scaling being disabled. doesn't do anything and should +// just get wiped out by the compiler. +template +class FragmentElementwiseScaler { +public: + CUTLASS_DEVICE + static Fragment apply(Fragment frag, FragmentScale const&) { return frag; } +}; +} // namespace + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + // BEGIN smem + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA, + /// whether or not to perform elementwise multiplication of A + // by another matrix (A_scale) that is also kept in shared memory prior + // to matmul A @ B + bool ScaleOperandA_, + // Accumulator type + typename AccumulatorSharedStorage, + // END smem + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter, + /// Used for partial specialization + typename Enable = bool> +class MmaPipelinedFromSharedMemory + : public MmaBaseFromSharedMemory { +public: + ///< Base class + using Base = MmaBaseFromSharedMemory; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + static constexpr bool ScaleOperandA = ScaleOperandA_; + + ///< loads fragments of A_scale from shared memory if operand A scaling is + ///< enabled. otherwise no-op. + using WarpIteratorAScale = typename cutlass::platform::conditional< + ScaleOperandA, + WarpIteratorA, + NoOpWarpIteratorScale>::type; + + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorB = SmemIteratorB_; + + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages == 2), "MmaPipelined requires kStages set to value 2"); + +private: + using WarpFragmentA = typename Operator::FragmentA; + + /// fragment type of OperandA elementwise scaling matrix. (almost) empty + /// if operand A scaling is disabled. + using WarpFragmentAScale = typename WarpIteratorAScale::Fragment; + + using WarpFragmentB = typename Operator::FragmentB; + + /// applies scaling factor to operand A fragment if operand A scaling is + /// enabled. otherwise no-op. + using FragmentAScaler = + FragmentElementwiseScaler; + +protected: + // /// Iterator to write threadblock-scoped tile of A operand to shared memory + // SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + /// Iterator to load a warp-scoped tile of A operand from intermediate + /// accumulator tile + WarpIteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of A_scale from intermediate + /// accumulator tile (only used if ScaleOperandA_ is true) + WarpIteratorAScale warp_tile_iterator_A_scale_; + +public: + /// constructor for MMA with operand A scaling enabled. + CUTLASS_DEVICE + MmaPipelinedFromSharedMemory( + // shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& shared_storage, + // warp iterator over A tile held in shared memory + WarpIteratorA warp_iter_a, + // warp iterator over A_scale tile held in shared memory + WarpIteratorAScale warp_iter_a_scale, + int thread_idx, + int warp_idx, + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A_(warp_iter_a), + warp_tile_iterator_A_scale_(warp_iter_a_scale), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_A_scale_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + /// Construct from tensor references + CUTLASS_DEVICE + MmaPipelinedFromSharedMemory( + typename Base::SharedStorage& shared_storage, ///< Shared storage needed for internal use + ///< by threadblock-scoped GEMM + AccumulatorSharedStorage& accumulator_shared_storage, + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx, ///< ID of each thread within a warp + int problem_size_0_n) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A_(accumulator_shared_storage.accum_ref(), lane_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + // For API compatibility with MmaMultistageFromSharedMemory + // but not supported as it worsens perf: older gpus < sm80 don't + // support async transfers and have to waste registers + CUTLASS_DEVICE + void set_prologue_done(bool value) {} + CUTLASS_DEVICE + static void prologue(typename Base::SharedStorage& shared_storage, + IteratorB iterator_B1, + int thread_idx, + int problem_size_0_n) + { + } + + CUTLASS_DEVICE + static void drain_cp_asyncs() {} + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()(int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC& accum, ///< destination accumulator tile + // IteratorA iterator_A, ///< iterator over A + // operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const& src_accum, ///< source accumulator tile + // TransformA transform_A = TransformA(), ///< transformation + // applied to A fragment + TransformB transform_B = TransformB()) + { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentB tb_frag_B; + + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_B.set_residual_tile(gemm_k_iterations == 1); + iterator_B.load(tb_frag_B); + + ++iterator_B; + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_B_; + + __syncthreads(); + + // remember that WarpFragmentAScale and WarpIteratorAScale are empty/no-op + // if scaling is disabled. + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentAScale warp_frag_A_scale[2]; + WarpFragmentB warp_frag_B[2]; + warp_frag_A[0].clear(); + warp_frag_A_scale[0].clear(); + warp_frag_B[0].clear(); + + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_A_scale_.load(warp_frag_A_scale[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_B.set_residual_tile(gemm_k_iterations == 2); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* + // issuing shared memory loads (which have the tightest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + bool hasNext = true; + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + // Write fragments to shared memory + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory SMEM: Don't reset iterator A, as + // we are continuing our iteration at this point + if (smem_write_stage_idx == 1) { + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } else { + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + } + + smem_write_stage_idx ^= 1; + hasNext = gemm_k_iterations > 1; + } + + // Only read the next if we need to + if (hasNext) { + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_scale_.load(warp_frag_A_scale[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + iterator_B.load(tb_frag_B); + + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_B.set_residual_tile(gemm_k_iterations == 3); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + } + + warp_mma(accum, + FragmentAScaler::apply(warp_frag_A[warp_mma_k % 2], + warp_frag_A_scale[warp_mma_k % 2]), + warp_frag_B[warp_mma_k % 2], + accum); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape1_, + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA1_, + /// whether or not to perform elementwise multiplication of A + // by another matrix (A_scale) that is also kept in shared memory prior + // to matmul A @ B + bool ScaleOperandA_, + // Accumulator type + typename AccumulatorSharedStorage, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB1_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB1_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB1, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy1_, + /// Number of stages, + int Stages_, + int kMaxK_, + /// Used for partial specialization + typename Enable = bool> +class MmaMultistageFromSharedMemory + : public MmaBaseFromSharedMemory { +public: + ///< Base class + using Base = MmaBaseFromSharedMemory; + + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape1 = Shape1_; + ///< Iterates over tiles of B operand in global memory + using IteratorB1 = IteratorB1_; + using IteratorB = IteratorB1; + ///< Policy describing tuning details + using Policy1 = Policy1_; + + using SmemIteratorB1 = SmemIteratorB1_; + using WarpIteratorA1 = WarpIteratorA1_; ///< Iterates over the intermediate + ///< accumulator tile in shared memory + static constexpr bool ScaleOperandA = ScaleOperandA_; + + ///< warp level iterator over A_scale matrix tile kept in shared memory. + ///< if elementwise A scaling is disabled then everything this does is no-op. + using WarpIteratorAScale = typename cutlass::platform::conditional< + ScaleOperandA, + WarpIteratorA1, + NoOpWarpIteratorScale>::type; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpB1 = CacheOpB1; + static constexpr bool kSmemContainsEntireB = Base::kSmemContainsEntireB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC1 = typename Policy1::Operator::FragmentC; + using FragmentC = FragmentC1; + + /// Warp-level Mma + using Operator1 = typename Policy1::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on B operand + static ComplexTransform const kTransformB1 = Operator1::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + static_assert(Base::kWarpGemmIterations1 > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand B + static int const TBLoadIterationsB1 = IteratorB1::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB1 = + (TBLoadIterationsB1 + Base::kWarpGemmIterations1 - 1) / Base::kWarpGemmIterations1; + }; + + static constexpr int kNumStagesConcurrentLoad = kSmemContainsEntireB ? Base::kStages + : Base::kStages - 1; + +private: + using WarpLoadedFragmentA1 = typename Operator1::FragmentA; + /// fragment of OperandA scale matrix. if operand A scaling is disabled this + /// is (almost) empty. + using WarpLoadedFragmentA1Scale = typename WarpIteratorAScale::Fragment; + using WarpLoadedFragmentB1 = typename Operator1::FragmentB; + using WarpTransformedFragmentA1 = typename Operator1::TransformedFragmentA; + using WarpTransformedFragmentB1 = typename Operator1::TransformedFragmentB; + + /// applies elementwise scaling to fragment of A. if operand A scaling is + /// disabled this is a no-op. + using FragmentAScaler = + FragmentElementwiseScaler; + +private: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A1 operand from intermediate + /// accumulator tile + WarpIteratorA1 warp_tile_iterator_A1_; + + /// Iterator to load a warp-scoped tile of A1_scale operand from shared memory + /// if operand A scaling is disabled everything this does is a no-op. + WarpIteratorAScale warp_tile_iterator_A1_scale_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB1 smem_iterator_B1_; + + bool prologue_done_; + +public: + /// constructor for MMA with operand A scaling enabled. + CUTLASS_DEVICE + MmaMultistageFromSharedMemory( + // shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& shared_storage, + // warp level iterator over operand A tile kept in shared memory + WarpIteratorA1 warp_tile_iterator_A1, + // warp level iterator over operand A elementwise scale tile kept in + // shared memory. + WarpIteratorAScale warp_tile_iterator_A1_scale, + int thread_idx, + int warp_idx, + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A1_(warp_tile_iterator_A1), + warp_tile_iterator_A1_scale_(warp_tile_iterator_A1_scale), + smem_iterator_B1_(shared_storage.operand_B_ref(), thread_idx), + prologue_done_(false) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + int warp_idx_mn_1 = warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + // Add per-warp offsets in units of warp-level tiles + warp_tile_iterator_A1_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + warp_tile_iterator_A1_scale_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); + } + + /// Construct from tensor references + CUTLASS_DEVICE + MmaMultistageFromSharedMemory( + typename Base::SharedStorage& shared_storage, ///< Shared storage needed for internal use + ///< by threadblock-scoped GEMM + AccumulatorSharedStorage& accumulator_shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx, + ///< GEMM0 N is used for accumulator extent + int problem_size_0_n) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A1_(accumulator_shared_storage.accum_ref(), lane_idx), + smem_iterator_B1_(shared_storage.operand_B_ref(), thread_idx), + prologue_done_(false) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn_1 = warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + // Add per-warp offsets in units of warp-level tiles + warp_tile_iterator_A1_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); + } + + CUTLASS_DEVICE + void set_prologue_done(bool value) { prologue_done_ = value; } + + CUTLASS_DEVICE + static void prologue(typename Base::SharedStorage& shared_storage, + IteratorB iterator_B1, + int thread_idx, + int problem_size_0_n) + { + SmemIteratorB1 smem_iterator_B1(shared_storage.operand_B_ref(), thread_idx); + _prologue(iterator_B1, + (problem_size_0_n + Base::Shape::kK - 1) / Base::Shape::kK, + smem_iterator_B1); + } + + CUTLASS_DEVICE + static void drain_cp_asyncs() + { + // commit and drain all pending and predicated cp.async pnz from the GEMM + // mainloop + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + __syncthreads(); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance_1(IteratorB1& iterator_B1, int group_start_B1 = 0) + { + iterator_B1.set_iteration_index(group_start_B1 * IteratorB1::kAccessesPerVector); + this->smem_iterator_B1_.set_iteration_index(group_start_B1); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB1; ++j) { + if (group_start_B1 + j < Detail::TBLoadIterationsB1) { + typename IteratorB1::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_B1_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB1::ThreadMap::kElementsPerAccess / + IteratorB1::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B1.get(); + + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B1.valid()); + + ++iterator_B1; + } + ++this->smem_iterator_B1_; + } + } + } + + CUTLASS_DEVICE + static void _prologue(IteratorB& iterator_B1, + int32_t gemm_k_iterations_1, + SmemIteratorB1& smem_iterator_B1_) + { + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; ++stage, --gemm_k_iterations_1) { + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); + iterator_B1.clear_mask(gemm_k_iterations_1 == 0); + + iterator_B1.set_iteration_index(0); + smem_iterator_B1_.set_iteration_index(0); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { + typename IteratorB1::AccessType* dst_ptr = + reinterpret_cast(smem_iterator_B1_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + int const kSrcBytes = sizeof_bits::value * + IteratorB1::ThreadMap::kElementsPerAccess / + IteratorB1::kAccessesPerVector / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B1.get(), iterator_B1.valid()); + + ++iterator_B1; + } + + ++smem_iterator_B1_; + } + + // Move to the next stage + iterator_B1.add_tile_offset({1, 0}); + + smem_iterator_B1_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + cutlass::arch::cp_async_fence(); + } + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); + iterator_B1.clear_mask(gemm_k_iterations_1 == 0); + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations_1_, + ///< destination accumulator tile + FragmentC1& accum, + ///< iterator over B1 operand in global memory + IteratorB1 iterator_B1, + ///< initial value of accumulator + FragmentC1 const& src_accum) + { + // 2nd Gemm + + // + // Prologue + // + // Perform accumulation in the 'd' output operand + accum = src_accum; + + if (!prologue_done_) { + _prologue(iterator_B1, gemm_k_iterations_1_, smem_iterator_B1_); + } else if (!kSmemContainsEntireB) { + // Restore the iterators increments + + int gemm_k_iterations_1 = gemm_k_iterations_1_; + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; ++stage, --gemm_k_iterations_1) { + iterator_B1.set_iteration_index(0); + this->smem_iterator_B1_.set_iteration_index(0); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { ++iterator_B1; } + ++this->smem_iterator_B1_; + } + iterator_B1.add_tile_offset({1, 0}); + this->smem_iterator_B1_.add_tile_offset({1, 0}); + } + iterator_B1.set_residual_tile(gemm_k_iterations_1 <= 1); + iterator_B1.clear_mask(gemm_k_iterations_1 <= 0); + } + + // DEPBAR+SYNC + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // remember that WarpFragmentAScale and WarpIteratorAScale are no-op/empty + // if scaling is disabled. + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA1 warp_loaded_frag_A1[2]; + WarpLoadedFragmentA1Scale warp_loaded_frag_A1_scale[2]; + WarpLoadedFragmentB1 warp_loaded_frag_B1[2]; + WarpTransformedFragmentA1 warp_transformed_frag_A1[2]; + WarpTransformedFragmentB1 warp_transformed_frag_B1[2]; + + Operator1 warp_mma1; + + warp_tile_iterator_A1_.load(warp_loaded_frag_A1[0]); + ++warp_tile_iterator_A1_; + + warp_tile_iterator_A1_scale_.load(warp_loaded_frag_A1_scale[0]); + ++warp_tile_iterator_A1_scale_; + + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B1[0]); + ++this->warp_tile_iterator_B_; + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma1.transform( + warp_transformed_frag_A1[0], + warp_transformed_frag_B1[0], + FragmentAScaler::apply(warp_loaded_frag_A1[0], warp_loaded_frag_A1_scale[0]), + warp_loaded_frag_B1[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC1 tmp_accum; + + if (platform::is_same::value || + platform::is_same::value) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + CUTLASS_PRAGMA_UNROLL + for (int gemm_k_iterations_1 = gemm_k_iterations_1_ - (Base::kStages - 1); + gemm_k_iterations_1 > (-Base::kStages + 1); + gemm_k_iterations_1--) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations1; ++warp_mma_k) { + // Load warp-level tile from accumulator fragment (A) + // or shared memory (operand B) + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations1); + // skip warp tile loading for the last kgroup (we are out of the buf) + if (gemm_k_iterations_1 > (-Base::kStages + 2) || + warp_mma_k < Base::kWarpGemmIterations1 - 1) { + warp_tile_iterator_A1_.load(warp_loaded_frag_A1[(warp_mma_k + 1) % 2]); + warp_tile_iterator_A1_scale_.load( + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); + } + ++warp_tile_iterator_A1_; + ++warp_tile_iterator_A1_scale_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma1.transform( + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + FragmentAScaler::apply(warp_loaded_frag_A1[warp_mma_k % 2], + warp_loaded_frag_A1_scale[warp_mma_k % 2]), + warp_loaded_frag_B1[warp_mma_k % 2]); + + if (platform::is_same::value || + platform::is_same::value) { + warp_mma1(tmp_accum, + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + tmp_accum); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma1(accum, + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + accum); + } + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations1 - 1) { + int group_start_iteration_B1; + + group_start_iteration_B1 = warp_mma_k * Detail::kAccessesPerGroupB1; + + if (!kSmemContainsEntireB) { + copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); + } + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations1) { + int group_start_iteration_B1; + group_start_iteration_B1 = (warp_mma_k + 1) * Detail::kAccessesPerGroupB1; + + if (!kSmemContainsEntireB) { + copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); + } + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_B1.add_tile_offset({1, 0}); + + this->smem_iterator_B1_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (!kSmemContainsEntireB) { + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_B1_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy1::kPartitionsK * + Base::kWarpGemmIterations1, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + } + + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 2); + iterator_B1.clear_mask(gemm_k_iterations_1 == 1); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations1) + warp_mma1.transform( + warp_transformed_frag_A1[(warp_mma_k + 1) % 2], + warp_transformed_frag_B1[(warp_mma_k + 1) % 2], + FragmentAScaler::apply(warp_loaded_frag_A1[(warp_mma_k + 1) % 2], + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]), + warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); + } + } + + if (platform::is_same::value || + platform::is_same::value) { + accum = plus_accum(accum, tmp_accum); + } + } +}; + +template +struct DefaultWarpIteratorAFromSharedMemory {}; + +// TensorOp - Ampere half +template +struct DefaultWarpIteratorAFromSharedMemory< + cutlass::gemm::GemmShape<32, 32, 32>, + cutlass::gemm::GemmShape<16, 8, 8>, + RegularWarpIterator, + Policy, + typename platform::enable_if<(sizeof_bits::value == 16 && + Policy::Operator::Policy::OpDelta::kRow == 1)>::type> { + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + using WarpShape = cutlass::MatrixShape<32, 32>; + + using WarpIterator = + cutlass::gemm::warp::WarpIteratorFromSmem; +}; + +// TensorOp - Ampere f32 +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<16, 8, 8>, + RegularWarpIterator, + Policy, + typename platform::enable_if<(sizeof_bits::value != 16 || + Policy::Operator::Policy::OpDelta::kRow != 1)>::type> { + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + + using WarpIterator = cutlass::gemm::warp::MmaTensorOpMultiplicandTileAccessIterator< + cutlass::MatrixShape, + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::layout::RowMajor, + cutlass::MatrixShape, + OpDelta::kRow, + kWarpSize>; +}; + +// TensorOp - Volta +template +struct DefaultWarpIteratorAFromSharedMemory, + RegularWarpIterator, + Policy> { + using InstructionShape = cutlass::gemm::GemmShape<16, 16, 4>; + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + + using WarpIterator = cutlass::gemm::warp::MmaVoltaTensorOpMultiplicandTileIterator< + cutlass::MatrixShape<32, 32>, // MatrixShape, + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>, + cutlass::MatrixShape<16, 4>, + OpDelta::kRow, + kWarpSize>; +}; + +// Simt +template +struct DefaultWarpIteratorAFromSharedMemory, + RegularWarpIterator, + Policy> { + using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; + static constexpr auto kWarpSize = 32; + + // We just use the same iterator, as we reproduced the same shared-memory + // schema. Just modify it to handle non-complete tiles. + using WarpIterator = RegularWarpIterator; +}; + +// Converts a "regular" Mma into their counterpart from shared memory +template +struct DefaultMmaFromSharedMemory; + +// Mma pipelined +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_, + /// Transformation applied to B operand + typename TransformB_, + typename AccumulatorSharedStorage_, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA> +struct DefaultMmaFromSharedMemory, + AccumulatorSharedStorage_, + kScaleOperandA, + kTransposeA> { + static constexpr int kWarpSize = 32; + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + + using RegularMma = MmaPipelined; + + using WarpShape = typename Policy_::Operator::Shape; + using InstructionShape = typename Policy_::Operator::InstructionShape; + using ArchMmaOperator = typename Policy_::Operator; + + static constexpr bool kIsTransposedA = false; + using WarpIteratorA = + typename DefaultWarpIteratorAFromSharedMemory::WarpIterator; + using IteratorB = + typename cutlass::transform::threadblock::MakeIteratorResidualLast::Iterator; + + using Mma = + typename cutlass::gemm::threadblock::MmaPipelinedFromSharedMemory; +}; + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + typename AccumulatorSharedStorage_, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA> +struct DefaultMmaFromSharedMemory, + AccumulatorSharedStorage_, + kScaleOperandA, + kTransposeA> { + static constexpr int kWarpSize = 32; + + using RegularMma = MmaMultistage; + + using WarpShape = typename Policy_::Operator::Shape; + using InstructionShape = typename Policy_::Operator::InstructionShape; + using WarpIteratorA_ = + typename DefaultWarpIteratorAFromSharedMemory::WarpIterator; + using WarpIteratorTranspose = TransposeWarpIterator; + static constexpr bool kIsTransposedA = WarpIteratorTranspose::kSupportsTranspose && kTransposeA; + using WarpIteratorA = typename platform:: + conditional::type; + + static int constexpr kMaxK = kIsTransposedA ? AccumulatorSharedStorage_::Shape::kM + : AccumulatorSharedStorage_::Shape::kN; + // Reduce the number of stages if we don't need that many + static int constexpr kStagesMax = (kMaxK + int(Shape_::kK) - 1) / int(Shape_::kK); + static int constexpr kStages = cutlass::const_min(Stages, kStagesMax); + + using IteratorB = + typename cutlass::transform::threadblock::MakeIteratorResidualLast::Iterator; + using Mma = typename cutlass::gemm::threadblock::MmaMultistageFromSharedMemory< + Shape_, + WarpIteratorA, + kScaleOperandA, + AccumulatorSharedStorage_, + IteratorB, + SmemIteratorB_, + RegularMma::kCacheOpB, + ElementC_, + LayoutC_, + Policy_, + kStages, + kMaxK>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct B2bGemm; + +// Tensor Cores >= Sm75 specialization (Ampere ...) +template < /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_, + typename Operator, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm< + cutlass::gemm::warp:: + MmaTensorOpAccumulatorTileIterator, + Operator, + scalar_t, + WarpShape_, + ThreadblockShape_> { + using IteratorC = typename cutlass::gemm::warp:: + MmaTensorOpAccumulatorTileIterator; + using FragmentC = typename IteratorC::Fragment; + using InstructionShape = InstructionShape_; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using accum_t = Element_; + using lse_scalar_t = float; + + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + + // Iterator to load accumulators (results of matmul in registers) + using FragmentIteratorAccumulator = cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape, + InstructionShape, + accum_t, + typename Operator::Policy::Operator::FragmentC, + cutlass::layout::RowMajor>; + + // Iterator to store to shared-memory + using SmemIteratorD0 = + typename cutlass::epilogue::warp::TileIteratorTensorOp; + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage; + // We need to provide an operation for the epilogue. Let's create an + // operation that does nothing (ScaleType::Nothing), just converts + // from accum_t (float) -> scalar_t (can be half) + using OutputOpNoOp = cutlass::epilogue::thread::LinearCombination< + typename SmemIteratorD0::Element, // ElementOutput + FragmentIteratorAccumulator::Fragment::kElements, + accum_t, // ElementAccumulator + typename SmemIteratorD0::Element, // ElementCompute + cutlass::epilogue::thread::ScaleType::Nothing>; + using Epilogue = cutlass::epilogue::threadblock::EpilogueSmemAccumulator< + SmemIteratorD0, + FragmentIteratorAccumulator, + SmemIteratorD0, // ScaleBiasIterator + // - not used + OutputOpNoOp>; + + // Epilogue 2: with LSE (for backwards pass) + static int const kElementsPerAccess = 2; // TODO: Why 2? + using IteratorAccumulatorLSE = cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + // Shape + cutlass::MatrixShape, + // WarpShape + cutlass::MatrixShape, + lse_scalar_t, + cutlass::layout::RowMajor, + kElementsPerAccess>>; + using EpilogueOpApplyLSE = cutlass::epilogue::thread::ApplyLogSumExp< + scalar_t, // ElementOutput_ + lse_scalar_t, // ElementLSE_ + accum_t, // ElementAccumulator_ + accum_t, // ElementCompute_ + 128 / cutlass::sizeof_bits::value + // FragmentIteratorAccumulator::Fragment::kElements + // InstructionShape::kM * InstructionShape::kN / 32 + >; + using EpilogueWithLSE = + cutlass::epilogue::threadblock::EpilogueSmemAccumulator; + + static void CUTLASS_DEVICE accumToSmem(AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) + { + SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); + smem_iterator_attn.add_tile_offset( + tile_coords * cutlass::MatrixCoord{SmemIteratorD0::TileIterations::kRow, + SmemIteratorD0::TileIterations::kColumn}); + Epilogue epilogue; + epilogue(OutputOpNoOp({}), smem_iterator_attn, accum); + } + + static void CUTLASS_DEVICE accumApplyLSEToSmem(AccumulatorSharedStorage& shared_storage, + FragmentC& accum, + lse_scalar_t const* lse, + int32_t lse_extents, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) + { + constexpr int32_t kAlignLSE = 32; + IteratorAccumulatorLSE iterator_lse( + lse, + {(int32_t)0, (int32_t)ceil_div(lse_extents, kAlignLSE) * kAlignLSE}, + thread_id, + warp_id, + cutlass::MatrixCoord{0, 0} // offset + ); + + SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); + smem_iterator_attn.add_tile_offset( + tile_coords * cutlass::MatrixCoord{SmemIteratorD0::TileIterations::kRow, + SmemIteratorD0::TileIterations::kColumn}); + EpilogueWithLSE epilogue; + EpilogueOpApplyLSE minus_lse_exp({}); + epilogue(minus_lse_exp, + smem_iterator_attn, + accum, + // scale - unused + iterator_lse, + // bias + iterator_lse); + } +}; + +// Volta Specialization +// only supported for f16 +template +struct B2bGemm, + float, + cutlass::layout::RowMajor, + cutlass::gemm::GemmShape<16, 16, 4>, + cutlass::MatrixShape<1, 1>>, + Operator, + cutlass::half_t, + WarpShape_, + ThreadblockShape_> { + using IteratorC = cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + cutlass::MatrixShape<32, 32>, + float, + cutlass::layout::RowMajor, + cutlass::gemm::GemmShape<16, 16, 4>, + cutlass::MatrixShape<1, 1>>; + using scalar_t = cutlass::half_t; + using accum_t = IteratorC::Element; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using FragmentC = IteratorC::Fragment; + using lse_scalar_t = float; + + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + using SmemIteratorD0 = + cutlass::epilogue::warp::TileIteratorVoltaTensorOp, + scalar_t, + SmemAccumulatorLayout>; + + // // Storage in shared-memory for Q.Kt + using AccumulatorSharedStorage = cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + scalar_t, + cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise< + 16, + 32>, // typename SmemIteratorD0::TensorLayout, + cutlass::MatrixShape<0, 0> // Padding + >; + + using OutputLayout = cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>; + using TensorRef = cutlass::TensorRef; + using Policy = typename IteratorC::Policy; + using Element = accum_t; + // Those are MmaVoltaTensorOpAccumulatorTileIterator private fields + // Let's copy their values + static int const kElementsPerPartial = 4; + using EleShapePerPatial = + typename cutlass::platform::conditional::value, + cutlass::MatrixShape<2, 2>, + cutlass::MatrixShape<1, 4>>::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; + + static void CUTLASS_DEVICE accumToSmem(AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) + { + // ctor - from MmaVoltaTensorOpAccumulatorTileIterator + TensorRef ref_(shared_storage.accum_ref()); + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (cutlass::platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + cutlass::MatrixCoord lane_offset(accum_m, accum_n); + + // Tile offset + ref_.add_coord_offset(tile_coords * cutlass::MatrixCoord({IteratorC::Shape::kRow, + IteratorC::Shape::kColumn})); + + using AccessType = cutlass::Array; + + // store - from MmaVoltaTensorOpAccumulatorTileIterator + CUTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + CUTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int mma_accum_start = (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + + mma_n) * + Policy::MmaIterations::kRow + + mma_m) * + kElementsPerMma; + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn / 2; + int r = (accum_m + lane_offset.row()); + AccessType to_store; + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + int c = (accum_n + n + lane_offset.column()); + to_store[n] = scalar_t(accum[idx]); + } + int c = (accum_n + lane_offset.column()); + assert(r < 32); + assert(c < 32); + *reinterpret_cast(ref_.data() + ref_.offset({r, c})) = + to_store; + } + } + } + } + } + } + } + + static void CUTLASS_DEVICE accumApplyLSEToSmem(AccumulatorSharedStorage& shared_storage, + typename IteratorC::Fragment& accum, + lse_scalar_t const* lse, + int lse_extent, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) + { + // Non-optimized way to apply LSE to registers + // NOTE: accum is attn.T + // TODO: Optimize for each architecture + static constexpr int WarpSize = 32; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator::Iterator; + auto lane_offset = AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); + + cutlass::Array lse_prefetched; + lse_prefetched.clear(); + int rowIdx = 0; + int colIdx = 0; + AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + ++rowIdx; + colIdx = 0; + }, + [&](int accum_m, int accum_n, int idx) { + if (rowIdx == 1) { + lse_prefetched[colIdx] = accum_n < lse_extent + ? lse[accum_n] + : platform::numeric_limits::infinity(); + } + accum[idx] = expf(accum[idx] - lse_prefetched[colIdx]); + ++colIdx; + }, + [&](int accum_m) {}); + accumToSmem(shared_storage, accum, lane_id, tile_coords); + } +}; + +// Simt Specialization +// for f32 on Sm70-Sm75 and f16/f32 below + +template +struct B2bGemm, + cutlass::gemm::Operand::kC, + float, + cutlass::layout::RowMajor, + OperatorPolicy, + 1, + 1>, + Operator, + scalar_t, + WarpShape_, + ThreadblockShape_> { + using IteratorC = cutlass::gemm::warp::MmaSimtTileIterator, + cutlass::gemm::Operand::kC, + float, + cutlass::layout::RowMajor, + OperatorPolicy, + 1, + 1>; + using accum_t = typename IteratorC::Element; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using FragmentC = typename IteratorC::Fragment; + using lse_scalar_t = float; + + // Storage in shared-memory for Q.Kt + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage // Padding + >; + + static void CUTLASS_DEVICE accumToSmem(AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) + { + using Policy = typename IteratorC::Policy; + using Element = typename IteratorC::Element; + using Iterations = typename IteratorC::Iterations; + using Delta = typename IteratorC::Delta; + + auto ref_ = shared_storage.accum_ref(); + // ctor - MmaSimtTileIterator + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); + + ref_.add_coord_offset(lane_offset); + + // Tile offset + ref_.add_coord_offset(tile_coords * cutlass::MatrixCoord({IteratorC::Shape::kRow, + IteratorC::Shape::kColumn})); + + // store - MmaSimtTileIterator + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + int r = Policy::LaneMmaShape::kM * (mma_m * Policy::WarpShape::kRow) + m; + int c = mma_n * Delta::kColumn + n; + int idx = n + Policy::LaneMmaShape::kN * + (mma_n + Iterations::kColumn * + (m + mma_m * Policy::LaneMmaShape::kM)); + ref_.at({r, c}) = scalar_t(accum[idx]); + } + } + } + } + } + + static void CUTLASS_DEVICE accumApplyLSEToSmem(AccumulatorSharedStorage& shared_storage, + typename IteratorC::Fragment& accum, + lse_scalar_t const* lse, + int lse_extent, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) + { + // Non-optimized way to apply LSE to registers + // NOTE: accum is attn.T + // TODO: Optimize for each architecture + static constexpr int WarpSize = 32; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator::Iterator; + auto lane_offset = AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); + + cutlass::Array lse_prefetched; + lse_prefetched.clear(); + int rowIdx = 0; + int colIdx = 0; + AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + ++rowIdx; + colIdx = 0; + }, + [&](int accum_m, int accum_n, int idx) { + if (rowIdx == 1) { + lse_prefetched[colIdx] = accum_n < lse_extent + ? lse[accum_n] + : platform::numeric_limits::infinity(); + } + accum[idx] = expf(accum[idx] - lse_prefetched[colIdx]); + ++colIdx; + }, + [&](int accum_m) {}); + accumToSmem(shared_storage, accum, lane_id, tile_coords); + } +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..c102234a4dfb040863bc5babad46e5bcfa1fd9b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h @@ -0,0 +1,253 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include "cutlass/arch/mma.h" + +template +struct CheckArch { + static constexpr bool isPreVolta = arch::kMinComputeCapability < 70; + static constexpr bool isPreAmpere = + arch::kMinComputeCapability < 80 && arch::kMinComputeCapability >= 70; + static constexpr bool isAmpere = arch::kMinComputeCapability >= 80; +#if defined(__CUDA_ARCH__) + static constexpr bool compiler_cc = arch::kMinComputeCapability * 10 <= __CUDA_ARCH__; +#else + static constexpr bool compiler_cc = true; +#endif + static constexpr bool value = (isPreVolta && std::is_same_v) || + (isPreAmpere && !std::is_same_v) || + isAmpere && compiler_cc; +}; + +#define DISPATCH_ARCHTAG(CC, func) \ + { \ + if constexpr (GPU_ARCH >= 80) { \ + if (CC >= 80) { \ + using ArchTag = cutlass::arch::Sm80; \ + func; \ + } else { \ + EVOFORMER_CHECK(false, "Compile flag error. Unexpected GPU"); \ + } \ + } else if constexpr (GPU_ARCH >= 75) { \ + if (CC >= 75) { \ + using ArchTag = cutlass::arch::Sm75; \ + func; \ + } else { \ + EVOFORMER_CHECK(false, "Compile flag error. Unexpected GPU"); \ + } \ + } else if constexpr (GPU_ARCH >= 70) { \ + if (CC >= 70) { \ + using ArchTag = cutlass::arch::Sm70; \ + func; \ + } else { \ + EVOFORMER_CHECK(false, "Compile flag error. Unexpected GPU"); \ + } \ + } else { \ + EVOFORMER_CHECK(false, "Only GPUs with Tensor Core are supported for now"); \ + } \ + } + +#define DISPATCH_TYPES(tensor, func) \ + { \ + if (tensor.scalar_type() == at::ScalarType::Half) { \ + using scalar_t = cutlass::half_t; \ + using torch_scalar_t = at::Half; \ + func; \ + } else if (tensor.scalar_type() == at::ScalarType::BFloat16) { \ + using scalar_t = cutlass::bfloat16_t; \ + using torch_scalar_t = at::BFloat16; \ + func; \ + } else { \ + EVOFORMER_CHECK(false, "Only fp16 and bf16 supported at the moment"); \ + } \ + } + +#define DISPATCH_BOOL(BOOL_V, BOOL_NAME, F) \ + { \ + if (BOOL_V) { \ + constexpr bool BOOL_NAME = true; \ + F(); \ + } else { \ + constexpr bool BOOL_NAME = false; \ + F(); \ + } \ + } + +#ifdef TORCH_CHECK +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + EVOFORMER_CHECK(uint64_t(PTR) % ALIGNMENT == 0, #PTR " is not correctly aligned") +#define EVOFORMER_CHECK TORCH_CHECK +#elif defined(__CUDACC_RTC__) +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + if (!(uint64_t(PTR) % ALIGNMENT == 0)) { return false; } +#define EVOFORMER_CHECK(COND, ERR) \ + if (!(COND)) { return false; } +#else +#include +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + if (!(uint64_t(PTR) % ALIGNMENT == 0)) { \ + std::cerr << #PTR " is not correctly aligned\n"; \ + return false; \ + } +#define EVOFORMER_CHECK(COND, ERR) \ + if (!(COND)) { \ + std::cerr << "[Evoformer Attention]" << "'" #COND "' failed: " << ERR << "\n"; \ + return false; \ + } +#endif + +namespace gemm_kernel_utils { + +template +constexpr CUTLASS_HOST_DEVICE integer ceil_div(integer n, integer m) +{ + return (n + m - 1) / m; +} + +template +constexpr CUTLASS_HOST_DEVICE integer align_up(integer n, integer m) +{ + return ((n + m - 1) / m) * m; +} + +//////////////////////////////////////////////////////////////////////////////// +// Determine the type of GEMM we do (TensorCores or not, Shapes ...) +// TODO: Maybe we could rely on Cutlass's DefaultGemm templates +//////////////////////////////////////////////////////////////////////////////// + +// Fallback to Simt (FMA on cuda cores) if not in a special case below +template +struct DefaultGemmType { + static constexpr int ThreadK = 8; + static constexpr int WarpK = 8; + static constexpr int kMinimumAlignment = 1; + using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; + using OpClass = cutlass::arch::OpClassSimt; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Specialization for tensorcores with f32 +template +struct DefaultGemmType< + ArchTag, + float, + typename cutlass::platform::enable_if= 80>::type> { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 4; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + using Operator = cutlass::arch::OpMultiplyAddFastF32; +}; + +// Specialization for tensorcores with f16/bf16 - Sm75+ +template +struct DefaultGemmType< + ArchTag, + scalar_t, + typename cutlass::platform::enable_if= 75 && + cutlass::sizeof_bits::value == 16>::type> { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 4; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Specialization for tensorcores with f16 - Volta +template <> +struct DefaultGemmType { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 2; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Enables to do +// `auto x = kCondition ? fa(arg) : fb(arg)` +// when `fa` and `fb` have different types +template +struct call_conditional; + +template +struct call_conditional { + template + static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) -> decltype(ta(arg)) + { + return ta(arg); + } +}; + +template +struct call_conditional { + template + static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) -> decltype(tb(arg)) + { + return tb(arg); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Mark a variable as warp-uniform - enables some compiler optimizations +// The cheapest way to do it is just to broadcast it from lane 0 +//////////////////////////////////////////////////////////////////////////////// + +CUTLASS_DEVICE int32_t warp_uniform(int32_t value) +{ + return (int32_t)__shfl_sync(0xffffffff, (unsigned)value, 0); +} + +template +CUTLASS_DEVICE T* warp_uniform(T* ptr) +{ + struct { + union { + T* ptr; + uint32_t asInt[2]; + }; + } p; + p.ptr = ptr; + p.asInt[0] = warp_uniform(p.asInt[0]); + p.asInt[1] = warp_uniform(p.asInt[1]); + return p.ptr; +} +} // namespace gemm_kernel_utils diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h new file mode 100644 index 0000000000000000000000000000000000000000..667f1982d30d295498b952f0a681af29796a1521 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h @@ -0,0 +1,691 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Epilogue iterator that supports prefetching + + Mostly copied from "cutlass/epilogue/threadblock/predicated_tile_iterator.h" +*/ + +#pragma once + +#include "cutlass/arch/arch.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/threadblock/output_tile_thread_map.h" +#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/transform/pitch_linear_thread_map.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in +/// epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | +/// ForwardTileIterator +/// +template +class PredicatedTileIteratorPrefetch { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); + static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); + static_assert(ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc()) + { + } + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { enable(); } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = false; } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = true; } + } + }; + +private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer + uint8_t* byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Scatter indices + int const* indices_; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + // + // Methods + // + +public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorPrefetch(PredicatedTileIteratorParams const& params, + Element* pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord(), + int const* indices = nullptr) + : params_(params), indices_(indices) + { + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + mask_.predicates[c] = + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { mask_.clear(); } + + if (ScatterD && !indices) { mask_.clear(); } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + + if (ScatterD) { + byte_pointer_ = + reinterpret_cast(pointer) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + } + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + CUTLASS_DEVICE + void prefetch_all() + { + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kIterations; ++iter) { + prefetch(); + ++(*this); + } + } + + CUTLASS_DEVICE + void prefetch() + { + uint8_t* byte_pointer = byte_pointer_; + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + AccessType* memory_pointer = reinterpret_cast(byte_pointer); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + // on windows using unsigned long here gives the error + // error: asm operand type size(4) does not match + // type/size implied by constraint 'l' + uint64_t addr = + (uint64_t)((void*)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess]); + asm volatile("prefetch.global.L1 [ %1 ];" : "=l"(addr) : "l"(addr)); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { byte_pointer += params_.increment_row; } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, int64_t byte_offset) const + { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * + LongIndex(params_.stride)); + } + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { byte_pointer += params_.increment_row; } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) const { load_with_byte_offset(frag, 0); } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const + { + uint8_t* byte_pointer = byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * + LongIndex(params_.stride)); + } + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask_.predicates[column]; + + if (UseCUDAStore) { + if (guard) { + memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess] = + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column]; + } + } else { + cutlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { byte_pointer += params_.increment_row; } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) const { store_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void downsample_load_with_byte_offset(Fragment& frag, + int64_t byte_offset, + int convolution_P, + int convolution_Q, + int add_P, + int add_Q, + int problem_N) const + { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + int output_row = row_offset + thread_start_row_; + int output_N = output_row / (convolution_P * convolution_Q); + int output_PQ = output_row % (convolution_P * convolution_Q); + int output_P = output_PQ / convolution_Q; + int output_Q = output_PQ % convolution_Q; + + int input_row = output_N * 2 * convolution_P * 2 * convolution_Q + + (2 * output_P + add_P) * 2 * convolution_Q + 2 * output_Q + + add_Q; + + int64_t byte_offset = (input_row - output_row) * problem_N * sizeof(float); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void upsample_load_with_byte_offset(Fragment& frag, + int64_t byte_offset, + int convolution_P, + int convolution_Q, + int add_P, + int add_Q, + int problem_N) const + { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + int output_row = row_offset + thread_start_row_; + int output_N = output_row / (convolution_P * convolution_Q); + int output_PQ = output_row % (convolution_P * convolution_Q); + int output_P = output_PQ / convolution_Q; + int output_Q = output_PQ % convolution_Q; + int row_add_P = add_P; + int row_add_Q = add_Q; + if (output_P > convolution_P - 2) row_add_P = 0; + if (output_Q > convolution_Q - 2) row_add_Q = 0; + + int input_row = output_N * (convolution_P / 2) * (convolution_Q / 2) + + ((output_P + row_add_P) / 2) * (convolution_Q / 2) + + (output_Q + row_add_Q) / 2; + + int64_t byte_offset = (input_row - output_row) * problem_N * sizeof(float); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + CUTLASS_DEVICE + MatrixCoord thread_start() const + { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_row() const { return thread_start_row_; } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_column() const { return thread_start_column_; } + + /// Extent of the matrix in rows + CUTLASS_DEVICE + Index extent_row() const { return extent_row_; } + + /// Extent of the matrix in columns + CUTLASS_DEVICE + Index extent_column() const { return extent_column_; } + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorPrefetch& operator++() + { + ++state_[0]; + + if (!ScatterD) { byte_pointer_ += params_.advance_row; } + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + + thread_start_row_ += + (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { mask_.clear(); } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { mask_.enable(); } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) const { mask = mask_; } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } +}; + +template +struct MakePrefetchableIterator { + using Iterator = PredicatedTileIteratorPrefetch; +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h new file mode 100644 index 0000000000000000000000000000000000000000..ff0e324c3a6c0899d8439d92658515fb255f2dd4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h @@ -0,0 +1,91 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "predicated_tile_access_iterator_residual_last.h" +#include "predicated_tile_iterator_residual_last.h" + +namespace cutlass { +namespace transform { +namespace threadblock { + +template +struct MakeIteratorResidualLast; + +template +struct MakeIteratorResidualLast< + PredicatedTileIterator> { + using Iterator = PredicatedTileIteratorResidualLast; +}; + +template +struct MakeIteratorResidualLast> { + using Iterator = PredicatedTileAccessIteratorResidualLast; +}; +} // namespace threadblock +} // namespace transform +} // namespace cutlass diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h new file mode 100644 index 0000000000000000000000000000000000000000..7f6a2430845a16aac5891257c5b0600c5226e9a9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h @@ -0,0 +1,1964 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "cutlass/array.h" +#include "cutlass/coord.h" +#include "cutlass/cutlass.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/predicate_vector.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/tensor_view.h" +#include "cutlass/transform/threadblock/predicated_tile_access_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIteratorResidualLast +/// +template +class PredicatedTileAccessIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for pitch-linear +/// data. +/// +template +class PredicatedTileAccessIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : PredicatedTileAccessIteratorParams { + using Base = PredicatedTileAccessIteratorParams; + + // Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : Base(layout.stride(0), + MakePredicatedTileAccessIteratorDesc()()) + { + } + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + +private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + +private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + Mask residual_tile_mask; + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Below is used when Gather is turned on. We need to record strided_offset + /// and contiguous_offset separated to compute the offset by using + /// + /// offset = contiguous_offset + indices[strided_offset] + /// + + /// Gather indices + int const* indices_; + + Index gather_offset_strided; + +private: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) + { + the_predicates.compute_predicates_(extent, is_steady_state); + } + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : params_(params), + pointer_(reinterpret_cast(const_cast(pointer))), + the_predicates(extent), + indices_(indices) + { + the_predicates.set_predicates(thread_id, threadblock_offset); + the_predicates.get_mask(residual_tile_mask); + + // Working around a weird compiler bug happening on P100 for the backward. + // I've seen together: the_predicates.predicates_[0] = 14 (instead of 15) + // residual_tile_mask[0] = 15 (correct) + // + // Adding prints when the value is calculated (in `compute_predicates_`) + // sometimes removes the bug. The consequence is that we skip some + // element of a tensor, leading to wrong results + // Setting `compute_predicates_`'s second argument (`is_steady_state`) to + // true also seems to get rid of the bug - at the cost of twice as many + // comparisons. +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 700) + constexpr bool kWorkAroundCompilerBug = false; +#else + constexpr bool kWorkAroundCompilerBug = true; +#endif + the_predicates.compute_predicates_(extent, true && !kWorkAroundCompilerBug); + + // update internal pointers + Layout layout(params_.stride_); + + if (!Gather) { + add_pointer_offset(layout(the_predicates.thread_offset_)); + } else { + gather_offset_strided = the_predicates.thread_offset_.strided(); + add_pointer_offset(layout(make_Coord(the_predicates.thread_offset_.contiguous(), 0))); + } + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { the_predicates.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool is_residual_tile) + { + if (is_residual_tile) { the_predicates.set_mask(residual_tile_mask); } + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + if (!Gather) { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } else { + add_pointer_offset(Shape::kContiguous * tile_offset.contiguous()); + gather_offset_strided += Shape::kStrided * tile_offset.strided(); + } + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const + { + if (Gather) { + assert(indices_); + + if (!valid()) { return nullptr; } + + LongIndex contiguous_offset = + the_predicates.iteration_contiguous_ * + (ThreadMap::Delta::kContiguous * sizeof_bits::value / 8) + + the_predicates.iteration_vector_; + int strided_index = gather_offset_strided + + the_predicates.iteration_strided_ * ThreadMap::Delta::kStrided; + + LongIndex strided_offset = indices_[strided_index] * LongIndex(params_.stride_) * + sizeof_bits::value / 8; + + return reinterpret_cast(pointer_ + contiguous_offset + strided_offset); + } + + return reinterpret_cast( + pointer_ + the_predicates.iteration_contiguous_ * + (ThreadMap::Delta::kContiguous * sizeof_bits::value) / + 8) + + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { return *this; } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + if (!Gather) { pointer_ += params_.inc_strided_; } + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + if (!Gather) { + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, + // this subtraction as well as the subsequent integer addition are both + // elided by the compiler. + pointer_ -= params_.inc_advance_; + } + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { the_predicates.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { the_predicates.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { the_predicates.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { the_predicates.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { return the_predicates.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for column-major +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType, + Gather>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()), + indices) + { + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { return reinterpret_cast(iterator_.get()); } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType, + Gather>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()), + indices) + { + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { return reinterpret_cast(iterator_.get()); } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank 2 +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingPredicates::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileAccessIteratorResidualLast; + + private: + /// stride of pitch-linear layout (units of Element) + Coord stride_; + /// amount (in byte) to increment pointer to move to next access along + /// contiguous dimension + LongIndex inc_contiguous_; + /// amount (in byte) to increment pointer from first access of current + /// contiguous dimension to first access of next one. + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access of current + /// contiguous dimension to first access of next one. + LongIndex inc_next_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + public: + // Default ctor + CUTLASS_HOST_DEVICE + Params() : stride_(0), inc_contiguous_(0), inc_strided_(0), inc_next_(0), inc_advance_(0) {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : stride_({layout.stride(0), layout.stride(1)}) + { + inc_contiguous_ = (LongIndex(stride_[0]) * ThreadMap::Delta::kContiguous) * + sizeof_bits::value / 8; + + inc_strided_ = (LongIndex(stride_[1]) * ThreadMap::Delta::kStrided) * + sizeof_bits::value / 8; + + inc_next_strided_ = + inc_strided_ - LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_; + + if (kAdvanceRank) { + // advance along strided dimension + inc_advance_ = + Shape::kStrided * LongIndex(stride_[1]) * sizeof_bits::value / 8; + } else { + // advance along contiguous dimension + inc_advance_ = Shape::kContiguous * stride_[0] * sizeof_bits::value / 8; + } + + inc_next_ = inc_advance_ - + LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_ - + LongIndex(ThreadMap::Iterations::kStrided - 1) * inc_strided_; + }; + }; + +private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + // + // Data members + // + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + UnderlyingPredicates the_predicates; + Mask residual_tile_mask; + +private: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) + { + the_predicates.compute_predicates_(extent, is_steady_state); + } + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : params_(params), + pointer_(reinterpret_cast(const_cast(pointer))), + the_predicates(extent) + { + the_predicates.set_predicates(thread_id, threadblock_offset); + + // update internal pointers + Layout layout(params_.stride_); + add_pointer_offset(layout(the_predicates.thread_offset_)); + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { the_predicates.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool is_residual_tile) + { + if (is_residual_tile) { the_predicates.set_mask(residual_tile_mask); } + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[1]); + pointer_ += Shape::kContiguous * tile_offset[0]; + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[0]); + pointer_ += Shape::kStrided * tile_offset[1]; + } + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const + { + return reinterpret_cast(pointer_) + the_predicates.iteration_vector_; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + the_predicates.operator++(); + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { return *this; } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + pointer_ += params_.inc_contiguous_; + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_next_strided_; + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { the_predicates.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { the_predicates.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { the_predicates.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { the_predicates.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return the_predicates.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank 2 +/// column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))){}; + }; + +private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())) + { + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + iterator_.add_tile_offset(make_Coord(tile_offset.row(), tile_offset.column())); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { return reinterpret_cast(iterator_.get()); } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank-2 +/// row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))){}; + }; + +private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())) + { + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + iterator_.add_tile_offset(make_Coord(tile_offset.column(), tile_offset.row())); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { return reinterpret_cast(iterator_.get()); } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for column-major +/// interleaved data. It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class PredicatedTileAccessIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) + { + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { return reinterpret_cast(iterator_.get()); } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major +/// interleaved data. +// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) + { + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast(params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) + { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) + { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { return reinterpret_cast(iterator_.get()); } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) + { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h new file mode 100644 index 0000000000000000000000000000000000000000..8d4173f1a6a2e5a073d7da4c129f471d4f394632 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h @@ -0,0 +1,886 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once +#include +#include +#include +namespace cutlass { +namespace epilogue { +namespace threadblock { + +template +struct atomic_store {}; + +template +struct atomic_store::value>::type> { + using Element = typename AccessType::Element; + static const int kCount = AccessType::kElements; + + CUTLASS_DEVICE + atomic_store(AccessType const& D, void* ptr, bool pred_guard) + { + static_assert(!(kCount % 2), "kCount must be even"); + half2* p = reinterpret_cast(ptr); + uint const* data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + : + : "r"((int)pred_guard)); + for (int i = 0; i < kCount / 2; i++) { + asm volatile(" @p red.relaxed.global.add.noftz.f16x2 [%0], %1;\n" + : + : "l"(p + i), "r"(data[i])); + } + asm volatile("}\n" ::); + } +}; + +template +struct atomic_store::value>::type> { + using Element = typename AccessType::Element; + static const int kCount = AccessType::kElements; + + CUTLASS_DEVICE + atomic_store(AccessType const& D, void* ptr, bool pred_guard) + { + Element* p = reinterpret_cast(ptr); + uint const* data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + : + : "r"((int)pred_guard)); + for (int i = 0; i < kCount; i++) { + asm volatile(" @p red.relaxed.global.add.f32 [%0], %1;\n" + : + : "l"(p + i), "r"(data[i])); + } + asm volatile("}\n" ::); + } +}; + +template +class PredicatedTileIteratorAffineRankNAtomic { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::AffineRankN; + using TensorRef = TensorRef; + using TensorView = TensorView; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = typename Layout::TensorCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); + static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); + static_assert(ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); + static_assert(!(Layout::kRank % 2), + "Layout rank must be even. This assumes the first half of the " + "modes correspond to the 'row' " + "and the second half of the modes correspond to the 'column'"); + + static bool const kBigEndian = false; + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Parameters structure + struct Params { + // + // Data members + // + + Layout layout; + + /// Stride in units of bytes along M modes + Coord stride_m; + + /// Stride in units of bytes along N modes + Coord stride_n; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank / 2 - 1)]; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank / 2 - 1)]; + + int64_t rank2_inc_col; + int64_t rank2_inc_row; + + // + // Methods + // + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(TensorCoord const& extent, Layout const& layout_) : layout(layout_) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i]); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); + } + + if (kBigEndian) { + // "Big Endian" scheme + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i + 1]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]); + } + } else { + // "Little Endian" scheme + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]); + } + } + } + + CUTLASS_HOST_DEVICE + Params(Layout const& layout_) : layout(layout_) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i]); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); + } + + rank2_inc_col = ThreadMap::Delta::kColumn * stride_n[0]; + rank2_inc_row = ThreadMap::Delta::kRow * stride_m[0]; + } + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { enable(); } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = false; } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = true; } + } + }; + +private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + Params params_; + + /// Byte-level pointer + uint8_t* byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in columns + Index extent_col_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column position (assuming steady-state predicates have + /// been computed) + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Offsets in columns, cached for performance + int64_t offset_modes_n_[ThreadMap::Iterations::kColumn]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + +private: + // + // Methods + // + +public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorAffineRankNAtomic( + Params const& params, + Element* pointer, + MatrixCoord extent, + int thread_idx, + MatrixCoord threadblock_offset = MatrixCoord(), + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : params_(params) + { + MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_col_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + if (Layout::kRank > 2) { + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + // + // Compute coordinate and decompose into N modes + // + + int coord_n = thread_start_column_ + c * ThreadMap::Delta::kColumn; + + mask_.predicates[c] = coord_n < extent.column(); + + Coord modes_n; + + int64_t offset_modes_n = 0; + + if (kBigEndian) { + modes_n = CoordinateDecomposition(coord_n, params_.divmod_n); + + offset_modes_n = dot(modes_n, params_.stride_n); + } else { + modes_n = CoordinateDecompositionLittleEndian( + coord_n, params_.divmod_n); + + offset_modes_n = dot(modes_n, params_.stride_n); + } + + offset_modes_n_[c] = offset_modes_n; + } + + if (!pointer) { mask_.clear(); } + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer); + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) + { + uint8_t* byte_pointer = byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + int64_t offset_modes_m = row_begin * params_.stride_m[0]; + + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + // + // Compute coordinate and decompose into M modes + // + + int coord_m = row * ThreadMap::Delta::kRow + row_begin; + + Coord modes_m; + + if (Layout::kRank > 2) { + if (kBigEndian) { + modes_m = CoordinateDecomposition(coord_m, + params_.divmod_m); + } else { + modes_m = CoordinateDecompositionLittleEndian( + coord_m, params_.divmod_m); + } + + offset_modes_m = dot(modes_m, params_.stride_m); + } + + // + // Compute the offset due to modes M + // + + bool row_guard = (coord_m < extent_row_); + int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0]; + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + // + // Compute coordinate and decompose into N modes + // + + if (Layout::kRank > 2) { offset_modes_n = offset_modes_n_[column]; } + + // + // Compute the pointer and access + // + bool guard; + if (Layout::kRank > 2) { + guard = row_guard && mask_.predicates[column]; + } else { + guard = (coord_m < extent_row_) && + ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < + extent_col_); + } + + atomic_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset), + guard); + + if (Layout::kRank == 2) { offset_modes_n += params_.rank2_inc_col; } + } + + if (Layout::kRank == 2) { offset_modes_m += params_.rank2_inc_row; } + } + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } + + CUTLASS_DEVICE + void load(Fragment& frag) {} + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorAffineRankNAtomic& operator++() + { + ++state_[0]; + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + + thread_start_row_ += + (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { state_[2] = 0; } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { mask_.clear(); } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { mask_.enable(); } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) { mask = mask_; } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } +}; + +template +class PredicatedTileIteratorAtomic { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static bool constexpr PermuteD = !layout::is_trivial_permute; + + static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); + static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); + static_assert(ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc()) + { + } + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { enable(); } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = false; } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = true; } + } + }; + +private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer. This pointer is usually for both load() and store(), + /// unless PermuteD is performed. When having PermuteD, byte_pointer_ is only + /// for load(). + uint8_t* byte_pointer_; + + /// Byte-level pointer for store(). Due to PermuteD Op, store_byte_pointer_ + /// may be with different address computation compared to byte_pointer_. + uint8_t* store_byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Scatter indices + int const* indices_; + + /// PermuteDLayout + PermuteDLayout permute_layout_; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + // + // Methods + // + +public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorAtomic(PredicatedTileIteratorParams const& params, + Element* pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord(), + int const* indices = nullptr) + : params_(params), + indices_(indices), + permute_layout_(PitchLinearCoord(extent.column(), extent.row()), + params_.stride * kElementsPerAccess / sizeof(AccessType)) + { + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + mask_.predicates[c] = + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { mask_.clear(); } + + if (ScatterD && !indices) { mask_.clear(); } + + // Initialize byte_pointer_ + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + + if (ScatterD) { + byte_pointer_ = + reinterpret_cast(pointer) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + } + + // store_byte_pointer_ is set to be the same with byte_pointer_ unless + // PermuteD is used. + store_byte_pointer_ = PermuteD ? reinterpret_cast(pointer) : byte_pointer_; + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + store_byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const + { + uint8_t* byte_pointer = store_byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * + LongIndex(params_.stride)); + } + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask_.predicates[column]; + + if (PermuteD) { + int col_offset = column * ThreadMap::Delta::kColumn; + + int col = col_offset + thread_start_column_; + int row = row_offset + thread_start_row_; + + // Locate memory_pointer + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + permute_layout_(PitchLinearCoord(col, row)) * sizeof(AccessType) / + kElementsPerAccess); + } + atomic_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[0], + guard); + + if (!PermuteD) { + memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD && !PermuteD) { byte_pointer += params_.increment_row; } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) const { store_with_byte_offset(frag, 0); } + + CUTLASS_DEVICE + void load(Fragment& frag) {} + + CUTLASS_DEVICE + MatrixCoord thread_start() const + { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_row() const { return thread_start_row_; } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_column() const { return thread_start_column_; } + + /// Extent of the matrix in rows + CUTLASS_DEVICE + Index extent_row() const { return extent_row_; } + + /// Extent of the matrix in columns + CUTLASS_DEVICE + Index extent_column() const { return extent_column_; } + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorAtomic& operator++() + { + ++state_[0]; + + if (!ScatterD && !PermuteD) { store_byte_pointer_ += params_.advance_row; } + + if (!ScatterD) { byte_pointer_ += params_.advance_row; } + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + store_byte_pointer_ += params_.advance_group; + + thread_start_row_ += + (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + store_byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + store_byte_pointer_ += params_.advance_tile; + + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * + ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile; + } + } + } + + return *this; + } + + /// Advances a number of positions to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorAtomic& operator+=(int increment) + { + // Row + state_[0] += increment; + int increment_row = state_[0] / ThreadMap::Count::kRow; + state_[0] = state_[0] % ThreadMap::Count::kRow; + + byte_pointer_ += (params_.advance_row * increment); + store_byte_pointer_ += (params_.advance_row * increment); + thread_start_row_ += (ThreadMap::Shape::kRow * increment); + + // Group + state_[1] += increment_row; + int increment_group = state_[1] / ThreadMap::Count::kGroup; + state_[1] = state_[1] % ThreadMap::Count::kGroup; + + byte_pointer_ += (params_.advance_group * increment_row); + store_byte_pointer_ += (params_.advance_group * increment_row); + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * + ThreadMap::Count::kRow * increment_row; + + // Cluster + state_[2] += increment_group; + int increment_cluster = state_[2] / ThreadMap::Count::kCluster; + state_[2] = state_[2] % ThreadMap::Count::kCluster; + + byte_pointer_ += (params_.advance_cluster * increment_group); + store_byte_pointer_ += (params_.advance_cluster * increment_group); + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow * increment_group; + + // Tile + byte_pointer_ += (params_.advance_tile * increment_cluster); + store_byte_pointer_ += (params_.advance_tile * increment_cluster); + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * + ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile * + increment_cluster; + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { mask_.clear(); } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { mask_.enable(); } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) const { mask = mask_; } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h new file mode 100644 index 0000000000000000000000000000000000000000..629047dbb057384ac8f0d7fa5d557e0c070cd830 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h @@ -0,0 +1,1938 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 + tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "cutlass/arch/memory.h" +#include "cutlass/transform/threadblock/predicated_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileIteratorResidualLast +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize +/// register liveness and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" +/// object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is +/// constructed. Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator +/// is constructed. Subsequent additions to logical coordinate offset may be +/// performed but are relatively expensive. +/// +/// Visitation order is intended to first visit a "residual" tile that may be +/// partially full in both the advance dimension and the steady-state dimension. +/// This is assumed to be the last tile in the iteration sequence. Advancing an +/// iterator that has just been constructed moves to the first tile that is full +/// in the advance dimension and recomputes predicates. Subsequent accesses may +/// be performed without updating internal predicates and are efficient in terms +/// of live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iterator will be dereferenced and advanced +/// at least once outside any looping structure to minimize integer arithmetic. +/// +/// Accesses out of bounds are safe so long as `clear_mask()` is called prior to +/// dereferencing the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update +// internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - +// subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to +// steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = +// transform::threadblock::PredicatedTileIteratorResidualLast; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template +class PredicatedTileIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = + AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout) {} + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : params_(base) {} + }; + +private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + +private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : address_iterator_(params.params_, pointer, extent, thread_id, threadblock_offset, indices) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { address_iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { address_iterator_.get_mask(mask); } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const* access_ptr = reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { *access_ptr = frag_ptr[idx]; } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize, + Gather>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()), + indices) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize, + Gather>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< Gather indices + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()), + indices) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank-2 data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = + AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout) {} + + CUTLASS_HOST_DEVICE + Params() {} + }; + +private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + +private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : address_iterator_(params.params_, pointer, extent, thread_id, threadblock_offset) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + if (kAdvanceRank) + address_iterator_.add_tile_offset(make_Coord(0, 1)); + else + address_iterator_.add_tile_offset(make_Coord(1, 0)); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { address_iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { address_iterator_.get_mask(mask); } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const* access_ptr = reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { *access_ptr = frag_ptr[idx]; } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 +/// column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) + { + } + }; + +private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 +/// row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) + { + } + }; + +private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for interleaved data. +/// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class PredicatedTileIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for interleaved-32 +/// data. It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h new file mode 100644 index 0000000000000000000000000000000000000000..2435c07f8989dcd6ad9c99a2fd530fe6092b1e08 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h @@ -0,0 +1,57 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "warp_iterator_from_smem.h" + +template +struct TransposeWarpIterator { + using Iterator = char; + static bool constexpr kSupportsTranspose = false; +}; + +template < + /// Operand identity + cutlass::gemm::Operand Operand, + /// Data type of A elements + typename Element, + bool kTranspose> +struct TransposeWarpIterator< + cutlass::gemm::warp::WarpIteratorFromSmem> { + using Iterator = cutlass::gemm::warp::WarpIteratorFromSmem; + static bool constexpr kSupportsTranspose = true; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h new file mode 100644 index 0000000000000000000000000000000000000000..7dd59832b4b0f278408552de38849b676d6b668a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Inspired from + "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" Loads tiles of GEMM + operands from a RowMajor shared-memory layout into registers to use by A100 + TensorCores. + + The difference with "mma_tensor_op_tile_access_iterator.h" is that: + (1) We use "ldmatrix" to load tiles, rather than manual loads (slightly + faster) (2) We support to transpose the operand (eg read `A.transpose()` when + the shared memory holds `A`) + + This is only implemented for the specific shapes. +*/ +#pragma once + +#include + +//////////////////////////////////////////////////////////////////////////////// +namespace cutlass { +namespace gemm { +namespace warp { + +template < + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + bool kTranspose = false> +class WarpIteratorFromSmem { +public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = cutlass::MatrixShape<32, 32>; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert( + kOperand == Operand::kA || kOperand == Operand::kB, + "WarpIteratorFromSmem may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + static_assert(sizeof_bits::value == 16, "Only supported for half"); + + /// Layout of source tile + using Layout = cutlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = cutlass::MatrixShape<16, 8>; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = 1; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 : 32 / sizeof_bits::value); + + using InstructionCount = MatrixShape; + + static int const kIterations = (kOperand == Operand::kA) ? InstructionCount::kColumn + : InstructionCount::kRow; + +public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + + /// Memory access type + // using AccessType = AlignedArray; + using AccessType = Array; + + static int constexpr kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? InstructionShape::kColumn : InstructionShape::kRow); + static int constexpr kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + static int const kTilesPerInstruction = InstructionShape::kRow / 8; + +private: + /// Underlying tensor reference + TensorRef ref_; + + /// Origin + MatrixCoord origin_; + + /// Iterations in a tile + int iterations_; + +public: + /// Constructor from TensorRef + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, int lane_id) + : WarpIteratorFromSmem(ref, {Shape::kRow, Shape::kColumn}, lane_id) + { + } + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, TensorCoord extent, int lane_id) + : ref_(ref), iterations_(0) + { + int ldsm_vec_num = (lane_id >> 3); + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id % 8, 0); + static_assert(InstructionCount::kRow * kAccessesInner * kTilesPerInstruction == 4, ""); + CUTLASS_PRAGMA_UNROLL + for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { + CUTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + CUTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < kTilesPerInstruction; + ++access_m_idx) { + int access_idx = + access_m_idx + + kTilesPerInstruction * (inner_idx + kAccessesInner * inst_m_idx); + + MatrixCoord offset(access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, + inner_idx * 4 * kElementsPerAccess); + + if (access_idx == ldsm_vec_num) { + if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } + origin_ += offset; + } + } + } + } + } else { + origin_ = MatrixCoord(0, lane_id % 8); + static_assert(InstructionCount::kColumn * kAccessesInner == 4, ""); + CUTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { + CUTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + + MatrixCoord offset(inner_idx * 4 * kElementsPerAccess, inst_n_idx * 8); + + if (access_idx == ldsm_vec_num) { + if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } + origin_ += offset; + } + } + } + } + + ref_.add_coord_offset(origin_); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& add_tile_offset(TensorCoord const& tile_offset) + { + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn); + if (kTranspose) { coord_offset = TensorCoord{coord_offset.column(), coord_offset.row()}; } + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_DEVICE + void advance() + { + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } else { + add_tile_offset({1, 0}); + } + + iterations_ = 0; + } + + /// increase iterations in a tile + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& operator++() + { + iterations_++; + + if (iterations_ >= kIterations) advance(); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + CUTLASS_DEVICE + void load(Fragment& frag) const + { + AccessType* access_ptr = reinterpret_cast(&frag); + using LoadLayout = + typename platform::conditional::type; + + MatrixCoord offset; + if (kOperand == Operand::kA) { + offset = MatrixCoord(0, iterations_ * InstructionShape::kColumn); + } else { + offset = MatrixCoord(iterations_ * InstructionShape::kRow, 0); + } + if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } + cutlass::arch::ldsm(access_ptr[0], ref_.data() + ref_.offset(offset)); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace cutlass +//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h new file mode 100644 index 0000000000000000000000000000000000000000..87e6df18bb04ad52278062f6bb03f40c36edc601 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h @@ -0,0 +1,1965 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/fast_math.h" +#include "cutlass/functional.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/vector.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" + +#include "gemm_kernel_utils.h" + +#include "cutlass/epilogue/thread/linear_combination_relu.h" +#include "cutlass/epilogue/threadblock/epilogue_smem_accumulator.h" +#include "cutlass/epilogue/warp/fragment_iterator_tensor_op.h" +#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h" +#include "cutlass/gemm/device/default_gemm_configuration.h" +#include "cutlass/gemm/kernel/default_gemm.h" +#include "cutlass/gemm/threadblock/default_mma.h" +#include "cutlass/gemm/threadblock/default_mma_core_simt.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/platform/platform.h" +#include "cutlass/transform/threadblock/predicated_tile_iterator.h" +#include "cutlass/transform/threadblock/vector_iterator.h" +#include "epilogue/epilogue_pipelined.h" +#include "iterators/epilogue_predicated_tile_iterator.h" + +#include "epilogue/epilogue_grad_bias.h" +#include "gemm/custom_mma.h" +#include "gemm/find_default_mma.h" +#include "gemm/mma_accum_lambda_iterator.h" +#include "gemm/mma_from_smem.h" +#include "transform/bias_broadcast.h" +#include "transform/tile_smem_loader.h" + +#include + +using namespace gemm_kernel_utils; + +namespace { + +template +struct GmemTile { + /* + Helper functions to efficient store/load RF to gmem + + GEMM accumulators have a particular format on A100, and + it takes some compute/shared-memory to rearrange them to + a RowMajor or ColumnMajor format in global memory through + an Epilogue. The same complexity goes for loading into RF. + + This class loads/stores RF as they are, and can be used for + efficient accumulation across gemms for instance: + + ``` + GmemTile tile; + for (int i = 0; i < N; ++i) { + // ... + + Fragment accum; + if (i == 0) { + accum.clear(); + } else { + tile.load(accum); + } + mma(accum, ...); + if (i < N-1) { + // Store for next GEMM + tile.store(accum); + } else { + // Store in tensor (eg RowMajor) + epilogue(accum); + } + + // ... + } + ``` + */ + + // 128bits per thread + using AccessType = cutlass::Array; + static constexpr int32_t kBytes = sizeof(AccessType); + static constexpr int32_t kStride = kNumThreads * AccessType::kElements; + static constexpr int32_t kNumIters = FragmentType::kElements / AccessType::kElements; + static constexpr int32_t kElementsStored = kNumThreads * FragmentType::kElements; + static_assert(FragmentType::kElements % AccessType::kElements == 0, + "fragment not aligned on 128 bits"); + + float* ptr; + + CUTLASS_DEVICE void load(FragmentType& fragment, int thread_id) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNumIters; ++i) { + AccessType* __restrict__ gmem_ptr = reinterpret_cast( + ptr + thread_id * AccessType::kElements + i * kStride); + AccessType sub_fragment; + cutlass::arch::global_load(sub_fragment, gmem_ptr, true); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < AccessType::kElements; ++j) { + fragment[i * AccessType::kElements + j] = sub_fragment[j]; + } + } + } + + CUTLASS_DEVICE void store(FragmentType const& fragment, int thread_id) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNumIters; ++i) { + AccessType* __restrict__ gmem_ptr = reinterpret_cast( + ptr + thread_id * AccessType::kElements + i * kStride); + AccessType sub_fragment; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < AccessType::kElements; ++j) { + sub_fragment[j] = fragment[i * AccessType::kElements + j]; + } + cutlass::arch::global_store(sub_fragment, gmem_ptr, true); + } + } +}; + +template +constexpr int getWarpsPerSm() +{ + constexpr bool is_half = !cutlass::platform::is_same::value; + if (Arch::kMinComputeCapability >= 80) { return is_half ? 12 : 8; } + return 8; +} +} // namespace + +template < + // which arch we target (eg `cutlass::arch::Sm80`) + typename ArchTag_, + // input/output type + typename scalar_t_, + // run optimized kernel because memory accesses will be aligned + bool kIsAligned_, + // use dropout if enabled + bool kApplyDropout_, + // when doing a GEMM, preload the next one (uses more shmem) + bool kPreload_, + // block dimensions + int kBlockSizeI_, + int kBlockSizeJ_, + // upperbound on `max(value.shape[-1], query.shape[-1])` + int kMaxK_ = (int)cutlass::platform::numeric_limits::max(), + template class Broadcast1_ = BroadcastNoLoad, + template class Broadcast2_ = BroadcastNoLoad> +struct AttentionBackwardKernel { + using scalar_t = scalar_t_; + using output_t = scalar_t; + using output_accum_t = float; + using lse_scalar_t = float; + using accum_t = float; + using ArchTag = ArchTag_; + static constexpr bool kIsAligned = kIsAligned_; + static constexpr bool kApplyDropout = kApplyDropout_; + static constexpr bool kPreload = kPreload_; + static constexpr int kBlockSizeI = kBlockSizeI_; + static constexpr int kBlockSizeJ = kBlockSizeJ_; + static constexpr int kMaxK = kMaxK_; + + struct Params { + // Input tensors + scalar_t* query_ptr; // [Mq, nH, K] + scalar_t* key_ptr; // [Mk, nH, K] + scalar_t* value_ptr; // [Mk, nH, Kv] + lse_scalar_t* logsumexp_ptr; // [nH, Mq] + scalar_t* output_ptr; // [Mq, nH, Kv] + scalar_t* grad_output_ptr; // [Mq, nH, Kv] + accum_t* delta_ptr; // [nH, Mq] + int32_t* cu_seqlens_q_ptr = nullptr; + int32_t* cu_seqlens_k_ptr = nullptr; + + // Output tensors + output_t* grad_query_ptr; // [Mq, nH, K] + output_t* grad_key_ptr; // [Mk, nH, K] + output_t* grad_value_ptr; // [Mk, nH, Kv] + + accum_t* grad_bias1_ptr = nullptr; + accum_t* grad_bias2_ptr = nullptr; + int32_t B = 0; + int32_t N = 0; + scalar_t* bias1_ptr = nullptr; + scalar_t* bias2_ptr = nullptr; + + // Accumulators + union { + output_accum_t* workspace = nullptr; // [Mq, Kq] + [Mkv, Kq] + [Mkv, Kv] + output_accum_t* workspace_gk; + }; + output_accum_t* workspace_gv; // (will be calculated by the kernel) + output_accum_t* workspace_gq; // (will be calculated by the kernel) + + // Scale + accum_t scale; + + // Dimensions/strides + int32_t head_dim = -1; + int32_t head_dim_value = -1; + int32_t num_queries = -1; + int32_t num_keys = -1; + int32_t num_heads = -1; + + int32_t q_strideM; + int32_t k_strideM; + int32_t v_strideM; + int32_t gO_strideM; + int32_t gB_strideM; + int8_t gQKV_strideM_multiplier = 1; // 3 for packed, 1 otherwise + + // RNG sequence offset based on batch_id and head_id + unsigned long long dropout_batch_head_rng_offset; + float dropout_prob = 0.0f; + + CUTLASS_HOST_DEVICE int32_t o_strideM() const { return head_dim_value * num_heads; } + CUTLASS_HOST_DEVICE int32_t gQ_strideM() const + { + return gQKV_strideM_multiplier * num_heads * head_dim; + } + CUTLASS_HOST_DEVICE int32_t gK_strideM() const + { + return gQKV_strideM_multiplier * num_heads * head_dim; + } + CUTLASS_HOST_DEVICE int32_t gV_strideM() const + { + return gQKV_strideM_multiplier * num_heads * head_dim_value; + } + + // Everything below is only used in `advance_to_block` + // and shouldn't use registers + int64_t o_strideH; + int32_t q_strideH; + int32_t k_strideH; + int32_t v_strideH; + int64_t o_strideB; + int64_t q_strideB; + int64_t k_strideB; + int64_t v_strideB; + int64_t lse_strideB; + int64_t lse_strideH; + int64_t delta_strideB; + int64_t delta_strideH; + int32_t num_batches; + + int64_t gO_strideB = 0; + int64_t gQ_strideB = 0; + int64_t gK_strideB = 0; + int64_t gV_strideB = 0; + int64_t gB_strideB = 0; + int64_t gO_strideH = 0; + int64_t gQ_strideH = 0; + int64_t gK_strideH = 0; + int64_t gV_strideH = 0; + int64_t gB_strideH = 0; + + CUTLASS_DEVICE bool advance_to_block() + { + int64_t batch_id = blockIdx.z; + int32_t head_id = blockIdx.y; + + if (kNeedsAccumGradQ || kNeedsAccumGradK || kNeedsAccumGradV) { + assert(workspace_size() == 0 || workspace != nullptr); + + workspace += (batch_id * num_heads + head_id) * workspace_strideBH(); + workspace = warp_uniform(workspace); + workspace_gv = workspace + workspace_elements_gk(); + workspace_gq = workspace_gv + workspace_elements_gv(); + } else { + workspace = nullptr; + } + + // Advance pointers that depend on the total concatenated + // number of queries, as `num_queries` is modified in the block + // below + dropout_batch_head_rng_offset = batch_id * (num_heads * num_queries * num_keys) + + head_id * (num_queries * num_keys); + logsumexp_ptr += batch_id * lse_strideB + head_id * lse_strideH; + + query_ptr += batch_id * q_strideB + head_id * q_strideH; + key_ptr += batch_id * k_strideB + head_id * k_strideH; + value_ptr += batch_id * v_strideB + head_id * v_strideH; + output_ptr += batch_id * o_strideB + head_id * o_strideH; + grad_output_ptr += batch_id * gO_strideB + head_id * gO_strideH; + delta_ptr += batch_id * delta_strideB + head_id * delta_strideH; + + grad_query_ptr += batch_id * gQ_strideB + head_id * gQ_strideH; + grad_key_ptr += batch_id * gK_strideB + head_id * gK_strideH; + grad_value_ptr += batch_id * gV_strideB + head_id * gV_strideH; + using broadcast_1 = Broadcast1_; + using broadcast_2 = Broadcast2_; + + if (broadcast_1::kEnable && grad_bias1_ptr) { + grad_bias1_ptr += batch_id * num_queries; + } + if (broadcast_2::kEnable && grad_bias2_ptr) { + auto strideB = num_heads * num_queries * num_keys; + auto strideH = num_queries * num_keys; + grad_bias2_ptr += (batch_id / N) * strideB + head_id * strideH; + } + if (broadcast_1::kEnable && bias1_ptr) { + bias1_ptr = broadcast_1::advance(bias1_ptr, + batch_id / N, + batch_id % N, + head_id, + num_queries * N, + num_queries, + 0); + } + if (broadcast_2::kEnable && bias2_ptr) { + auto strideB = num_heads * num_queries * num_keys; + auto strideH = num_queries * num_keys; + bias2_ptr = broadcast_2::advance( + bias2_ptr, batch_id / N, batch_id % N, head_id, strideB, 0, strideH); + } + + num_queries = warp_uniform(num_queries); + num_keys = warp_uniform(num_keys); + + query_ptr = warp_uniform(query_ptr); + key_ptr = warp_uniform(key_ptr); + value_ptr = warp_uniform(value_ptr); + logsumexp_ptr = warp_uniform(logsumexp_ptr); + output_ptr = warp_uniform(output_ptr); + grad_output_ptr = warp_uniform(grad_output_ptr); + delta_ptr = warp_uniform(delta_ptr); + + grad_query_ptr = warp_uniform(grad_query_ptr); + grad_key_ptr = warp_uniform(grad_key_ptr); + grad_value_ptr = warp_uniform(grad_value_ptr); + if (broadcast_1::kEnable) { + grad_bias1_ptr = warp_uniform(grad_bias1_ptr); + bias1_ptr = warp_uniform(bias1_ptr); + } + if (broadcast_2::kEnable) { + grad_bias2_ptr = warp_uniform(grad_bias2_ptr); + bias2_ptr = warp_uniform(bias2_ptr); + } + + return true; + } + + __host__ dim3 getBlocksGrid() const { return dim3(1, num_heads, num_batches); } + __host__ dim3 getThreadsGrid() const { return dim3(kWarpSize * kNumWarpsPerBlock, 1, 1); } + CUTLASS_HOST_DEVICE int64_t workspace_elements_gk() const + { + if (!kNeedsAccumGradK) { return 0; } + return align_up(num_keys, (int32_t)kBlockSizeJ) * + align_up(head_dim, (int32_t)kBlockSizeI); + } + CUTLASS_HOST_DEVICE int64_t workspace_elements_gv() const + { + if (!kNeedsAccumGradV) { return 0; } + return align_up(num_keys, (int32_t)kBlockSizeJ) * + align_up(head_dim_value, (int32_t)kBlockSizeI); + } + CUTLASS_HOST_DEVICE int64_t workspace_elements_gq() const + { + if (!kNeedsAccumGradQ) { return 0; } + if (num_keys <= kBlockSizeJ) { return 0; } + return align_up(num_queries, (int32_t)kBlockSizeI) * + align_up(head_dim, (int32_t)kBlockSizeJ); + } + CUTLASS_HOST_DEVICE int64_t workspace_strideBH() const + { + // Aligned on 128bits + return align_up( + workspace_elements_gk() + workspace_elements_gv() + workspace_elements_gq(), + int64_t(4)); + } + CUTLASS_HOST_DEVICE int64_t workspace_size() const + { + // Returns size of buffer we need to run this kernel + return num_batches * num_heads * workspace_strideBH() * sizeof(float); + } + }; + + static constexpr int64_t kWarpSize = 32; + + // If this is true, we store and accumulate dK/dV in RF + // rather than going back to gmem every time + static constexpr bool kIsHalf = cutlass::sizeof_bits::value <= 16; + static constexpr bool kOutputInRF = kIsHalf && kMaxK <= kBlockSizeI; + static_assert(!kPreload || (kIsHalf && ArchTag::kMinComputeCapability >= 80 && kOutputInRF), + "preload MMA not supported"); + static constexpr bool kPrologueQK = kPreload; + static constexpr bool kPrologueGV = kPreload; + static constexpr bool kPrologueDOV = kPreload; + static constexpr bool kPrologueGQ = kPreload; + static constexpr bool kPrologueGK = kPreload; + + static constexpr int64_t kNumWarpsPerBlock = (kBlockSizeI * kBlockSizeJ) / (32 * 32); + + // Compute delta for the f16 kernels + // TODO: Figure out why it's slower on the f32 kernels + // (something due to RF pressure?) + // TODO: Remove condition on `kOutputInRF` - this is needed to work + // around a compiler bug on V100, not exactly sure why but I spent + // too much time on this already. Reproducible with + // (B, Mq, Mkv, K) = (1, 1, 1, 136) for instance + static constexpr bool kKernelComputesDelta = + kIsHalf && (kOutputInRF || ArchTag::kMinComputeCapability != 70); + + static constexpr bool kNeedsAccumGradQ = + !cutlass::platform::is_same::value; + static constexpr bool kNeedsAccumGradK = + !kOutputInRF && !cutlass::platform::is_same::value; + static constexpr bool kNeedsAccumGradV = + !kOutputInRF && !cutlass::platform::is_same::value; + + // Launch bounds + static constexpr int64_t kNumThreads = kWarpSize * kNumWarpsPerBlock; + static constexpr int64_t kMinBlocksPerSm = + getWarpsPerSm() / kNumWarpsPerBlock; + + using GemmType = DefaultGemmType; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration; + static constexpr auto kOptimalAlignement = + cutlass::platform::max(DefaultConfig::kAlignmentA, DefaultConfig::kAlignmentB); + static constexpr auto kMinimumAlignment = GemmType::kMinimumAlignment; + + struct MatmulQK { + /* + attn_T = k_j @ q_i.transpose(-2, -1) # matmul + attn_T = (attn_T - logsumexp[i_start:i_end].unsqueeze(1).transpose(-2, + -1)).exp() # epilogue + + with attn_T.shape = (kBlockSizeJ, kBlockSizeI) + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using DefaultMma = typename cutlass::gemm::threadblock::DefaultMma< + scalar_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, + scalar_t, // ElementB + cutlass::layout::ColumnMajor, // LayoutB + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + accum_t, // ElementC + cutlass::layout::RowMajor, // LayoutC + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + DefaultConfig::kStages, + typename GemmType::Operator, + false, // AccumulatorsInRowMajor = false, + cutlass::gemm::SharedMemoryClearOption::kNone>; + using MmaCore = typename DefaultMma::MmaCore; + using Mma = typename MakeCustomMma::Mma; + + // used for efficient load of bias tile (Bij) from global memory to shared + // memory + using BiasLoader = + TileSmemLoader, + MmaCore::kThreads, + // input restriction: kv_len has to be a multiple of this value + 128 / cutlass::sizeof_bits::value>; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = + typename cutlass::gemm::threadblock::B2bGemm; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator::Iterator; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MatmulGradV { + /* + grad_v[j_start:j_end] += attn_T @ do_i # matmul + + Dimensions: (kBlockSizeJ * kNumWarpsPerBlock, kBlockSizeI, K) + (we might need to iterate multiple times on K) + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = + cutlass::gemm::kernel::DefaultGemm; + + // if dropout: + // for computing dVj += (Pij.T * Zij) @ dOi + // Pij_dropped.T = Pij.T * Zij is computed on the fly as fragments of + // Pij.T are loaded in. The reason we do it this way is because Pij.T and + // Zij are reused in later steps, while Pij_dropped.T is only needed in + // this step. computing Pij_dropped.T on the fly allows us to avoid + // keeping all 3 of Pij_dropped.T, Pij.T, and Zij in shared memory at the + // same time. + // if no dropout: + // for computing dVj += Pij.T @ dOi + using DefaultMmaFromSmem = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + typename MatmulQK::AccumulatorSharedStorage, + kApplyDropout>; // kScaleOperandA + + using Mma = typename DefaultMmaFromSmem::Mma; + using WarpIteratorA = typename DefaultMmaFromSmem::WarpIteratorA; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + struct MatmulDOIVJ { + /* + doi_t_vj = do_i @ v_j.transpose(-2, -1) # matmul + tmp = (doi_t_vj - Di.unsqueeze(1)) * attn # inplace / epilogue? + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + + using ElementC = accum_t; // CSY: Change it for better accuracy + using ElementAccum = accum_t; + + // no-op output op - epilogue just stores result to global memory + using BiasGradEpilogueOutputOp = typename cutlass::epilogue::thread::LinearCombination< + ElementC, + DefaultConfig::EpilogueOutputOp::kCount, + typename DefaultConfig::EpilogueOutputOp::ElementAccumulator, + typename DefaultConfig::EpilogueOutputOp::ElementCompute, + cutlass::epilogue::thread::ScaleType::Nothing>; + + using DefaultGemm = typename cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, + scalar_t, // ElementB + cutlass::layout::ColumnMajor, // LayoutB + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + ElementC, // ElementC + cutlass::layout::RowMajor, // LayoutC + ElementAccum, // ElementAccumulator + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + BiasGradEpilogueOutputOp, // EpilogueOutputOp + void, // ThreadblockSwizzle (not used) + // multiple preloads, dropout Zij tile, and 3 stages push us over shared + // memory capacity on A100. set a ceiling on number of stages to save + // shared memory if dropout is in use. + kPreload && kApplyDropout && (kBlockSizeI * kBlockSizeJ > 64 * 64) + ? cutlass::const_min(2, DefaultConfig::kStages) + : DefaultConfig::kStages, // Stages + false, // SplitKSerial + typename GemmType::Operator, + cutlass::gemm::SharedMemoryClearOption::kNone>; + using Mma = typename MakeCustomMma::Mma; + + // epilogue used to write bias gradient, which is just the output of this + // matmul with some operations applied to the fragment + using BiasGradEpilogue = typename DefaultGemm::Epilogue; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = + typename cutlass::gemm::threadblock::B2bGemm; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MatmulGradQ { + // grad_q <- tmp @ k_j + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = + cutlass::gemm::kernel::DefaultGemm; + + using DefaultMmaFromSmem = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + typename MatmulDOIVJ::AccumulatorSharedStorage, + false>; // kScaleOperandA + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + struct MatmulGradK { + // grad_k <- tmp.transpose(-2, -1) @ q_i + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = + cutlass::gemm::kernel::DefaultGemm; + + using DefaultMmaFromSmemN = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + typename MatmulQK::AccumulatorSharedStorage, + false>; // kScaleOperandA + using DefaultMmaFromSmemT = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + typename MatmulDOIVJ::AccumulatorSharedStorage, + false, // kScaleOperandA + kPreload>; // kTransposeA + using DefaultMmaFromSmem = + typename cutlass::platform::conditional::type; + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + using broadcast_1 = Broadcast1_; + using broadcast_2 = Broadcast2_; + + // shared storage for keeping Zij matrix. not needed if we aren't using + // dropout, in which case we use an empty array to save shared memory + using ZijSharedStorage = typename cutlass::platform::conditional< + kApplyDropout, + typename MatmulQK::AccumulatorSharedStorage, + // dummy shared storage object that takes up no space. + typename cutlass::gemm::threadblock::AccumulatorSharedStorage< +#ifdef _WIN32 + // windows builds throw the error: + // "type containing an unknown-size array is not allowed" + // if we try to make Zij shared storage zero-sized. + // To get around this just make it sized 1 on windows. + typename cutlass::gemm::GemmShape<1, 1, 0>, +#else + typename cutlass::gemm::GemmShape<0, 0, 0>, +#endif + typename MatmulQK::AccumulatorSharedStorage::Element, + typename MatmulQK::AccumulatorSharedStorage::Layout, + typename cutlass::MatrixShape<0, 0>>>::type; + + struct SharedStoragePrologue { + struct { + cutlass::Array di; // (do_i * o_i).sum(-1) + typename MatmulQK::Mma::SharedStorageA mm_qk_k; + } persistent; + union { + struct { + // part1 - after Q.K / dV / dO.V + union { + // 1. efficient load of bias tile Bij, which is then applied to Pij + // typename MatmulQK::BiasLoader::SmemTile bias; + cutlass::AlignedBuffer bias; + // 4. store Pij. it is needed: + // - in dVj += (Pij.T * Zij) @ dOi + // - in dSij = Pij * (dPij - Di) + // 6. dVj += (Pij.T * Zij) @ dOi + // 10. write to fragment + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + }; + // 5. store Zij. it is needed: + // - to compute Pij_dropped = Pij * Zij on the fly as fragments of Pij + // are loaded for the computation of dVj. + // - to compute dPij = (dOi @ Vj.T) * Zij + // 6. used in dVj += (Pij.T * Zij) @ dOi + // 9. used in dPij = dPij_dropped * Zij + ZijSharedStorage zij; + + union { + // 2. prologue for dVj + // 6. workspace for dVj += (Pij.T * Zij) @ dOi + typename MatmulGradV::Mma::SharedStorage mm_gradV; + // 7. dVj epilogue + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; + }; + + // 3. prologue for dPij_dropped + // 8. used in dPij_dropped = dOi @ Vj.T + typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; + } part1; + + struct { + // part2 - dQ + union { + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part1) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + }; + typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) + typename MatmulGradQ::Mma::SharedStorage mm_gradQ; // (preload) + union { + // store dB = dSij to global memory + typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; + }; + + } part2; + + struct { + // part3 - after last iteration on dQ's epilogue / dK + union { + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part1) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + }; + typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue_lastIter; + + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; + } part3; + + struct { + // part4 - after last iteration on dK's epilogue / preload next K.Q_t + typename MatmulQK::Mma::SharedStorageB mm_qk_q; + + // If we reach end of current key, dump RF->gmem with "final" epilogues + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue_final; + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue_final; + } part4; + }; +// =========================================== +#define FIELD(INSIDE_STRUCT, FIELDNAME) \ + CUTLASS_DEVICE auto& FIELDNAME() { return INSIDE_STRUCT.FIELDNAME; } + + FIELD(persistent, di) + FIELD(persistent, mm_qk_k) + FIELD(part1, bias) + FIELD(part1, attn_shared_storage) + FIELD(part1, zij) + FIELD(part1, mm_gradV) + FIELD(part1, gradV_epilogue) + FIELD(part1, mm_doivj) + FIELD(part2, mm_gradK) + FIELD(part2, mm_gradQ) + FIELD(part2, gradB_epilogue) + FIELD(part2, gradQ_epilogue) + FIELD(part2, tmp_shared_storage) + FIELD(part3, tmpT_shared_storage) + FIELD(part3, gradQ_epilogue_lastIter) + FIELD(part3, gradK_epilogue) + FIELD(part4, mm_qk_q) + FIELD(part4, gradK_epilogue_final) + FIELD(part4, gradV_epilogue_final) + }; + + struct SharedStorageNoPrologue { + struct { + cutlass::Array di; // (do_i * o_i).sum(-1) + } persistent; + union { + struct { + // part1 - Q.K matmul + typename MatmulQK::Mma::SharedStorageA mm_qk_k; + typename MatmulQK::Mma::SharedStorageB mm_qk_q; + } part1; + + struct { + // part2 - compute gradV + union { + // 1. efficient load of bias tile Bij, which is then applied to Pij + cutlass::AlignedBuffer bias; + // 2. store Pij to shared memory. it is needed: + // - in this step, where it is used in dVj += (Pij.T * Zij) @ dOi + // - in next step where it is used in dSij = Pij * (dPij - Di) + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + }; + // 3. store Zij. it is needed: + // - in this step, where it is used to compute Pij_dropped = Pij * Zij + // on the + // fly as fragments of Pij are loaded for the computation of dVj. + // - later to compute dPij = (dOi @ Vj.T) * Zij + ZijSharedStorage zij; + + union { + typename MatmulGradV::Mma::SharedStorage mm_gradV; + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; + }; + } part2; + + struct { + // part3 - DO.V matmul + union { + // first compute dPij = (dOi @ Vj.T) * Zij + // and dSij = Pij * (dPij - Di) + struct { + // (from part2) - Pij for computing dSij = Pij * (dPij - Di) + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + // (from part2) - Zij for computing dPij = dPij_dropped * Zij + ZijSharedStorage zij; + // matmul to compute dOiVj + typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; + }; + // then store dB = dSij to global memory + typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; + }; + } part3; + + struct { + // part4 - compute gradQ + typename MatmulQK::AccumulatorSharedStorage tmpT_shared_storage; // (from part2) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + union { + typename MatmulGradQ::Mma::SharedStorage mm_gradQ; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue_lastIter; + }; + } part4; + + struct { + // part5 - compute gradK + typename MatmulQK::AccumulatorSharedStorage tmpT_shared_storage; // (from part2) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + union { + typename MatmulGradK::Mma::SharedStorage mm_gradK; + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; + }; + } part5; + + struct { + // part6 - store RF accumulated into gmem + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue_final; + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue_final; + } part6; + }; +// =========================================== +#define FIELD(INSIDE_STRUCT, FIELDNAME) \ + CUTLASS_DEVICE auto& FIELDNAME() { return INSIDE_STRUCT.FIELDNAME; } + + FIELD(persistent, di) + FIELD(part1, mm_qk_k) + FIELD(part1, mm_qk_q) + FIELD(part2, bias) + FIELD(part2, attn_shared_storage) + FIELD(part2, zij) + FIELD(part2, mm_gradV) + FIELD(part2, gradV_epilogue) + FIELD(part3, mm_doivj) + FIELD(part3, gradB_epilogue) + FIELD(part4, tmpT_shared_storage) + FIELD(part4, tmp_shared_storage) + FIELD(part4, mm_gradQ) + FIELD(part4, gradQ_epilogue) + FIELD(part4, gradQ_epilogue_lastIter) + FIELD(part5, mm_gradK) + FIELD(part5, gradK_epilogue) + FIELD(part6, gradK_epilogue_final) + FIELD(part6, gradV_epilogue_final) + }; + + using SharedStorage = typename cutlass::platform:: + conditional::type; + + struct OutputFragments { + typename MatmulGradV::Mma::FragmentC gradV; + typename MatmulGradK::Mma::FragmentC gradK; + + CUTLASS_DEVICE void clear() + { + gradV.clear(); + gradK.clear(); + } + }; + + static bool __host__ check_supported(Params const& p) + { + CHECK_ALIGNED_PTR(p.query_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.key_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.value_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.output_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.grad_output_ptr, kMinimumAlignment); + EVOFORMER_CHECK(p.lse_strideH % 8 == 0, "LSE is not correctly aligned"); + EVOFORMER_CHECK(p.lse_strideB % 8 == 0, "LSE is not correctly aligned"); + EVOFORMER_CHECK(p.num_heads <= 1 || p.q_strideH % kMinimumAlignment == 0, + "query is not correctly aligned (strideH)"); + EVOFORMER_CHECK(p.num_heads <= 1 || p.k_strideH % kMinimumAlignment == 0, + "key is not correctly aligned (strideH)"); + EVOFORMER_CHECK(p.num_heads <= 1 || p.v_strideH % kMinimumAlignment == 0, + "value is not correctly aligned (strideH)"); + EVOFORMER_CHECK(p.num_batches <= 1 || p.q_strideB % kMinimumAlignment == 0, + "query is not correctly aligned (strideB)"); + EVOFORMER_CHECK(p.num_batches <= 1 || p.k_strideB % kMinimumAlignment == 0, + "key is not correctly aligned (strideB)"); + EVOFORMER_CHECK(p.num_batches <= 1 || p.v_strideB % kMinimumAlignment == 0, + "value is not correctly aligned (strideB)"); + EVOFORMER_CHECK(p.q_strideM % kMinimumAlignment == 0, + "query is not correctly aligned (strideM)"); + EVOFORMER_CHECK(p.k_strideM % kMinimumAlignment == 0, + "key is not correctly aligned (strideM)"); + EVOFORMER_CHECK(p.v_strideM % kMinimumAlignment == 0, + "value is not correctly aligned (strideM)"); + EVOFORMER_CHECK(p.dropout_prob <= 1.0f && p.dropout_prob >= 0.0f, + "Invalid value for `dropout_prob`"); + EVOFORMER_CHECK(kApplyDropout || p.dropout_prob == 0.0f, + "Set `kApplyDropout`=True to support `dropout_prob > 0`"); + EVOFORMER_CHECK(p.head_dim > 0, "Invalid value for `head_dim`"); + EVOFORMER_CHECK(p.head_dim_value > 0, "Invalid value for `head_dim_value`"); + EVOFORMER_CHECK(p.num_queries > 0, "Invalid value for `num_queries`"); + EVOFORMER_CHECK(p.num_keys > 0, "Invalid value for `num_keys`"); + EVOFORMER_CHECK(p.num_heads > 0, "Invalid value for `num_heads`"); + EVOFORMER_CHECK(p.num_batches > 0, "Invalid value for `num_batches`"); + EVOFORMER_CHECK(p.head_dim <= kMaxK, "kMaxK: Expected `head_dim < kMaxK`"); + EVOFORMER_CHECK(p.head_dim_value <= kMaxK, "kMaxK: Expected `head_dim_value < kMaxK`"); + return true; + } + + static CUTLASS_DEVICE void attention_kernel(Params p) + { + extern __shared__ char smem_buffer[]; + SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); + + uint16_t thread_id = threadIdx.x; + uint8_t warp_id = warp_uniform(thread_id / 32); + uint8_t lane_id = thread_id % 32; + + if (kPrologueQK) { + prologueQkNextIteration(shared_storage, p, 0, 0, warp_id, lane_id); + } + + // Computes (dO*out).sum(-1) and writes it to `p.delta_ptr` + if (kKernelComputesDelta) { + constexpr int kOptimalElements = 128 / cutlass::sizeof_bits::value; + if (p.head_dim_value % kOptimalElements == 0) { + for (int query_start = 0; query_start < p.num_queries; query_start += kBlockSizeI) { + computeDelta(p, query_start, warp_id, lane_id); + } + } else { + for (int query_start = 0; query_start < p.num_queries; query_start += kBlockSizeI) { + computeDelta<1>(p, query_start, warp_id, lane_id); + } + } + __syncthreads(); + } + + OutputFragments output_frags; + + int32_t key_start = 0; + int32_t key_end = p.num_keys / kBlockSizeJ * kBlockSizeJ; + for (; key_start < key_end; key_start += kBlockSizeJ) { + output_frags.clear(); + int32_t query_start = getQueryStart(p, key_start); + int32_t query_end = + query_start + (p.num_queries - query_start) / kBlockSizeI * kBlockSizeI; + for (; query_start < query_end; query_start += kBlockSizeI) { + processBlockIJ( + shared_storage, output_frags, p, query_start, key_start, warp_id, lane_id); + } + // last (partial) query + if (query_start < p.num_queries) { + processBlockIJ( + shared_storage, output_frags, p, query_start, key_start, warp_id, lane_id); + } + if (kOutputInRF) { + writeFragsToGmem( + shared_storage, output_frags, p, key_start, warp_id, lane_id); + } else if (getQueryStart(p, key_start) >= p.num_queries) { + zfillGradKV(p, key_start, warp_id, lane_id); + } + __syncthreads(); + } + // Last (partial) key + if (key_start != p.num_keys) { + output_frags.clear(); + int32_t query_start = getQueryStart(p, key_start); + for (; query_start < p.num_queries; query_start += kBlockSizeI) { + warp_id = warp_uniform(warp_id); + processBlockIJ( + shared_storage, output_frags, p, query_start, key_start, warp_id, lane_id); + } + if (kOutputInRF) { + writeFragsToGmem( + shared_storage, output_frags, p, key_start, warp_id, lane_id); + } else if (getQueryStart(p, key_start) >= p.num_queries) { + zfillGradKV(p, key_start, warp_id, lane_id); + } + } + } + + static CUTLASS_DEVICE void loadDi(cutlass::Array& di, + Params const& p, + int32_t query_start) + { + int32_t thread_id = threadIdx.x + threadIdx.y * blockDim.x; + if (thread_id < kBlockSizeI) { + accum_t di_rf = accum_t(0); + if (query_start + thread_id < p.num_queries) { + di_rf = p.delta_ptr[query_start + thread_id]; + } + di[thread_id] = di_rf; + } + } + + template + static CUTLASS_DEVICE void zfillGradKV(Params const& p, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) + { + constexpr int kThreadsPerKey = 8; + constexpr int kParallelKeys = kNumThreads / kThreadsPerKey; + static_assert(kBlockSizeJ % kParallelKeys == 0, ""); + // This function is not really optimized, but should rarely be used + // It's only used when some keys are "useless" and don't attend to + // any query, due to causal masking + int thread_id = 32 * warp_id + lane_id; + int k_shift = lane_id % kThreadsPerKey; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < kBlockSizeJ; j += kParallelKeys) { + int key = key_start + j + (thread_id / kThreadsPerKey); + if (!skipBoundsChecks && key >= p.num_keys) { continue; } + auto gv_ptr = p.grad_value_ptr + key * p.gV_strideM(); + auto gk_ptr = p.grad_key_ptr + key * p.gK_strideM(); + + for (int k = k_shift; k < p.head_dim_value; k += kThreadsPerKey) { + gv_ptr[k] = scalar_t(0); + } + for (int k = k_shift; k < p.head_dim; k += kThreadsPerKey) { gk_ptr[k] = scalar_t(0); } + } + } + + template + static CUTLASS_DEVICE void processBlockIJ(SharedStorage& shared_storage, + OutputFragments& output_frags, + Params& p, + int32_t query_start, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) + { + cutlass::MatrixCoord no_offset{0, 0}; + accum_t scale = p.scale; + int16_t thread_id = 32 * warp_id + lane_id; + auto rematerializeThreadIds = [&]() { + // Prevents `nvcc` from keeping values deduced from + // `thread_id`, `warp_id`, ... in RF - to reduce register pressure + warp_id = warp_uniform(thread_id / 32); + lane_id = thread_id % 32; + thread_id = 32 * warp_id + lane_id; + }; + + bool isFirstQuery = (query_start == getQueryStart(p, key_start)); + int32_t next_query, next_key; + incrIteration(p, query_start, key_start, next_query, next_key); + bool isLastQuery = next_key != key_start; + __syncthreads(); + loadDi(shared_storage.di(), p, query_start); + + int32_t num_queries_in_block = + skipBoundsChecks ? MatmulQK::Mma::Shape::kN + : warp_uniform(cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kN, + p.num_queries - query_start)); + int32_t num_keys_in_block = + skipBoundsChecks ? MatmulQK::Mma::Shape::kM + : warp_uniform(cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, + p.num_keys - key_start)); + + auto prologueGradV = [&](int col) { + typename MatmulGradV::Mma::IteratorB iterator_dO( + {int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM + col, + {num_queries_in_block, p.head_dim_value - col}, + thread_id, + no_offset); + MatmulGradV::Mma::prologue( + shared_storage.mm_gradV(), iterator_dO, thread_id, num_queries_in_block); + }; + auto prologueGradQ = [&](int col) { + typename MatmulGradQ::Mma::IteratorB iterator_K( + {int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM + col, + {num_keys_in_block, p.head_dim - col}, + thread_id, + no_offset); + MatmulGradQ::Mma::prologue( + shared_storage.mm_gradQ(), iterator_K, thread_id, num_keys_in_block); + }; + auto prologueGradK = [&](int col) { + typename MatmulGradK::Mma::IteratorB iterator_Q( + {int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM + col, + {num_queries_in_block, p.head_dim - col}, + thread_id, + no_offset); + MatmulGradK::Mma::prologue( + shared_storage.mm_gradK(), iterator_Q, thread_id, num_queries_in_block); + }; + auto prologueDOV = [&]() { + typename MatmulDOIVJ::Mma::IteratorA iterator_A( + {int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM, + {num_queries_in_block, p.head_dim_value}, + thread_id, + no_offset); + typename MatmulDOIVJ::Mma::IteratorB iterator_B({int32_t(p.v_strideM)}, + p.value_ptr + key_start * p.v_strideM, + {p.head_dim_value, num_keys_in_block}, + thread_id, + no_offset); + MatmulDOIVJ::Mma::prologue( + shared_storage.mm_doivj(), iterator_A, iterator_B, thread_id, p.head_dim_value); + }; + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // MatmulQK + ///////////////////////////////////////////////////////////////////////////////////////////////// + { + using Mma = typename MatmulQK::Mma; + + cutlass::gemm::GemmCoord problem_size(num_keys_in_block, + num_queries_in_block, + p.head_dim // k + ); + + // k_j + typename Mma::IteratorA iterator_A({int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM, + {problem_size.m(), problem_size.k()}, + thread_id, + no_offset); + + // q_i.transpose(-2, -1) + typename Mma::IteratorB iterator_B({int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM, + {problem_size.k(), problem_size.n()}, + thread_id, + no_offset); + + Mma mma( + shared_storage.mm_qk_k(), shared_storage.mm_qk_q(), thread_id, warp_id, lane_id); + + typename Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma.set_prologue_done(kPrologueQK); + mma.set_zero_outside_bounds(!skipBoundsChecks); + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + + // Epilogue: add LSE + exp and store that to our shared memory buffer + // shmem <- (matmul_result - + // logsumexp[i_start:i_end].unsqueeze(1)).exp() + int warp_idx_mn_0 = warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % Mma::Base::WarpCount::kM, warp_idx_mn_0 / Mma::Base::WarpCount::kM}; + + if (broadcast_1::kEnable || broadcast_2::kEnable) { + cutlass::TensorRef bias_tensor_ref( + shared_storage.bias().data(), + cutlass::layout::RowMajor(MatmulQK::ThreadblockShape::kM)); + using Shape = cutlass::MatrixShape; + AttentionBiasEpilogue + bias_epilogue; + bias_epilogue(bias_tensor_ref, + p.bias1_ptr + key_start, + p.bias2_ptr + query_start * p.num_keys + key_start, + thread_id, + {num_queries_in_block, num_keys_in_block}, + p.num_keys); + // Pij += Bij, Pij is in register fragment and Bij is in shared memory + auto lane_offset = MatmulQK::AccumLambdaIterator::get_lane_offset( + lane_id, warp_id, output_tile_coords); + MatmulQK::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_n) {}, + [&](int accum_m, int accum_n, int idx) { + // remember we are transposed + accum[idx] = accum[idx] * scale + bias_tensor_ref.at({accum_n, accum_m}); + }, + [&](int accum_n) {}); + } else { + accum = cutlass::multiplies()(scale, accum); + } + + __syncthreads(); + if (kPrologueGV) { prologueGradV(0); } + if (kPrologueDOV) { prologueDOV(); } + + MatmulQK::B2bGemm::accumApplyLSEToSmem(shared_storage.attn_shared_storage(), + accum, + p.logsumexp_ptr + query_start, + problem_size.n(), + thread_id, + warp_id, + lane_id, + output_tile_coords); + + __syncthreads(); + } + rematerializeThreadIds(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradV matmul + // + // grad_v[j_start:j_end] += attn_T @ do_i + ///////////////////////////////////////////////////////////////////////////////////////////////// + constexpr bool kSingleIterationGradV = kMaxK <= MatmulGradV::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradV ? 1 : p.head_dim_value); + col += MatmulGradV::ThreadblockShape::kN) { + using Mma = typename MatmulGradV::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + cutlass::gemm::GemmCoord problem_size( + num_keys_in_block, p.head_dim_value - col, num_queries_in_block); + auto createEpilogueIter = [&]() { + return typename MatmulGradV::OutputTileIterator( + typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, + p.grad_value_ptr + key_start * p.gV_strideM() + col, + {num_keys_in_block, p.head_dim_value - col}, + thread_id); + }; + typename Mma::IteratorB iterator_B({int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM + col, + {num_queries_in_block, p.head_dim_value - col}, + thread_id, + no_offset); + + // if dropout: dVj += (Pij.T * Zij) @ dOi + // otherwise: dVj += Pij.T @ dOi + Mma mma(shared_storage.mm_gradV(), + // operand A: Pij + typename MatmulGradV::WarpIteratorA( + shared_storage.attn_shared_storage().accum_ref(), lane_id), + // if we're using dropout, operand A is Pij_dropped = Pij * Zij + // which is computed on the fly as fragments of Pij are loaded in + typename Mma::WarpIteratorAScale(shared_storage.zij().accum_ref(), lane_id), + thread_id, + warp_id, + lane_id); + + int storage_id = col / MatmulGradV::ThreadblockShape::kN; + AccumTileGmem gmem_tile{p.workspace_gv + storage_id * AccumTileGmem::kElementsStored}; + if (!kOutputInRF) { + if (isFirstQuery || !kNeedsAccumGradV) { + output_frags.gradV.clear(); + } else { + gmem_tile.load(output_frags.gradV, thread_id); + } + } + mma.set_prologue_done(kPrologueGV); + + auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + + mma(gemm_k_iterations, output_frags.gradV, iterator_B, output_frags.gradV); + __syncthreads(); + if (kPrologueGV && !kSingleIterationGradV && + col + MatmulGradV::ThreadblockShape::kN < p.head_dim_value) { + prologueGradV(col + MatmulGradV::ThreadblockShape::kN); + } + + if (!kOutputInRF) { + if (kNeedsAccumGradV && !isLastQuery) { + gmem_tile.store(output_frags.gradV, thread_id); + } else { + accumulateInGmem(shared_storage.gradV_epilogue(), + output_frags.gradV, + createEpilogueIter(), + isFirstQuery || kNeedsAccumGradV, + warp_id, + lane_id); + } + } + } + __syncthreads(); + ///////////////////////////////////////////////////////////////////////////////////////////////// + // MatmulDOIVJ + ///////////////////////////////////////////////////////////////////////////////////////////////// + { + using Mma = typename MatmulDOIVJ::Mma; + // do_i + typename Mma::IteratorA iterator_A({int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM, + {num_queries_in_block, p.head_dim_value}, + thread_id, + no_offset); + + // v_j.transpose(-2, -1) + typename Mma::IteratorB iterator_B({int32_t(p.v_strideM)}, + p.value_ptr + key_start * p.v_strideM, + {p.head_dim_value, num_keys_in_block}, + thread_id, + no_offset); + + Mma mma(shared_storage.mm_doivj(), thread_id, warp_id, lane_id); + mma.set_prologue_done(kPrologueDOV); + mma.set_zero_outside_bounds(!skipBoundsChecks); + + typename Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = (p.head_dim_value + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + __syncthreads(); + if (kPrologueGQ) { prologueGradQ(0); } + if (kPrologueGK) { prologueGradK(0); } + + int warp_idx_mn_0 = warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % Mma::Base::WarpCount::kM, warp_idx_mn_0 / Mma::Base::WarpCount::kM}; + // TODO: This must be terribly inefficient. There must be a better way + // tmp [RF] <- (accum [RF] - Di [smem] ) * attn_T.T [smem] + // attn_shared_storage [smem] <- tmp.T + // tmp_shared_storage [smem] <- tmp + { + using LambdaIterator = + typename DefaultMmaAccumLambdaIterator::Iterator; + auto lane_offset = + LambdaIterator::get_lane_offset(lane_id, warp_id, output_tile_coords); + + auto attn_T = shared_storage.attn_shared_storage().accum_ref(); + accum_t current_di; + // dSij = (dPij - Di) * Pij + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { current_di = shared_storage.di()[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + if (skipBoundsChecks || + (accum_m < num_queries_in_block && accum_n < num_keys_in_block)) { + accum_t attn = attn_T.at({accum_n, accum_m}); + accum[idx] = (accum[idx] - current_di) * attn; + } else { + accum[idx] = 0; + } + }, + [&](int accum_m) { + + }); + + using DefaultGemm = typename MatmulDOIVJ::DefaultGemm; + using OutputOp = typename MatmulDOIVJ::BiasGradEpilogueOutputOp; + if (broadcast_1::kEnable && p.grad_bias1_ptr) { + using Epilogue = + typename BiasGradEpilogueAffineRankN::Epilogue; + cutlass::layout::AffineRankN<2> layout({0, 1}); + auto dst_ptr = p.grad_bias1_ptr + key_start; + typename Epilogue::OutputTileIterator output_iter( + {layout}, + dst_ptr, + {num_queries_in_block, num_keys_in_block}, + (int)thread_id); + Epilogue epilogue(shared_storage.gradB_epilogue(), + (int)thread_id, + (int)warp_id, + (int)lane_id); + epilogue(OutputOp(1), output_iter, accum); + } + + if (broadcast_2::kEnable && p.grad_bias2_ptr) { + if (broadcast_1::kEnable) { __syncthreads(); } + using Epilogue = + typename BiasGradEpilogue::Epilogue; + typename Epilogue::OutputTileIterator::Params params{p.num_keys}; + auto dst_ptr = p.grad_bias2_ptr + query_start * p.num_keys + key_start; + typename Epilogue::OutputTileIterator output_iter( + params, dst_ptr, {num_queries_in_block, num_keys_in_block}, (int)thread_id); + Epilogue epilogue(shared_storage.gradB_epilogue(), + (int)thread_id, + (int)warp_id, + (int)lane_id); + epilogue(OutputOp(1), output_iter, accum); + } + + accum = accum * scale; + + __syncthreads(); + if (!MatmulGradK::DefaultMmaFromSmem::kIsTransposedA) { + auto tmpT = shared_storage.tmpT_shared_storage().accum_ref(); + // attn <- attn_T.T + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + tmpT.at({accum_n, accum_m}) = scalar_t(accum[idx]); + }, + [&](int accum_m) {}); + } + } + + MatmulDOIVJ::B2bGemm::accumToSmem( + shared_storage.tmp_shared_storage(), accum, lane_id, output_tile_coords); + __syncthreads(); + } + p.head_dim = warp_uniform(p.head_dim); + p.k_strideM = warp_uniform(p.k_strideM); + rematerializeThreadIds(); + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradQ matmul + // + // grad_q[i_start:i_end] += tmp @ k_j + ///////////////////////////////////////////////////////////////////////////////////////////////// + // Skip the loop & associated branches if we know at compile time the number + // of iterations + constexpr bool kSingleIterationGradQ = kMaxK <= MatmulGradQ::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradQ ? 1 : p.head_dim); + col += MatmulGradQ::ThreadblockShape::kN) { + using Mma = typename MatmulGradQ::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + cutlass::gemm::GemmCoord problem_size( + num_queries_in_block, + false ? MatmulGradQ::ThreadblockShape::kN : p.head_dim - col, + num_keys_in_block); + + // k_j + typename Mma::IteratorB iterator_B({int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM + col, + {problem_size.k(), problem_size.n()}, + thread_id, + no_offset); + + auto a = shared_storage.tmp_shared_storage().accum_ref(); + Mma mma(shared_storage.mm_gradQ(), + shared_storage.tmp_shared_storage(), + thread_id, + warp_id, + lane_id, + problem_size.k()); + + typename Mma::FragmentC accum; + + bool isFirst = key_start == 0; + int col_id = col / MatmulGradQ::ThreadblockShape::kN; + int num_cols = + kSingleIterationGradQ ? 1 : ceil_div(p.head_dim, MatmulGradQ::ThreadblockShape::kN); + int storage_id = (col_id + query_start / kBlockSizeI * num_cols); + AccumTileGmem gmem_tile{p.workspace_gq + storage_id * AccumTileGmem::kElementsStored}; + if (isFirst || !kNeedsAccumGradQ) { + accum.clear(); + } else { + gmem_tile.load(accum, thread_id); + } + + auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + mma.set_prologue_done(kPrologueGQ); + mma(gemm_k_iterations, accum, iterator_B, accum); + __syncthreads(); + bool isLastColumn = kSingleIterationGradQ || + (col + MatmulGradQ::ThreadblockShape::kN >= p.head_dim); + if (kPrologueGQ && !isLastColumn) { + prologueGradQ(col + MatmulGradQ::ThreadblockShape::kN); + } + + // Output results + int32_t next_query, next_key; + incrIteration(p, p.num_queries, key_start, next_query, next_key); + bool isLast = next_query > query_start || next_key >= p.num_keys; + if (kNeedsAccumGradQ && !isLast) { + gmem_tile.store(accum, thread_id); + } else { + typename MatmulGradQ::OutputTileIterator output_it( + typename MatmulGradQ::OutputTileIterator::Params{p.gQ_strideM()}, + p.grad_query_ptr + query_start * p.gQ_strideM() + col, + {problem_size.m(), problem_size.n()}, + thread_id); + accumulateInGmem(isLastColumn + ? shared_storage.gradQ_epilogue_lastIter() + : shared_storage.gradQ_epilogue(), + accum, + output_it, + isFirst || kNeedsAccumGradQ, + warp_id, + lane_id); + } + } + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradK matmul + // + // grad_k[i_start:i_end] += tmp.transpose(-2, -1) @ q_i + ///////////////////////////////////////////////////////////////////////////////////////////////// + rematerializeThreadIds(); + + constexpr bool kSingleIterationGradK = kMaxK <= MatmulGradK::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradK ? 1 : p.head_dim); + col += MatmulGradK::ThreadblockShape::kN) { + using Mma = typename MatmulGradK::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + cutlass::gemm::GemmCoord problem_size( + num_keys_in_block, + false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col, + num_queries_in_block); + auto createEpilogueIter = [&]() { + return typename MatmulGradK::OutputTileIterator( + typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, + p.grad_key_ptr + key_start * p.gK_strideM() + col, + {num_keys_in_block, + false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col}, + thread_id); + }; + + // q_i + typename Mma::IteratorB iterator_B({int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM + col, + {problem_size.k(), problem_size.n()}, + thread_id, + no_offset); + + auto getTmp = [&](int) { return &shared_storage.tmp_shared_storage(); }; + auto getTmpT = [&](int) { return &shared_storage.tmpT_shared_storage(); }; + // this is basically: + // opA = kIsTransposedA ? getTmp() : getTmpT(); + bool constexpr kIsTransposedA = MatmulGradK::DefaultMmaFromSmem::kIsTransposedA; + auto& opA = + *call_conditional::apply( + getTmp, getTmpT, 0); + Mma mma(shared_storage.mm_gradK(), opA, thread_id, warp_id, lane_id, problem_size.k()); + + int storage_id = col / MatmulGradK::ThreadblockShape::kN; + AccumTileGmem gmem_tile{p.workspace_gk + storage_id * AccumTileGmem::kElementsStored}; + if (!kOutputInRF) { + if (isFirstQuery || !kNeedsAccumGradK) { + output_frags.gradK.clear(); + } else { + gmem_tile.load(output_frags.gradK, thread_id); + } + } + mma.set_prologue_done(kPrologueGK); + + auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + + mma(gemm_k_iterations, output_frags.gradK, iterator_B, output_frags.gradK); + __syncthreads(); + bool isLastColumn = kSingleIterationGradK || + col + MatmulGradK::ThreadblockShape::kN >= p.head_dim; + if (kPrologueGK && !isLastColumn) { + prologueGradK(col + MatmulGradK::ThreadblockShape::kN); + } + + if (kPrologueQK && isLastColumn) { + int32_t next_query, next_key; + incrIteration(p, query_start, key_start, next_query, next_key); + DISPATCH_BOOL(next_key != key_start, kForceReloadK, ([&]() { + prologueQkNextIteration( + shared_storage, p, next_query, next_key, warp_id, lane_id); + })); + } + + // Output results + if (!kOutputInRF) { + if (kNeedsAccumGradK && !isLastQuery) { + gmem_tile.store(output_frags.gradK, thread_id); + } else { + accumulateInGmem(isLastColumn + ? shared_storage.gradK_epilogue_final() + : shared_storage.gradK_epilogue(), + output_frags.gradK, + createEpilogueIter(), + isFirstQuery || kNeedsAccumGradK, + warp_id, + lane_id); + __syncthreads(); + } + } + } + } + + static CUTLASS_DEVICE int32_t getQueryStart(Params const& p, int32_t key_start) { return 0; }; + + static CUTLASS_DEVICE void incrIteration(Params const& p, + int32_t query_start, + int32_t key_start, + int32_t& next_query, + int32_t& next_key) + { + next_query = query_start + kBlockSizeI; + next_key = key_start; + if (next_query >= p.num_queries) { + next_key = key_start + kBlockSizeJ; + next_query = getQueryStart(p, next_key); + } + } + + template + static CUTLASS_DEVICE void prologueQkNextIteration(SharedStorage& shared_storage, + Params const& p, + int32_t query_start, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) + { + if (query_start >= p.num_queries || key_start >= p.num_keys) { return; } + + static constexpr bool kReloadK = kForceReloadK || !MatmulQK::Mma::kSmemContainsEntireMat; + int thread_id = 32 * warp_id + lane_id; + typename MatmulQK::Mma::IteratorA iterator_A({int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM, + {p.num_keys - key_start, p.head_dim}, + thread_id, + cutlass::MatrixCoord{0, 0}); + + typename MatmulQK::Mma::IteratorB iterator_B({int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM, + {p.head_dim, p.num_queries - query_start}, + thread_id, + cutlass::MatrixCoord{0, 0}); + + MatmulQK::Mma::prologue(shared_storage.mm_qk_k(), + shared_storage.mm_qk_q(), + iterator_A, + iterator_B, + thread_id, + p.head_dim); + } + + template + static CUTLASS_DEVICE void writeFragsToGmem(SharedStorage& shared_storage, + OutputFragments& output_frags, + Params const& p, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) + { + uint16_t thread_id = 32 * warp_id + lane_id; + int32_t num_keys_in_block = + skipBoundsChecks + ? MatmulQK::Mma::Shape::kM + : cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, p.num_keys - key_start); + typename MatmulGradV::OutputTileIterator outputV_it( + typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, + p.grad_value_ptr + key_start * p.gV_strideM(), + {num_keys_in_block, p.head_dim_value}, + thread_id); + accumulateInGmem(shared_storage.gradV_epilogue_final(), + output_frags.gradV, + outputV_it, + true, + warp_id, + lane_id); + + typename MatmulGradK::OutputTileIterator outputK_it( + typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, + p.grad_key_ptr + key_start * p.gK_strideM(), + {num_keys_in_block, false ? MatmulGradK::ThreadblockShape::kN : p.head_dim}, + thread_id); + accumulateInGmem(shared_storage.gradK_epilogue_final(), + output_frags.gradK, + outputK_it, + true, + warp_id, + lane_id); + } + + template + static CUTLASS_DEVICE void accumulateInGmem( + typename MatmulT::DefaultEpilogue::SharedStorage& epilogue_smem, + typename MatmulT::Mma::FragmentC const& accum, + typename MatmulT::OutputTileIterator output_it, + bool first, + uint8_t warp_id, + uint8_t lane_id) + { + using DefaultEpilogue = typename MatmulT::DefaultEpilogue; + using DefaultOutputOp = typename MatmulT::DefaultOutputOp; + using Mma = typename MatmulT::Mma; + int thread_id = 32 * warp_id + lane_id; + DISPATCH_BOOL( + first, kIsFirst, ([&]() { + static constexpr auto ScaleType = + kIsFirst ? cutlass::epilogue::thread::ScaleType::Nothing + : cutlass::epilogue::thread::ScaleType::NoBetaScaling; + using EpilogueOutputOp = typename cutlass::epilogue::thread::LinearCombination< + typename DefaultOutputOp::ElementOutput, + DefaultOutputOp::kCount, + typename DefaultOutputOp::ElementAccumulator, + typename DefaultOutputOp::ElementCompute, + ScaleType>; + using Epilogue = typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename MatmulT::OutputTileIterator, + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true // IterationsUnroll + >; + EpilogueOutputOp rescale({1, 1}); + Epilogue epilogue(epilogue_smem, thread_id, warp_id, lane_id); + epilogue(rescale, output_it, accum, output_it); + })); + } + + template + static CUTLASS_DEVICE void computeDelta(Params const& p, + int32_t query_start, + uint8_t warp_id, + uint8_t lane_id) + { + // Each thread computes one value for Delta + // Depending on warp configuration, we might have multiple + // threads of the same warp working on the same row + using AccessType = cutlass::Array; + static_assert(kNumThreads >= kBlockSizeI, ""); + static constexpr int kNumThreadsPerLine = kNumThreads / kBlockSizeI; + int16_t thread_id = 32 * warp_id + lane_id; + + int16_t laneFirstCol = kElementsPerAccess * (lane_id % kNumThreadsPerLine); + int16_t laneRow = thread_id / kNumThreadsPerLine; + bool rowPred = (query_start + laneRow) < p.num_queries; + bool pred = rowPred; + + // on windows, previous syntax __restrict__ AccessType* + // resulted in error: "restrict" is not allowed + const AccessType* __restrict__ grad_output_ptr = reinterpret_cast( + p.grad_output_ptr + (query_start + laneRow) * p.gO_strideM + laneFirstCol); + const AccessType* __restrict__ output_ptr = reinterpret_cast( + p.output_ptr + (query_start + laneRow) * p.o_strideM() + laneFirstCol); + + static constexpr int64_t kMaxIters = kMaxK / (kElementsPerAccess * kNumThreadsPerLine); + constexpr int kPipelineStages = 2; + accum_t delta_value = accum_t(0); + using GlobalLoad = cutlass::arch::global_load; + AccessType frag_grad_output[kPipelineStages]; + AccessType frag_output[kPipelineStages]; + + auto loadAndIncrement = [&](int ld_pos, bool is_valid) { + frag_grad_output[ld_pos].clear(); + frag_output[ld_pos].clear(); + GlobalLoad(frag_grad_output[ld_pos], grad_output_ptr, is_valid); + GlobalLoad(frag_output[ld_pos], output_ptr, is_valid); + grad_output_ptr += kNumThreadsPerLine; + output_ptr += kNumThreadsPerLine; + }; + + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kPipelineStages - 1; ++iter) { + int ld_pos = iter % kPipelineStages; + pred = pred && (laneFirstCol + iter * kElementsPerAccess * kNumThreadsPerLine) < + p.head_dim_value; + loadAndIncrement(ld_pos, pred); + } + auto columnIteration = [&](int iter) { + // Load for next iter + int ld_pos = (iter + kPipelineStages - 1) % kPipelineStages; + pred = pred && (laneFirstCol + (iter + kPipelineStages - 1) * kElementsPerAccess * + kNumThreadsPerLine) < p.head_dim_value; + loadAndIncrement(ld_pos, pred); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < AccessType::kElements; ++i) { + delta_value += accum_t(frag_output[iter % kPipelineStages][i]) * + accum_t(frag_grad_output[iter % kPipelineStages][i]); + } + }; + + // If we have a small lower-bound for K, we can unroll the loop + if (kMaxK <= 256) { + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kMaxIters; ++iter) { columnIteration(iter); } + } else { + int num_iters = ceil_div(p.head_dim_value, kElementsPerAccess * kNumThreadsPerLine) * + (kElementsPerAccess * kNumThreadsPerLine); + for (int iter = 0; iter < num_iters; ++iter) { columnIteration(iter); } + } + + // Reduce between workers + static_assert(kNumThreadsPerLine == 1 || kNumThreadsPerLine == 2 || kNumThreadsPerLine == 4, + ""); + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kNumThreadsPerLine; i *= 2) { + delta_value = delta_value + __shfl_xor_sync(0xffffffff, delta_value, i); + } + + // Store in gmem + if (rowPred) { p.delta_ptr[query_start + laneRow] = delta_value; } + } +}; + +template +__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) + attention_kernel_backward_batched_impl(typename AK::Params p) +{ + if (!p.advance_to_block()) { return; } + AK::attention_kernel(p); +} + +template +__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) + attention_kernel_backward_batched(typename AK::Params params); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h new file mode 100644 index 0000000000000000000000000000000000000000..e3b11ebcc6615156179c5f9047acb0d65688be7b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h @@ -0,0 +1,986 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include + +#include "cutlass/bfloat16.h" +#include "cutlass/fast_math.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/vector.h" +#include "cutlass/matrix.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" + +#include "cutlass/epilogue/threadblock/default_epilogue_simt.h" +#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "cutlass/gemm/device/default_gemm_configuration.h" +#include "cutlass/gemm/kernel/default_gemm.h" +#include "cutlass/gemm/threadblock/default_mma.h" +#include "cutlass/gemm/threadblock/default_mma_core_simt.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" +#include "cutlass/gemm/threadblock/threadblock_swizzle.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/platform/platform.h" +#include "cutlass/transform/threadblock/predicated_tile_iterator.h" +#include "epilogue/epilogue_pipelined.h" +#include "epilogue/epilogue_rescale_output.h" +#include "gemm/find_default_mma.h" +#include "gemm/mma_from_smem.h" +#include "gemm_kernel_utils.h" +#include "transform/bias_broadcast.h" +#include "transform/tile_smem_loader.h" + +#include + +using namespace gemm_kernel_utils; + +namespace { +template +constexpr int getWarpsPerSm() +{ + return (Arch::kMinComputeCapability >= 80 && !cutlass::platform::is_same::value + ? 16 + : 12); +} +static CUTLASS_DEVICE float atomicMaxFloat(float* addr, float value) +{ + // source: https://stackoverflow.com/a/51549250 + return (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) + : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value))); +} +} // namespace + +template < + // The datatype of Q/K/V + typename scalar_t_, + // Architecture we are targeting (eg `cutlass::arch::Sm80`) + typename ArchTag, + // If Q/K/V are correctly aligned in memory and we can run a fast kernel + bool isAligned_, + int kQueriesPerBlock, + int kKeysPerBlock_, + bool kSingleValueIteration_, // = `value.shape[-1] <= kKeysPerBlock` + // This is quite slower on V100 for some reason + // Set to false if you know at compile-time you will never need dropout + bool kSupportsBias_ = false, + template class Broadcast1_ = BroadcastNoLoad, + template class Broadcast2_ = BroadcastNoLoad> +struct AttentionKernel { + using scalar_t = scalar_t_; + using accum_t = float; + using lse_scalar_t = float; + using output_t = scalar_t; + // Accumulator between 2 iterations + // Using `accum_t` improves perf on f16 at the cost of + // numerical errors + using output_accum_t = accum_t; + static constexpr bool kSupportsBias = kSupportsBias_; + static constexpr int kKeysPerBlock = kKeysPerBlock_; + static constexpr bool kIsAligned = isAligned_; + static constexpr bool kSingleValueIteration = kSingleValueIteration_; + static constexpr int32_t kAlignLSE = 32; // block size of backward + static constexpr bool kPreloadV = + ArchTag::kMinComputeCapability >= 80 && cutlass::sizeof_bits::value == 16; + static constexpr bool kKeepOutputInRF = kSingleValueIteration; + static constexpr bool kNeedsOutputAccumulatorBuffer = + !kKeepOutputInRF && !cutlass::platform::is_same::value; + + static_assert(kQueriesPerBlock % 32 == 0, ""); + static_assert(kKeysPerBlock % 32 == 0, ""); + static constexpr int kNumWarpsPerBlock = kQueriesPerBlock * kKeysPerBlock / (32 * 32); + static constexpr int kWarpSize = 32; + + // Launch bounds + static constexpr int kNumThreads = kWarpSize * kNumWarpsPerBlock; + static constexpr int kMinBlocksPerSm = getWarpsPerSm() / kNumWarpsPerBlock; + + struct Params { + // Input tensors + scalar_t* query_ptr; // [num_queries, num_heads, head_dim] + scalar_t* key_ptr; // [num_keys, num_heads, head_dim] + scalar_t* value_ptr; // [num_keys, num_heads, head_dim_value] + + // Output tensors + output_t* output_ptr; // [num_queries, num_heads, head_dim_value] + output_accum_t* output_accum_ptr; // [num_queries, num_heads, head_dim_value] + lse_scalar_t* logsumexp_ptr; // [num_heads, num_queries] - can be null + + // Scale + accum_t scale; + + // Dimensions/strides + int32_t head_dim; + int32_t head_dim_value; + int32_t num_queries; + int32_t num_keys; + + int32_t q_strideM; + int32_t k_strideM; + int32_t v_strideM; + // int32_t bias_strideM = 0; + + int32_t o_strideM = 0; + + // Everything below is only used in `advance_to_block` + // and shouldn't use registers + int32_t q_strideH; + int32_t k_strideH; + int32_t v_strideH; + // int32_t bias_strideH = 0; + + int64_t q_strideB; + int64_t k_strideB; + int64_t v_strideB; + // int32_t bias_strideB = 0; + + int32_t num_batches; + int32_t num_heads; + + // Parameters for biases + scalar_t* bias1_ptr = nullptr; + scalar_t* bias2_ptr = nullptr; + int32_t B = 0; + int32_t N = 0; + + // Moves pointers to what we should process + // Returns "false" if there is no work to do + CUTLASS_DEVICE bool advance_to_block() + { + auto batch_id = blockIdx.z; + auto head_id = blockIdx.y; + auto query_start = blockIdx.x * kQueriesPerBlock; + + auto lse_dim = ceil_div((int32_t)num_queries, kAlignLSE) * kAlignLSE; + + query_ptr += batch_id * q_strideB; + key_ptr += batch_id * k_strideB; + value_ptr += batch_id * v_strideB; + output_ptr += int64_t(batch_id * num_queries) * o_strideM; + if (output_accum_ptr != nullptr) { + output_accum_ptr += int64_t(batch_id * num_queries) * (head_dim_value * num_heads); + } + + int64_t q_start = 0, k_start = 0; + // Advance to the current batch / head / query_start + query_ptr += (q_start + query_start) * q_strideM + head_id * q_strideH; + key_ptr += k_start * k_strideM + head_id * k_strideH; + + value_ptr += k_start * v_strideM + head_id * v_strideH; + output_ptr += int64_t(q_start + query_start) * o_strideM + head_id * head_dim_value; + + if (output_accum_ptr != nullptr) { + output_accum_ptr += int64_t(q_start + query_start) * (head_dim_value * num_heads) + + head_id * head_dim_value; + } else { + // Accumulate directly in the destination buffer (eg for f32) + output_accum_ptr = (accum_t*)output_ptr; + } + + if (logsumexp_ptr != nullptr) { + // lse[batch_id, head_id, query_start] + logsumexp_ptr += batch_id * lse_dim * num_heads + head_id * lse_dim + query_start; + } + + using broadcast_1 = Broadcast1_; + if (kSupportsBias && broadcast_1::kEnable && bias1_ptr) { + bias1_ptr = broadcast_1::advance(bias1_ptr, + batch_id / N, + batch_id % N, + head_id, + num_queries * N, + num_queries, + 0); + } + using broadcast_2 = Broadcast2_; + if (kSupportsBias && broadcast_2::kEnable && bias2_ptr) { + auto strideB = num_heads * num_queries * num_keys; + auto strideH = num_queries * num_keys; + bias2_ptr = broadcast_2::advance( + bias2_ptr, batch_id / N, batch_id % N, head_id, strideB, 0, strideH); + } + + num_queries -= query_start; + num_batches = 0; // no longer used after + + // If num_queries == 1, and there is only one key head we're wasting + // 15/16th of tensor core compute In that case : + // - we only launch kernels for head_id % kQueriesPerBlock == 0 + // - we iterate over heads instead of queries (strideM = strideH) + if (num_queries == 1 && k_strideH == 0 && v_strideH == 0) { + if (head_id % kQueriesPerBlock != 0) return false; + q_strideM = q_strideH; + num_queries = num_heads; + num_heads = 1; // unused but here for intent + o_strideM = head_dim_value; + } + + // Make sure the compiler knows these variables are the same on all + // the threads of the warp. + query_ptr = warp_uniform(query_ptr); + key_ptr = warp_uniform(key_ptr); + value_ptr = warp_uniform(value_ptr); + output_ptr = warp_uniform(output_ptr); + output_accum_ptr = warp_uniform(output_accum_ptr); + logsumexp_ptr = warp_uniform(logsumexp_ptr); + num_queries = warp_uniform(num_queries); + num_keys = warp_uniform(num_keys); + num_heads = warp_uniform(num_heads); + head_dim = warp_uniform(head_dim); + head_dim_value = warp_uniform(head_dim_value); + o_strideM = warp_uniform(o_strideM); + if (kSupportsBias && broadcast_1::kEnable) { bias1_ptr = warp_uniform(bias1_ptr); } + if (kSupportsBias && broadcast_2::kEnable) { bias2_ptr = warp_uniform(bias2_ptr); } + return true; + } + + __host__ dim3 getBlocksGrid() const + { + return dim3(ceil_div(num_queries, (int32_t)kQueriesPerBlock), num_heads, num_batches); + } + + __host__ dim3 getThreadsGrid() const { return dim3(kWarpSize, kNumWarpsPerBlock, 1); } + }; + + struct MM0 { + /* + In this first matmul, we compute a block of `Q @ K.T`. + While the calculation result is still hot in registers, we update + `mi`, `m_prime`, `s_prime` in shared-memory, and then store this value + into a shared-memory ("AccumulatorSharedStorage") that is used later as + operand A for the second matmul (see MM1) + */ + using GemmType = DefaultGemmType; + + using OpClass = typename GemmType::OpClass; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration; + static constexpr int kAlignmentA = kIsAligned ? DefaultConfig::kAlignmentA + : GemmType::kMinimumAlignment; + static constexpr int kAlignmentB = kIsAligned ? DefaultConfig::kAlignmentB + : GemmType::kMinimumAlignment; + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using DefaultMma = typename cutlass::gemm::threadblock::FindDefaultMma< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::ColumnMajor, // LayoutB, + kAlignmentB, + accum_t, + cutlass::layout::RowMajor, // LayoutC, + OpClass, + ArchTag, // ArchTag + ThreadblockShape, // ThreadblockShape + WarpShape, // WarpShape + typename GemmType::InstructionShape, // InstructionShape + DefaultConfig::kStages, // Should use `DefaultConfig::kStages`, but that + // uses too much smem + typename GemmType::Operator // Operator + >::DefaultMma; + using MmaCore = typename DefaultMma::MmaCore; + using IteratorA = typename DefaultMma::IteratorA; + using IteratorB = typename DefaultMma::IteratorB; + using Mma = typename DefaultMma::ThreadblockMma; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator::Iterator; + static_assert(MmaCore::WarpCount::kM * MmaCore::WarpCount::kN * MmaCore::WarpCount::kK == + kNumWarpsPerBlock, + ""); + + // used for efficient load of bias tile Bij from global to shared memory + using BiasLoader = + TileSmemLoader, + MmaCore::kThreads, + // input restriction: kv_len has to be a multiple of this value + 128 / cutlass::sizeof_bits::value>; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = + typename cutlass::gemm::threadblock::B2bGemm; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MM1 { + /** + Second matmul: perform `attn @ V` where `attn` is the attention (not + normalized) and stored in shared memory + */ + using GemmType = DefaultGemmType; + + using OpClass = typename GemmType::OpClass; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration; + static constexpr int kAlignmentA = DefaultConfig::kAlignmentA; // from smem + static constexpr int kAlignmentB = kIsAligned ? DefaultConfig::kAlignmentB + : GemmType::kMinimumAlignment; + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using LayoutB = cutlass::layout::RowMajor; + using DefaultGemm = + cutlass::gemm::kernel::DefaultGemm; + + using DefaultMmaFromSmem = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + typename MM0::AccumulatorSharedStorage, + false>; // kScaleOperandA + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + static_assert(WarpCount::kM * WarpCount::kN * WarpCount::kK == kNumWarpsPerBlock, ""); + + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = typename cutlass::epilogue::threadblock::PredicatedTileIterator< + typename DefaultEpilogue::OutputTileIterator::ThreadMap, + output_t>; + using OutputTileIteratorAccum = + typename cutlass::epilogue::threadblock::PredicatedTileIterator< + typename DefaultEpilogue::OutputTileIterator::ThreadMap, + output_accum_t>; + + struct SharedStorageMM1 { + typename Mma::SharedStorage mm; + }; + }; + + static constexpr int64_t kAlignmentQ = MM0::kAlignmentA; + static constexpr int64_t kAlignmentK = MM0::kAlignmentB; + static constexpr int64_t kAlignmentV = 1; + + // Shared storage - depends on kernel params + struct ScalingCoefs { + cutlass::Array m_prime; + cutlass::Array s_prime; + cutlass::Array mi; + }; + + struct SharedStorageEpilogueAtEnd : ScalingCoefs { + struct SharedStorageAfterMM0 { + // Everything here might be overwritten during MM0 + union { + // typename MM0::BiasLoader::SmemTile bias; + cutlass::AlignedBuffer bias; + typename MM0::AccumulatorSharedStorage si; + }; + typename MM1::SharedStorageMM1 mm1; + }; + + union { + typename MM0::Mma::SharedStorage mm0; + SharedStorageAfterMM0 after_mm0; + typename MM1::DefaultEpilogue::SharedStorage epilogue; + }; + + CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& epilogue_shared_storage() + { + return epilogue; + } + }; + + struct SharedStorageEpilogueInLoop : ScalingCoefs { + struct SharedStorageAfterMM0 { + // Everything here might be overwritten during MM0 + union { + // typename MM0::BiasLoader::SmemTile bias; + cutlass::AlignedBuffer bias; + typename MM0::AccumulatorSharedStorage si; + }; + typename MM1::SharedStorageMM1 mm1; + typename MM1::DefaultEpilogue::SharedStorage epilogue; + }; + + union { + typename MM0::Mma::SharedStorage mm0; + SharedStorageAfterMM0 after_mm0; + }; + + CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& epilogue_shared_storage() + { + return after_mm0.epilogue; + } + }; + + using SharedStorage = + typename cutlass::platform::conditional::type; + + static bool __host__ check_supported(Params const& p) + { + CHECK_ALIGNED_PTR(p.query_ptr, kAlignmentQ); + CHECK_ALIGNED_PTR(p.key_ptr, kAlignmentK); + CHECK_ALIGNED_PTR(p.value_ptr, kAlignmentV); + EVOFORMER_CHECK(p.q_strideM % kAlignmentQ == 0, "query is not correctly aligned (strideM)"); + EVOFORMER_CHECK(p.k_strideM % kAlignmentK == 0, "key is not correctly aligned (strideM)"); + EVOFORMER_CHECK(p.v_strideM % kAlignmentV == 0, "value is not correctly aligned (strideM)"); + EVOFORMER_CHECK(p.num_heads <= 1 || p.q_strideH % kAlignmentQ == 0, + "query is not correctly aligned (strideH)"); + EVOFORMER_CHECK(p.num_heads <= 1 || p.k_strideH % kAlignmentK == 0, + "key is not correctly aligned (strideH)"); + EVOFORMER_CHECK(p.num_heads <= 1 || p.v_strideH % kAlignmentV == 0, + "value is not correctly aligned (strideH)"); + return true; + } + + static void CUTLASS_DEVICE attention_kernel(Params& p) + { + // In this block, we will only ever: + // - read query[query_start:query_end, :] + // - write to output[query_start:query_end, :] + + extern __shared__ char smem_buffer[]; + SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); + auto& m_prime = shared_storage.m_prime; + auto& s_prime = shared_storage.s_prime; + auto& mi = shared_storage.mi; + const uint32_t query_start = blockIdx.x * kQueriesPerBlock; + + static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); + if (thread_id() < kQueriesPerBlock) { + s_prime[thread_id()] = accum_t(0); + m_prime[thread_id()] = -cutlass::platform::numeric_limits::infinity(); + mi[thread_id()] = -cutlass::platform::numeric_limits::infinity(); + } + typename MM1::Mma::FragmentC accum_o; + accum_o.clear(); + + auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator { + using OutputTileIterator = typename MM1::OutputTileIterator; + return OutputTileIterator( + typename OutputTileIterator::Params{(int32_t)p.o_strideM}, + p.output_ptr, + typename OutputTileIterator::TensorCoord{p.num_queries, p.head_dim_value}, + thread_id(), + {0, col}); + }; + + auto createOutputAccumIter = [&](int col) -> typename MM1::OutputTileIteratorAccum { + using OutputTileIteratorAccum = typename MM1::OutputTileIteratorAccum; + return OutputTileIteratorAccum( + typename OutputTileIteratorAccum::Params{(int32_t)(p.head_dim_value * p.num_heads)}, + p.output_accum_ptr, + typename OutputTileIteratorAccum::TensorCoord{p.num_queries, p.head_dim_value}, + thread_id(), + {0, col}); + }; + + // Iterate through keys + for (int32_t iter_key_start = 0; iter_key_start < p.num_keys; + iter_key_start += kKeysPerBlock) { + int32_t problem_size_0_m = cutlass::fast_min((int32_t)kQueriesPerBlock, p.num_queries); + int32_t problem_size_0_n = + cutlass::fast_min(int32_t(kKeysPerBlock), p.num_keys - iter_key_start); + int32_t const& problem_size_0_k = p.head_dim; + int32_t const& problem_size_1_n = p.head_dim_value; + int32_t const& problem_size_1_k = problem_size_0_n; + + auto prologueV = [&](int blockN) { + typename MM1::Mma::IteratorB iterator_V( + typename MM1::IteratorB::Params{MM1::LayoutB(p.v_strideM)}, + p.value_ptr + iter_key_start * p.v_strideM, + {problem_size_1_k, problem_size_1_n}, + thread_id(), + cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); + MM1::Mma::prologue( + shared_storage.after_mm0.mm1.mm, iterator_V, thread_id(), problem_size_1_k); + }; + + __syncthreads(); // Need to have shared memory initialized, and `m_prime` + // updated from end of prev iter + // + // MATMUL: Q.K_t + // + // Computes the block-matrix product of: + // (a) query[query_start:query_end, :] + // with + // (b) key[iter_key_start:iter_key_start + kKeysPerBlock] + // and stores that into `shared_storage.si` + // + + // Compute threadblock location + cutlass::gemm::GemmCoord tb_tile_offset = {0, 0, 0}; + + cutlass::MatrixCoord tb_offset_A{tb_tile_offset.m() * MM0::Mma::Shape::kM, + tb_tile_offset.k()}; + + cutlass::MatrixCoord tb_offset_B{tb_tile_offset.k(), + tb_tile_offset.n() * MM0::Mma::Shape::kN}; + + // Construct iterators to A and B operands + typename MM0::IteratorA iterator_A( + typename MM0::IteratorA::Params(typename MM0::MmaCore::LayoutA(p.q_strideM)), + p.query_ptr, + {problem_size_0_m, problem_size_0_k}, + thread_id(), + tb_offset_A); + + typename MM0::IteratorB iterator_B( + typename MM0::IteratorB::Params(typename MM0::MmaCore::LayoutB(p.k_strideM)), + p.key_ptr + iter_key_start * p.k_strideM, + {problem_size_0_k, problem_size_0_n}, + thread_id(), + tb_offset_B); + + auto my_warp_id = warp_id(); + auto my_lane_id = lane_id(); + + // Construct thread-scoped matrix multiply + typename MM0::Mma mma(shared_storage.mm0, thread_id(), my_warp_id, my_lane_id); + + typename MM0::Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = + (problem_size_0_k + MM0::Mma::Shape::kK - 1) / MM0::Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + __syncthreads(); + + if (kPreloadV) { + prologueV(0); + } else { + MM1::Mma::drain_cp_asyncs(); + } + + typename MM0::Mma::Operator::IteratorC::TensorCoord iteratorC_tile_offset = { + (tb_tile_offset.m() * MM0::Mma::WarpCount::kM) + + (my_warp_id % MM0::Mma::WarpCount::kM), + (tb_tile_offset.n() * MM0::Mma::WarpCount::kN) + + (my_warp_id / MM0::Mma::WarpCount::kM)}; + + // multiply by scaling factor + // if (kSupportsBias) { + // accum = + // cutlass::multiplies()(p.scale, + // accum); + // } + + if (kSupportsBias) { + cutlass::TensorRef bias_tensor_ref( + shared_storage.after_mm0.bias.data(), + cutlass::layout::RowMajor(MM0::ThreadblockShape::kN)); + using Shape = + cutlass::MatrixShape; + AttentionBiasEpilogue + bias_epilogue; + bias_epilogue(bias_tensor_ref, + p.bias1_ptr + iter_key_start, + p.bias2_ptr + query_start * p.num_keys + iter_key_start, + thread_id(), + {problem_size_0_m, problem_size_0_n}, + p.num_keys); + // Pij += Bij, Pij is in register fragment and Bij is in shared memory + auto lane_offset = MM0::AccumLambdaIterator::get_lane_offset( + lane_id(), warp_id(), iteratorC_tile_offset); + MM0::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + if (accum_m < problem_size_0_m && accum_n < problem_size_0_n) { + accum[idx] = + accum[idx] * p.scale + bias_tensor_ref.at({accum_m, accum_n}); + } + }, + [&](int accum_m) {}); + } + + DISPATCH_BOOL(iter_key_start == 0, kIsFirst, ([&] { + DISPATCH_BOOL( + p.num_keys - iter_key_start >= kKeysPerBlock, kFullColumns, ([&] { + // Update `mi` from accum stored in registers + // Also does accum[i] <- exp(accum[i] - mi) + iterative_softmax(accum_o, + accum, + mi, + m_prime, + s_prime, + lane_id(), + thread_id(), + warp_id(), + p.num_keys - iter_key_start, + iteratorC_tile_offset, + kSupportsBias ? 1.0f : p.scale); + })); + })); + + // Output results to shared-memory + int warp_idx_mn_0 = + my_warp_id % (MM0::Mma::Base::WarpCount::kM * MM0::Mma::Base::WarpCount::kN); + auto output_tile_coords = + cutlass::MatrixCoord{warp_idx_mn_0 % MM0::Mma::Base::WarpCount::kM, + warp_idx_mn_0 / MM0::Mma::Base::WarpCount::kM}; + + MM0::B2bGemm::accumToSmem( + shared_storage.after_mm0.si, accum, my_lane_id, output_tile_coords); + + __syncthreads(); + + // + // MATMUL: Attn . V + // Run the matmul `attn @ V` for a block of attn and V. + // `attn` is read from shared memory (in `shared_storage_si`) + // `V` is read from global memory (with iterator_B) + // + + const int64_t nBlockN = + kSingleValueIteration + ? 1 + : ceil_div((int64_t)problem_size_1_n, int64_t(MM1::ThreadblockShape::kN)); + for (int blockN = 0; blockN < nBlockN; ++blockN) { + int gemm_k_iterations = + (problem_size_1_k + MM1::Mma::Shape::kK - 1) / MM1::Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add and store it in accum + // (in registers) + if (!kPreloadV) { + __syncthreads(); // we share shmem between mma and epilogue + } + + typename MM1::Mma::IteratorB iterator_V( + typename MM1::IteratorB::Params{MM1::LayoutB(p.v_strideM)}, + p.value_ptr + iter_key_start * p.v_strideM, + {problem_size_1_k, problem_size_1_n}, + thread_id(), + cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); + typename MM1::Mma mma_pv(shared_storage.after_mm0.mm1.mm, + shared_storage.after_mm0.si, + (int)thread_id(), + (int)warp_id(), + (int)lane_id(), + (int)problem_size_1_k); + mma_pv.set_prologue_done(kPreloadV); + if (!kKeepOutputInRF) { accum_o.clear(); } + mma_pv(gemm_k_iterations, accum_o, iterator_V, accum_o); + __syncthreads(); + + if (kPreloadV && !kSingleValueIteration && blockN + 1 < nBlockN) { + prologueV(blockN + 1); + } + + if (!kKeepOutputInRF) { + MM1::Mma::drain_cp_asyncs(); + DISPATCH_BOOL( + iter_key_start == 0, kIsFirst, ([&] { + DISPATCH_BOOL( + (iter_key_start + kKeysPerBlock) >= p.num_keys, kIsLast, ([&] { + using DefaultEpilogue = typename MM1::DefaultEpilogue; + using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; + using ElementCompute = typename DefaultOp::ElementCompute; + using EpilogueOutputOp = typename cutlass::epilogue::thread:: + MemoryEfficientAttentionNormalize< + typename cutlass::platform:: + conditional:: + type, + output_accum_t, + DefaultOp::kCount, + typename DefaultOp::ElementAccumulator, + ElementCompute, + kIsFirst, + kIsLast, + cutlass::Array>; + using Epilogue = + typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename MM1::Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename cutlass::platform::conditional< + kIsLast, + typename MM1::OutputTileIterator, + typename MM1::OutputTileIteratorAccum>::type, + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true, // IterationsUnroll + typename MM1::OutputTileIteratorAccum // Read + // iterator + >; + + int col = blockN * MM1::Mma::Shape::kN; + auto source_iter = createOutputAccumIter(col); + auto dest_iter = + call_conditional:: + apply(createOutputIter, createOutputAccumIter, col); + EpilogueOutputOp rescale(s_prime, m_prime); + Epilogue epilogue(shared_storage.epilogue_shared_storage(), + thread_id(), + warp_id(), + lane_id()); + epilogue(rescale, dest_iter, accum_o, source_iter); + })); + })); + if (!kSingleValueIteration) { __syncthreads(); } + } + } + __syncthreads(); // we modify `m_prime` after + } + + if (kKeepOutputInRF) { + constexpr bool kIsFirst = true; + constexpr bool kIsLast = true; + using DefaultEpilogue = typename MM1::DefaultEpilogue; + using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; + using ElementCompute = typename DefaultOp::ElementCompute; + using EpilogueOutputOp = + typename cutlass::epilogue::thread::MemoryEfficientAttentionNormalize< + output_t, // output + output_accum_t, // source + DefaultOp::kCount, + typename DefaultOp::ElementAccumulator, // accum + output_accum_t, // compute + kIsFirst, + kIsLast, + cutlass::Array>; + using Epilogue = typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename MM1::Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename MM1::OutputTileIterator, // destination + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true, // IterationsUnroll + typename MM1::OutputTileIteratorAccum // source tile + >; + auto dest_iter = createOutputIter(0); + EpilogueOutputOp rescale(s_prime, m_prime); + Epilogue epilogue( + shared_storage.epilogue_shared_storage(), thread_id(), warp_id(), lane_id()); + MM1::Mma::drain_cp_asyncs(); + epilogue(rescale, dest_iter, accum_o); + } + + // 7. Calculate logsumexp + // To make the backward easier, we pad logsumexp with `inf` + // this avoids a few bound checks, and is not more expensive during fwd + static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); + if (p.logsumexp_ptr && thread_id() < kQueriesPerBlock) { + auto lse_dim = ceil_div((int32_t)p.num_queries, kAlignLSE) * kAlignLSE; + if (thread_id() < p.num_queries) { + p.logsumexp_ptr[thread_id()] = + accum_t(mi[thread_id()]) + cutlass::fast_log(accum_t(s_prime[thread_id()])); + } else if (thread_id() < lse_dim) { + p.logsumexp_ptr[thread_id()] = + cutlass::platform::numeric_limits::infinity(); + } + } + } + + template + CUTLASS_DEVICE static void iterative_softmax( + typename WarpIteratorC::Fragment& frag_o, // output so far + typename WarpIteratorC::Fragment& frag, + cutlass::Array& mi, + cutlass::Array& m_prime, + cutlass::Array& s_prime, + int8_t lane_id, + int8_t thread_id, + int8_t warp_id, + int16_t max_col, + typename WarpIteratorC::TensorCoord const& tile_offset, + float scaling) + { + /* Iterates on the accumulator and corresponding position on result matrix + + (1) Update `mi[r]` to the max value of the row `r` + (2) In a second iteration do the following: + (a) accum <- exp(accum - mi) + (b) m_prime <- exp(m_prime - mi) + (c) s_prime <- s_prime * m_prime + sum(accum) + + All of this is done on registers, before we store all of this + on shared memory for the next matmul with Value. + */ + using Fragment = typename WarpIteratorC::Fragment; + using LambdaIterator = + typename DefaultMmaAccumLambdaIterator::Iterator; + // Convert to `accum_t` (rather than double) + constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E + if (!kIsFirst) { + if (thread_id < kQueriesPerBlock) { m_prime[thread_id] = mi[thread_id]; } + __syncthreads(); + } + + auto lane_offset = LambdaIterator::get_lane_offset(lane_id, warp_id, tile_offset); + + // First update `mi` to the max per-row + { + accum_t max; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { max = -cutlass::platform::numeric_limits::infinity(); }, + [&](int accum_m, int accum_n, int idx) { + if (kFullColumns || accum_n < max_col) { + max = cutlass::fast_max(max, frag[idx]); + } + }, + [&](int accum_m) { + // Having 4x atomicMax seems faster than reduce within warp + // first... + atomicMaxFloat(&mi[accum_m], max * scaling); + }); + } + frag = cutlass::multiplies()(scaling * kLog2e, frag); + + // Make sure we all share the update values for `mi` + __syncthreads(); + + if (thread_id < kQueriesPerBlock) { + auto m_prime_exp = exp2f(kLog2e * (m_prime[thread_id] - mi[thread_id])); + m_prime[thread_id] = m_prime_exp; + s_prime[thread_id] *= m_prime_exp; + } + __syncthreads(); // Update output fragments + if (kKeepOutputInRF && !kIsFirst) { + accum_t mp; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { mp = m_prime[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { frag_o[idx] *= mp; }, + [&](int accum_m) {}); + __syncthreads(); + } + // Update accum_m, accum_n, ... + { + accum_t mi_row, total_row; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { mi_row = kLog2e * mi[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + frag[idx] = (kFullColumns || accum_n < max_col) ? exp2f(frag[idx] - mi_row) + : accum_t(0.0); + }, + [&](int accum_m) {}); + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { total_row = 0.0; }, + [&](int accum_m, int accum_n, int idx) { total_row += frag[idx]; }, + [&](int accum_m) { + if (LambdaIterator::reduceSameRow( + lane_id, total_row, [](accum_t a, accum_t b) { return a + b; })) { + atomicAdd(&s_prime[accum_m], total_row); + } + }); + } + } + + static CUTLASS_DEVICE int8_t lane_id() { return threadIdx.x; } + static CUTLASS_DEVICE int8_t warp_id() { return threadIdx.y; } + static CUTLASS_DEVICE int16_t thread_id() { return threadIdx.x + threadIdx.y * blockDim.x; } +}; + +template +__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) + attention_kernel_batched_impl(typename AK::Params p) +{ + if (!p.advance_to_block()) { return; } + AK::attention_kernel(p); +} + +template +__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) + attention_kernel_batched(typename AK::Params params); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h new file mode 100644 index 0000000000000000000000000000000000000000..0f15a43574cfe3eba123e03b7835227275e762f4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h @@ -0,0 +1,148 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +// This does nothing. +template +struct BroadcastNoLoad { + using Fragment = + cutlass::Array; + static const bool kEnable = false; + CUTLASS_DEVICE static void load(Fragment& frag, + scalar_t* ptr, + int thread_id, + const cutlass::MatrixCoord& extent, + int stride) + { + } + CUTLASS_DEVICE static scalar_t* + advance(scalar_t* ptr, int B_id, int N_id, int H_id, int strideB, int strideN, int strideH) + { + return ptr; + } +}; + +// This is to load the bias matrix from the global memory with on-the-fly +// broadcast. The shape in global memory is [B, N, 1, 1, L]. Each time we load +// the last dimension as a L row vector, and we further broadcast the L vector +// to a tile of size [L, L] by repeating the L vector L times +template +struct BroadcastA : public BroadcastNoLoad { + using Base = BroadcastNoLoad; + static const bool kEnable = true; + using layout = cutlass::layout::AffineRank2RowMajor; + + using GmemTileIterator = cutlass::transform::threadblock:: + PredicatedTileIterator; + using Fragment = typename GmemTileIterator::Fragment; + + CUTLASS_DEVICE static void load(Fragment& frag, + scalar_t* ptr, + int thread_id, + const cutlass::MatrixCoord& extent, + int stride) + { + GmemTileIterator iter({layout(0, 1)}, ptr, extent, thread_id); + iter.load(frag); + } + + CUTLASS_DEVICE static scalar_t* + advance(scalar_t* ptr, int B_id, int N_id, int H_id, int strideB, int strideN, int strideH) + { + return ptr + B_id * strideB + N_id * strideN; + } +}; + +// This is to load the bias matrix from the global memory with on-the-fly +// broadcast. The shape in global memory is [B, 1, H, L, L]. Each time we load +// a [L, L] matrix. Different N use the same bias matrix when B and H are the +// same. +template +struct BroadcastB : public BroadcastNoLoad { + using Base = BroadcastNoLoad; + static const bool kEnable = true; + using layout = cutlass::layout::RowMajor; + + using GmemTileIterator = cutlass::transform::threadblock:: + PredicatedTileIterator; + using Fragment = typename GmemTileIterator::Fragment; + + CUTLASS_DEVICE static void load(Fragment& frag, + scalar_t* ptr, + int thread_id, + const cutlass::MatrixCoord& extent, + int stride) + { + GmemTileIterator iter({layout(stride)}, ptr, extent, thread_id); + iter.load(frag); + } + + CUTLASS_DEVICE static scalar_t* + advance(scalar_t* ptr, int B_id, int N_id, int H_id, int strideB, int strideN, int strideH) + { + return ptr + B_id * strideB + H_id * strideH; + } +}; + +template + class Broadcast1_, + template + class Broadcast2_> +struct AttentionBiasEpilogue { + using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap< + cutlass::layout::PitchLinearShape, + kThreads, + 1>; + + using Broadcast1 = Broadcast1_; + using Broadcast2 = Broadcast2_; + + Broadcast1 broadcast1; + Broadcast2 broadcast2; + + using Ref = cutlass::TensorRef; + using SmemTileIterator = cutlass::transform::threadblock:: + RegularTileIterator; + + CUTLASS_DEVICE void operator()(const Ref& ref, + scalar_t* ptr1, + scalar_t* ptr2, + int thread_id, + const cutlass::MatrixCoord& extent, + int stride) + { + static_assert(Broadcast1::Fragment::kElements == Broadcast2::Fragment::kElements, + "The two broadcast fragments must have the same number of " + "elements"); + typename SmemTileIterator::Fragment frag; + frag.clear(); + float* frag_ptr = reinterpret_cast(&frag); + if (Broadcast1::kEnable) { + typename Broadcast1::Fragment frag1; + frag1.clear(); + broadcast1.load(frag1, ptr1, thread_id, extent, stride); + scalar_t* frag1_ptr = reinterpret_cast(&frag1); + for (int i = 0; i < Broadcast1::Fragment::kElements; ++i) { + frag_ptr[i] += static_cast(frag1_ptr[i]); + } + } + if (Broadcast2::kEnable) { + typename Broadcast2::Fragment frag2; + frag2.clear(); + broadcast2.load(frag2, ptr2, thread_id, extent, stride); + scalar_t* frag2_ptr = reinterpret_cast(&frag2); + for (int i = 0; i < Broadcast2::Fragment::kElements; ++i) { + frag_ptr[i] += static_cast(frag2_ptr[i]); + } + } + SmemTileIterator iter(ref, thread_id); + iter.store(frag); + __syncthreads(); + } +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h new file mode 100644 index 0000000000000000000000000000000000000000..5f3f0ee507e53bb0ad7a0f8fc3fcd25982c177f7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h @@ -0,0 +1,93 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once +#include +#include "cutlass/aligned_buffer.h" +#include "cutlass/array.h" +#include "cutlass/coord.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/numeric_types.h" +#include "cutlass/platform/platform.h" +#include "cutlass/transform/pitch_linear_thread_map.h" +#include "cutlass/transform/threadblock/predicated_tile_iterator.h" +#include "cutlass/transform/threadblock/regular_tile_iterator.h" + +template // thread access width in elements +class TileSmemLoader { +public: + using Shape = ThreadblockTileShape; + using SmemTile = cutlass::AlignedBuffer; + + using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap< + cutlass::layout::PitchLinearShape, // strided + Threads, // Threads + ElementsPerAccess>; // ElementsPerAccess + + using GmemTileIterator = cutlass::transform::threadblock::PredicatedTileIterator< + ThreadblockTileShape, // Shape + scalar_t, // Element + cutlass::layout::RowMajor, // Layout + 0, // AdvanceRank + ThreadMap>; // ThreadMap + + using SmemTileIterator = + cutlass::transform::threadblock::RegularTileIterator; // ThreadMap + + using Fragment = typename GmemTileIterator::Fragment; + + /// load a tile from global memory into shared memory + CUTLASS_DEVICE + static void load(GmemTileIterator tile_load_iter, SmemTileIterator tile_store_iter) + { + Fragment tb_frag; + tb_frag.clear(); + tile_load_iter.load(tb_frag); + tile_store_iter.store(tb_frag); + + __syncthreads(); + } +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1a887b50e1a3d06c7e6f554b6df4399de071c712 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "fp_quantize.h" + +#include +#include +#include + +#define DISPATCH_QUANTIZE(T_TYPE, C_TYPE, mantisa, exponent) \ + if (val.options().dtype() == torch::T_TYPE) { \ + launch_quantization((C_TYPE*)val.data_ptr(), \ + (uint8_t*)out.data_ptr(), \ + num_groups, \ + group_size, \ + at::cuda::getCurrentCUDAStream(), \ + q_range, \ + q_bits, \ + q_mantisa_bits, \ + stochastic_rounding); \ + } + +at::Tensor quantize(torch::Tensor& out, + torch::Tensor& val, + int group_size, + int stochastic_rounding, + int q_bits, + int q_mantisa_bits) +{ + int total_elems = at::numel(val); + float q_range = q_bits == 8 ? (q_mantisa_bits == 3 ? 480.0 : 114688.0) : // fp8 ranges + (q_bits == 12 ? 510.0 : // fp12 range + (q_bits == 6 ? 28.0 : // fp6 range + 6.0)); // fp4 range (using power 2); TODO (Reza): add the power-4 + // in case accuracy is not matching! + int num_groups = total_elems / group_size; + + DISPATCH_QUANTIZE(kHalf, __half, 23, 8); +#ifdef BF16_AVAILABLE + DISPATCH_QUANTIZE(kBFloat16, __nv_bfloat16, 23, 8); +#endif + + return out; +} + +#define DISPATCH_DEQUANTIZE(T_TYPE, C_TYPE, mantisa) \ + if (val.options().dtype() == torch::T_TYPE) { \ + launch_dequantization((uint8_t*)val_q.data_ptr(), \ + (C_TYPE*)val.data_ptr(), \ + num_groups, \ + group_size, \ + q_mantisa_bits, \ + q_exponent_bits, \ + at::cuda::getCurrentCUDAStream()); \ + return; \ + } + +void dequantize(torch::Tensor& val, + torch::Tensor& val_q, + int group_size, + int q_mantisa_bits, + int q_exponent_bits) +{ + int total_elems = at::numel(val); + + int num_groups = total_elems / group_size; + + DISPATCH_DEQUANTIZE(kHalf, __half, 10); +#ifdef BF16_AVAILABLE + DISPATCH_DEQUANTIZE(kBFloat16, __nv_bfloat16, 7); +#endif +} + +#define DISPATCH_DEQUANTIZE_INDEX(T_TYPE, C_TYPE, mantisa) \ + if (val.options().dtype() == torch::T_TYPE) { \ + launch_selective_dequantization((uint8_t*)val_q.data_ptr(), \ + (C_TYPE*)val.data_ptr(), \ + (int32_t*)indexes.data_ptr(), \ + num_groups, \ + group_size, \ + num_indexes, \ + q_mantisa_bits, \ + q_exponent_bits, \ + at::cuda::getCurrentCUDAStream()); \ + return; \ + } +void selective_dequantize(torch::Tensor& val, + torch::Tensor& val_q, + torch::Tensor& indexes, + int group_size, + int q_mantisa_bits, + int q_exponent_bits) +{ + int total_elems = at::numel(val); + int num_indexes = indexes.size(0); + int num_groups = total_elems / group_size; + + DISPATCH_DEQUANTIZE_INDEX(kHalf, __half, 10); +#ifdef BF16_AVAILABLE + DISPATCH_DEQUANTIZE_INDEX(kBFloat16, __nv_bfloat16, 7); +#endif +} + +at::Tensor get_scales(torch::Tensor& out, int num_groups) +{ + auto options = at::TensorOptions() + .dtype(torch::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto scales = + torch::from_blob(out.data_ptr(), {num_groups, 1}, {out.stride(0) / 4, 1}, options); + return scales; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("quantize", &quantize, "quantize function"); + m.def("dequantize", &dequantize, "dequantize function"); + m.def("get_scales", &get_scales, "get scales function"); + m.def("selective_dequantize", &selective_dequantize, "selective dequantize function"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..66ea7392e0115ab9ebbeb66baad01f9d6fa80750 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu @@ -0,0 +1,532 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "context.h" +#include "fp_quantize.h" +#include "memory_access_utils.h" +#include "reduction_utils.h" + +#include +#include + +#include +#include + +#ifdef BF16_AVAILABLE +#include +#endif +#include + +using ROp = reduce::ROpType; + +namespace quantization { + +constexpr int access_granularity = 16; +constexpr int quanitzed_access_granularity = 4; +constexpr int quanitzed_access_granularity_6bits = 2; +constexpr int threads = 256; +constexpr int warps = threads / 32; + +} // namespace quantization + +template +__device__ void round(uint32_t& mantisa, uint32_t& dst_exponent, curandStatePhilox4_32_10_t* state) +{ + constexpr uint32_t mantisa_mask = (1 << (_mantisa_bits - q_mantisa_bits)) - 1; + uint32_t offset = stochastic_rounding ? (curand_poisson(state, 10) & mantisa_mask) + : 1 << (_mantisa_bits - q_mantisa_bits - 1); + mantisa += offset; + dst_exponent += (((mantisa & ~mantisa_mask) == (1 << _mantisa_bits)) ? 1 : 0); +} + +template +__device__ void clip(uint32_t& exponent, uint32_t& mantisa) +{ + constexpr uint32_t max_exponent = (1 << (q_exponent_bits - 1)) + (1 << (_exponent_bits - 1)); + constexpr uint32_t min_exponent = + (1 << (_exponent_bits - 1)) - ((1 << (q_exponent_bits - 1)) - 1); + if (exponent > max_exponent) { + exponent = max_exponent; + mantisa = (((uint32_t)-1) >> (32 - q_mantisa_bits)) << 1; //.11 .. 10 + } + if (exponent < min_exponent) { + exponent = min_exponent; + mantisa = 0; + } +} + +template +__global__ void apply_quantization(T* val, + uint8_t* q_val, + int group_size, + std::pair seed, + float q_range) +{ + int tidx = threadIdx.x; + int wid = tidx >> 5; + int lane = tidx & 0x1f; + int gid = blockIdx.x * quantization::warps + wid; + + constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; + constexpr uint32_t _mantisa_mask = (1 << _mantisa_bits) - 1; + constexpr uint32_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; + constexpr uint32_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); + // CG helpers + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); + constexpr uint32_t load_stride = vector_size * hw_warp_size; + constexpr uint32_t store_stride = (total_q_bits * vector_size / 8) * hw_warp_size; + const uint32_t thread_offset = lane * vector_size; + const uint32_t store_thread_offset = lane * (total_q_bits * vector_size / 8); + const uint32_t base_load_offset = gid * group_size + thread_offset; + const uint32_t base_store_offset = + gid * ((group_size * total_q_bits / 8) + 4) + + store_thread_offset; // 4-byte for saving the scale per group + const T* load_base_ptr = val + base_load_offset; + T tmp_buf[unroll * vector_size]; + T cur_max; + reduce::init(&cur_max); + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + +#pragma unroll + for (int i = 0; i < unroll; i++) { + if (i * load_stride + thread_offset < group_size) { + mem_access::load_global( + &tmp_buf[vector_size * i], load_base_ptr + i * load_stride); + for (int j = 0; j < vector_size; j++) + cur_max = reduce::element(cur_max, __habs(tmp_buf[i * vector_size + j])); + } + } + reduce::_block(tb, warp, &cur_max); + + int mantisa_mask = ((1 << q_mantisa_bits) - 1); + mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); + + uint8_t* store_base_ptr = q_val + base_store_offset; + float scale = (float)q_range / conversion::to(cur_max); +#pragma unroll + for (int i = 0; i < unroll; i++) { + if (i * load_stride + thread_offset < group_size) { + uint64_t q_buf = 0; + uint64_t q_buf1 = 0; +#pragma unroll + for (int j = 0; j < vector_size; j++) { + float val_f = conversion::to(tmp_buf[i * vector_size + j]) * scale; + uint32_t* data = reinterpret_cast(&val_f); + uint32_t sign = (data[0] & _sign_mask) >> (_mantisa_bits + _exponent_bits); + uint32_t cur_exponent = (data[0] & _exponent_mask) >> _mantisa_bits; + uint32_t dst_mantisa = (data[0] & _mantisa_mask); + + uint32_t dst_exponent = cur_exponent; + + round<_mantisa_bits, q_mantisa_bits, stochastic_rounding>( + dst_mantisa, dst_exponent, &state); + if (cur_exponent != 0) + clip<_mantisa_bits, _exponent_bits, q_mantisa_bits, q_exponent_bits>( + dst_exponent, dst_mantisa); + + dst_mantisa = (dst_mantisa & mantisa_mask) >> (_mantisa_bits - q_mantisa_bits); + + if (dst_exponent != (1 << q_exponent_bits) - 1) + dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + + (1 << (q_exponent_bits - 1)) - 1; + if (total_q_bits == 8 || total_q_bits == 4 || total_q_bits == 6) + q_buf = q_buf | + ((uint64_t)((uint8_t)(sign << (q_exponent_bits + q_mantisa_bits) | + (dst_exponent << q_mantisa_bits) | dst_mantisa)) + << j * total_q_bits); + else if (total_q_bits == 12) { + if (j < 5) + q_buf = + q_buf | + ((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) | + (dst_exponent << q_mantisa_bits) | dst_mantisa)) + << j * total_q_bits); + else + q_buf1 = + q_buf1 | + ((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) | + (dst_exponent << q_mantisa_bits) | dst_mantisa)) + << (j - 5) * total_q_bits); + } + } + if (total_q_bits == 12) { + uint64_t last_nibble_mask = 0xf; + last_nibble_mask = q_buf1 & last_nibble_mask; + q_buf = (last_nibble_mask << 60) | q_buf; + q_buf1 >>= 4; + } + uint8_t* int8_data = reinterpret_cast(&q_buf); + uint8_t* int8_data1 = reinterpret_cast(&q_buf1); + if (total_q_bits == 6) { + mem_access::store_global( + store_base_ptr + i * store_stride, int8_data); + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity_6bits, + int8_data + quantization::quanitzed_access_granularity_6bits); + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity_6bits * 2, + int8_data + 2 * quantization::quanitzed_access_granularity_6bits); + } else { + mem_access::store_global( + store_base_ptr + i * store_stride, int8_data); + + if (total_q_bits > 4) { + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity, + int8_data + quantization::quanitzed_access_granularity); + if (total_q_bits == 12) { + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity * 2, + int8_data1); + } + } + } + } + } + if (lane == 0) { + float q_scale = conversion::to(cur_max) / (float)q_range; + uint8_t* scale_as_int8 = reinterpret_cast(&q_scale); + uint32_t scale_offset = + gid * ((group_size * total_q_bits / 8) + 4) + (group_size * total_q_bits / 8); + if (total_q_bits != 6) + mem_access::store_global( + q_val + scale_offset, scale_as_int8); + else { + mem_access::store_global( + q_val + scale_offset, scale_as_int8); + mem_access::store_global( + q_val + scale_offset + quantization::quanitzed_access_granularity_6bits, + scale_as_int8 + quantization::quanitzed_access_granularity_6bits); + } + } +} + +template +__global__ void apply_dequantization(uint8_t* val, T* q_val, int group_size, int total_num_elements) +{ + constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); + int tidx = (blockIdx.x * blockDim.x + threadIdx.x) * vector_size; + + constexpr int quantized_bits = _mantisa_bits + _exponent_bits + 1; + constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; + constexpr uint16_t _mantisa_mask = (1 << _mantisa_bits) - 1; + constexpr uint16_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; + constexpr uint16_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); + const uint32_t g_index = (tidx / group_size); + const uint32_t group_size_bytes = (group_size * quantized_bits / 8); + const uint8_t* load_base_ptr = + val + g_index * (group_size_bytes + 4) + (tidx % group_size) * quantized_bits / 8; + + int mantisa_mask = ((1 << q_mantisa_bits) - 1); + mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); + + T* store_base_ptr = q_val + tidx; + float scale; + + uint8_t* scale_as_int8 = reinterpret_cast(&scale); + if (quantized_bits == 6) { + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + mem_access::load_global( + scale_as_int8 + quantization::quanitzed_access_granularity_6bits, + val + g_index * (group_size_bytes + 4) + group_size_bytes + + quantization::quanitzed_access_granularity_6bits); + } else + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + + if (tidx < total_num_elements) { + uint64_t q_buf_in; + uint64_t q_buf_in1; + uint8_t* int8_data = reinterpret_cast(&q_buf_in); + uint8_t* int8_data1 = reinterpret_cast(&q_buf_in1); + if (quantized_bits == 6) { + mem_access::load_global( + int8_data, load_base_ptr); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits, + load_base_ptr + quantization::quanitzed_access_granularity_6bits); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits * 2, + load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2); + + } else { + mem_access::load_global(int8_data, + load_base_ptr); + if (quantized_bits > 4) { + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity, + load_base_ptr + quantization::quanitzed_access_granularity); + if (quantized_bits == 12) { + mem_access::load_global( + int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2); + } + } + } + T store_buf[vector_size]; + uint16_t* q_buf = reinterpret_cast(store_buf); +#pragma unroll + for (int j = 0; j < vector_size; j++) { + uint16_t new_data; + if (j < 5 || quantized_bits != 12) { + new_data = (uint16_t)(q_buf_in >> (j * quantized_bits)); + } else { + if (j == 5) { + new_data = (uint16_t)(q_buf_in1); + new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60)); + } else + new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8)); + } + + uint16_t sign = (new_data & _sign_mask) >> (_mantisa_bits + _exponent_bits); + uint16_t dst_exponent = (new_data & _exponent_mask) >> _mantisa_bits; + uint16_t dst_mantisa = (new_data & _mantisa_mask); + + if (dst_exponent != (1 << q_exponent_bits) - 1) + dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + + (1 << (q_exponent_bits - 1)) - 1; + + q_buf[j] = + ((sign << (q_exponent_bits + q_mantisa_bits)) | (dst_exponent << q_mantisa_bits) | + (dst_mantisa << (q_mantisa_bits - _mantisa_bits))); + float up_cast = conversion::to(store_buf[j]); + store_buf[j] = conversion::to(up_cast * scale); + } + mem_access::store_global(store_base_ptr, store_buf); + } +} + +#define LAUNCH_FOR_QUANTIZATION_UNROLL(COUNT) \ + case COUNT: \ + apply_quantization \ + <<>>(val, q_val, group_size, seed, q_range); \ + break; + +template +void launch_quantization(T* val, + uint8_t* q_val, + int num_groups, + int group_size, + cudaStream_t stream, + float q_range, + int q_bits, + int q_mantisa_bits, + int stochastic_rounding) +{ + const dim3 grid((num_groups + quantization::warps - 1) / quantization::warps); + const dim3 block(quantization::threads); + + std::pair seed = FPContext::Instance().IncrementOffset(16); + + constexpr int vals_per_unroll = hw_warp_size * quantization::access_granularity / sizeof(T); + + const int copy_unroll = (group_size + vals_per_unroll - 1) / vals_per_unroll; + QUANT_SWITCH((q_bits - q_mantisa_bits - 1) * q_mantisa_bits + stochastic_rounding, [&] { + switch (copy_unroll) { + LAUNCH_FOR_QUANTIZATION_UNROLL(1) + LAUNCH_FOR_QUANTIZATION_UNROLL(2) + LAUNCH_FOR_QUANTIZATION_UNROLL(3) + LAUNCH_FOR_QUANTIZATION_UNROLL(4) + LAUNCH_FOR_QUANTIZATION_UNROLL(5) + LAUNCH_FOR_QUANTIZATION_UNROLL(6) + } + }); +} +#define INSTANTIATE_LAUNCH_QUANTIZATION(T, mantisa, exponent) \ + template void launch_quantization( \ + T*, uint8_t*, int, int, cudaStream_t, float q_range, int, int, int); +// fp8(E4M3), nearest-rounding +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_QUANTIZATION(__nv_bfloat16, 23, 8); +#endif +INSTANTIATE_LAUNCH_QUANTIZATION(__half, 23, 8); + +template +void launch_dequantization(uint8_t* val, + T* q_val, + int num_groups, + int group_size, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream) +{ + int blocks = ((num_groups * group_size) - 1) / + (quantization::threads * (quantization::access_granularity / sizeof(T))) + + 1; + const dim3 grid(blocks); + const dim3 block(quantization::threads); + DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] { + apply_dequantization + <<>>(val, q_val, group_size, (num_groups * group_size)); + }); +} +#define INSTANTIATE_LAUNCH_DEQUANTIZATION(T, mantisa) \ + template void launch_dequantization(uint8_t*, T*, int, int, int, int, cudaStream_t); +// fp8(E4M3) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_DEQUANTIZATION(__nv_bfloat16, 7); +#endif +INSTANTIATE_LAUNCH_DEQUANTIZATION(__half, 10); + +template +__global__ void apply_selective_dequantization(uint8_t* val, + T* q_val, + int32_t* indexes, + int group_size, + int total_num_elements) +{ + int index = indexes[blockIdx.x]; + constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); + int tidx = (blockIdx.y * blockDim.x + threadIdx.x) * vector_size; + int input_index = index * total_num_elements + tidx; + constexpr int quantized_bits = _mantisa_bits + _exponent_bits + 1; + constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; + constexpr uint16_t _mantisa_mask = (1 << _mantisa_bits) - 1; + constexpr uint16_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; + constexpr uint16_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); + const uint32_t g_index = (input_index / group_size); + const uint32_t group_size_bytes = (group_size * quantized_bits / 8); + const uint8_t* load_base_ptr = + val + g_index * (group_size_bytes + 4) + (input_index % group_size) * quantized_bits / 8; + + int mantisa_mask = ((1 << q_mantisa_bits) - 1); + mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); + + T* store_base_ptr = q_val + tidx + blockIdx.x * total_num_elements; + float scale; + + uint8_t* scale_as_int8 = reinterpret_cast(&scale); + if (quantized_bits == 6) { + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + mem_access::load_global( + scale_as_int8 + quantization::quanitzed_access_granularity_6bits, + val + g_index * (group_size_bytes + 4) + group_size_bytes + + quantization::quanitzed_access_granularity_6bits); + } else + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + + if (tidx < total_num_elements) { + uint64_t q_buf_in; + uint64_t q_buf_in1; + uint8_t* int8_data = reinterpret_cast(&q_buf_in); + uint8_t* int8_data1 = reinterpret_cast(&q_buf_in1); + if (quantized_bits == 6) { + mem_access::load_global( + int8_data, load_base_ptr); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits, + load_base_ptr + quantization::quanitzed_access_granularity_6bits); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits * 2, + load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2); + } else { + mem_access::load_global(int8_data, + load_base_ptr); + if (quantized_bits > 4) { + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity, + load_base_ptr + quantization::quanitzed_access_granularity); + if (quantized_bits == 12) { + mem_access::load_global( + int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2); + } + } + } + T store_buf[vector_size]; + uint16_t* q_buf = reinterpret_cast(store_buf); +#pragma unroll + for (int j = 0; j < vector_size; j++) { + uint16_t new_data; + if (j < 5 || quantized_bits != 12) { + new_data = (uint16_t)(q_buf_in >> (j * quantized_bits)); + } else { + if (j == 5) { + new_data = (uint16_t)(q_buf_in1); + new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60)); + } else + new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8)); + } + + uint16_t sign = (new_data & _sign_mask) >> (_mantisa_bits + _exponent_bits); + uint16_t dst_exponent = (new_data & _exponent_mask) >> _mantisa_bits; + uint16_t dst_mantisa = (new_data & _mantisa_mask); + + if (dst_exponent != (1 << q_exponent_bits) - 1) + dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + + (1 << (q_exponent_bits - 1)) - 1; + + q_buf[j] = + ((sign << (q_exponent_bits + q_mantisa_bits)) | (dst_exponent << q_mantisa_bits) | + (dst_mantisa << (q_mantisa_bits - _mantisa_bits))); + float up_cast = conversion::to(store_buf[j]); + store_buf[j] = conversion::to(up_cast * scale); + } + mem_access::store_global(store_base_ptr, store_buf); + } +} + +template +void launch_selective_dequantization(uint8_t* val, + T* q_val, + int32_t* indexes, + int num_groups, + int group_size, + int num_indexes, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream) +{ + int total_elements_per_index = (num_groups / num_indexes) * group_size; + int blocks = (total_elements_per_index - 1) / + (quantization::threads * (quantization::access_granularity / sizeof(T))) + + 1; + const dim3 grid(num_indexes, blocks); + const dim3 block(quantization::threads); + DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] { + apply_selective_dequantization + <<>>(val, q_val, indexes, group_size, total_elements_per_index); + }); +} +#define INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(T, mantisa) \ + template void launch_selective_dequantization( \ + uint8_t*, T*, int32_t*, int, int, int, int, int, cudaStream_t); +// fp8(E4M3) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__nv_bfloat16, 7); +#endif +INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__half, 10); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h new file mode 100644 index 0000000000000000000000000000000000000000..5bd9badbcb4fa06091a8c01b0b0d386153a24589 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h @@ -0,0 +1,66 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include +#include "cublas_v2.h" +#include "cuda.h" +#include "curand.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#define WARP_SIZE 32 + +class FPContext { +public: + FPContext() : _seed(42) + { + curandCreateGenerator(&_gen, CURAND_RNG_PSEUDO_DEFAULT); + curandSetPseudoRandomGeneratorSeed(_gen, 123); + } + + virtual ~FPContext() {} + + static FPContext& Instance() + { + static FPContext _ctx; + return _ctx; + } + + curandGenerator_t& GetRandGenerator() { return _gen; } + + cudaStream_t GetCurrentStream() + { + // get current pytorch stream. + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + return stream; + } + + std::pair IncrementOffset(uint64_t offset_inc) + { + uint64_t offset = _curr_offset; + _curr_offset += offset_inc; + return std::pair(_seed, offset); + } + + void SetSeed(uint64_t new_seed) { _seed = new_seed; } + +private: + curandGenerator_t _gen; + cublasHandle_t _cublasHandle; + uint64_t _seed; + uint64_t _curr_offset; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h new file mode 100644 index 0000000000000000000000000000000000000000..60c75541f603f8b58558ecdf1d5658407f4d2857 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include + +#include + +#ifdef BF16_AVAILABLE +#include +#endif +#include +#include + +#define QUANT_SWITCH(Q_BITS, ...) \ + [&] { \ + if (12 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 3; \ + __VA_ARGS__(); \ + } else if (13 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 3; \ + __VA_ARGS__(); \ + } else if (10 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (11 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (28 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 12; \ + constexpr int CONST_Q_MANTISA_BITS = 7; \ + __VA_ARGS__(); \ + } else if (29 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 12; \ + constexpr int CONST_Q_MANTISA_BITS = 7; \ + __VA_ARGS__(); \ + } else if (6 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 6; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (7 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 6; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (2 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 4; \ + constexpr int CONST_Q_MANTISA_BITS = 1; \ + __VA_ARGS__(); \ + } else { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 4; \ + constexpr int CONST_Q_MANTISA_BITS = 1; \ + __VA_ARGS__(); \ + } \ + }() + +#define DEQUANT_SWITCH(Q_MANTISA_EXPONENT_BITS, ...) \ + [&] { \ + if (12 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 3; \ + constexpr int CONST_Q_EXPONENT_BITS = 4; \ + __VA_ARGS__(); \ + } else if (10 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + constexpr int CONST_Q_EXPONENT_BITS = 5; \ + __VA_ARGS__(); \ + } else if (28 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 7; \ + constexpr int CONST_Q_EXPONENT_BITS = 4; \ + __VA_ARGS__(); \ + } else if (6 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + constexpr int CONST_Q_EXPONENT_BITS = 3; \ + __VA_ARGS__(); \ + } else { \ + constexpr int CONST_Q_MANTISA_BITS = 1; \ + constexpr int CONST_Q_EXPONENT_BITS = 2; \ + __VA_ARGS__(); \ + } \ + }() + +template +void launch_quantization(T* val, + uint8_t* q_val, + int num_groups, + int group_size, + cudaStream_t stream, + float q_range, + int q_bits, + int q_mantisa_bits, + int stochastic_rounding); + +template +void launch_dequantization(uint8_t* val, + T* q_val, + int num_groups, + int group_size, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream); + +template +void launch_selective_dequantization(uint8_t* val, + T* q_val, + int32_t* indexes, + int num_groups, + int group_size, + int num_indexes, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..460330b93d31f1a992efe02a9f3bddcff297f4ee --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp @@ -0,0 +1,154 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_gds_op.h" + +using namespace std; + +// For when there is more than 1 device +static std::map> base_ptr_registry; + +static void _safe_handle_register(const int fd, CUfileDescr_t& cf_descr, CUfileHandle_t& cf_handle) +{ + memset((void*)&cf_descr, 0, sizeof(CUfileDescr_t)); + cf_descr.handle.fd = fd; + cf_descr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; + CUfileError_t status = cuFileHandleRegister(&cf_handle, &cf_descr); + if (status.err != CU_FILE_SUCCESS) { + std::cerr << "file register error:" << cuFileGetErrorString(status) << std::endl; + close(fd); + exit(EXIT_FAILURE); + } +} + +static void* _find_base_ptr(const int64_t device, char* buf_ptr) +{ + void* base_ptr = nullptr; + int64_t last = -1; + int64_t ptr_diff; + for (const auto& value : base_ptr_registry[device]) { + ptr_diff = buf_ptr - (char*)value; + if (last == -1 && ptr_diff >= 0) { + last = ptr_diff; + base_ptr = value; + } else if (ptr_diff < last && ptr_diff >= 0) { + last = ptr_diff; + base_ptr = value; + } + } + if (!base_ptr || buf_ptr < base_ptr) { + std::cerr << "BASE PTR ERROR :" << base_ptr << " BUF PTR " << (void*)buf_ptr << std::endl; + for (const auto& value : base_ptr_registry[device]) { + std::cerr << "BASE PTR AVAIL :" << value << std::endl; + } + exit(EXIT_FAILURE); + } + + return base_ptr; +} + +void gds_op_desc_t::add_buffer_to_registry(const torch::Tensor& buffer) +{ + const int64_t device = buffer.get_device(); + void* reg_ptr = buffer.data_ptr(); + + // TODO: add checking to make sure pointer isn't already in set + const auto it = base_ptr_registry.find(device); + if (it == base_ptr_registry.end()) { + std::set new_ptr_set; + new_ptr_set.insert(reg_ptr); + base_ptr_registry.insert(std::pair>(device, new_ptr_set)); + } else { + base_ptr_registry[device].insert(reg_ptr); + } + + check_cudaruntimecall(cudaSetDevice(device)); + CUfileError_t status = cuFileBufRegister(reg_ptr, buffer.nbytes(), 0); + if (status.err != CU_FILE_SUCCESS) { + std::cerr << "buffer register failed:" << cuFileGetErrorString(status) << std::endl; + exit(EXIT_FAILURE); + } +} + +void gds_op_desc_t::remove_buffer_from_registry(const torch::Tensor& buffer) +{ + const int64_t device = buffer.get_device(); + void* reg_ptr = buffer.data_ptr(); + + // std::cout << "DEREG PTR " << reg_ptr << std::endl; + check_cudaruntimecall(cudaSetDevice(device)); + cuFileBufDeregister(reg_ptr); + + // Remove from tracked registry + base_ptr_registry[device].erase(reg_ptr); +} + +gds_op_desc_t::gds_op_desc_t(const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const int intra_op_parallelism, + const bool validate, + const int64_t file_offset) + : io_op_desc_t(read_op, buffer, fd, filename, intra_op_parallelism, validate, file_offset) +{ + _contiguous_buffer = _buffer.contiguous(); + const int64_t device = _buffer.get_device(); + check_cudaruntimecall(cudaSetDevice(device)); + _base_ptr = _find_base_ptr(device, (char*)_contiguous_buffer.data_ptr()); + + _safe_handle_register(fd, _cf_descr, _cf_handle); +} + +char* gds_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } + +void gds_op_desc_t::finish() { cuFileHandleDeregister(_cf_handle); } + +void gds_op_desc_t::validate() +{ + check_cudaruntimecall(cudaSetDevice(_buffer.get_device())); + const auto cpu_buffer = _buffer.to(torch::kCPU); + const auto num_io_bytes = static_cast(_contiguous_buffer.nbytes()); + validate_aio_operation( + _read_op, _filename.c_str(), (char*)(cpu_buffer.data_ptr()), num_io_bytes); +} + +void gds_op_desc_t::run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config) +{ + assert(tid < _intra_op_parallelism); + check_cudaruntimecall(cudaSetDevice(_buffer.get_device())); + const auto buf_offset = data_ptr() + (_num_bytes_per_thread * tid) - (char*)_base_ptr; + const auto tid_file_offset = _file_offset + (_num_bytes_per_thread * tid); + + if (_read_op) { + auto ret = + cuFileRead(_cf_handle, _base_ptr, _num_bytes_per_thread, tid_file_offset, buf_offset); + if (ret < 0) { _report_error(ret, errno, tid_file_offset); } + } else { + auto ret = + cuFileWrite(_cf_handle, _base_ptr, _num_bytes_per_thread, tid_file_offset, buf_offset); + if (ret < 0) { _report_error(ret, errno, tid_file_offset); } + } +} + +void gds_op_desc_t::_report_error(const ssize_t return_code, + const int error_num, + const off_t offset) +{ + const auto op_string = _read_op ? "read failed with " : "write failed with "; + const auto error_string = IS_CUFILE_ERR(return_code) ? "cuFile error: " : "posix error: "; + const auto error_code = IS_CUFILE_ERR(return_code) ? cuFileGetErrorString(return_code) + : cuFileGetErrorString(error_num); + std::cerr << op_string << error_string << error_code << " return code = " << return_code + << " filename = " << _filename << " num bytes = " << _num_bytes_per_thread + << " offset = " << offset << std::endl; + exit(EXIT_FAILURE); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h new file mode 100644 index 0000000000000000000000000000000000000000..fe2d3cafb8ef7e706330663bc49a559b7b7cc343 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include +#include + +#include "deepspeed_aio_op_desc.h" +#include "deepspeed_gds_utils.h" + +struct gds_op_desc_t : io_op_desc_t { + CUfileDescr_t _cf_descr; + CUfileHandle_t _cf_handle; + void* _base_ptr; + + gds_op_desc_t(const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const int intra_op_parallelism, + const bool validate, + const int64_t file_offset); + + void run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config); + + char* data_ptr() const; + + void validate(); + + void finish(); + + void _report_error(const ssize_t return_code, const int error_num, const off_t offset); + + static void add_buffer_to_registry(const torch::Tensor& buffer); + + static void remove_buffer_from_registry(const torch::Tensor& buffer); +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..12b014d909880da286ea76b264a72ee15d05d2a7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +// CUDA/cuFile includes +#include +#include +#include "cufile.h" + +// Macro for checking cuda errors following a cuda launch or api call +#define cudaCheckError() \ + { \ + cudaError_t e = cudaGetLastError(); \ + if (e != cudaSuccess) { \ + printf("Cuda failure %s:%d: '%s'\n", __FILE__, __LINE__, cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ + } + +#define check_cudadrivercall(fn) \ + do { \ + CUresult res = fn; \ + if (res != CUDA_SUCCESS) { \ + const char* str = nullptr; \ + cuGetErrorName(res, &str); \ + std::cerr << "cuda driver api call failed " << #fn << " res : " << res << ", " \ + << __LINE__ << ":" << str << std::endl; \ + std::cerr << "EXITING program!!!" << std::endl; \ + exit(1); \ + } \ + } while (0) + +#define check_cudaruntimecall(fn) \ + do { \ + cudaError_t res = fn; \ + if (res != cudaSuccess) { \ + const char* str = cudaGetErrorName(res); \ + std::cerr << "cuda runtime api call failed " << #fn << __LINE__ << ":" << str \ + << std::endl; \ + std::cerr << "EXITING program!!!" << std::endl; \ + exit(1); \ + } \ + } while (0) + +#define check_cuFileCall(fn, api_msg) \ + do { \ + CUfileError_t status = fn; \ + if (status.err != CU_FILE_SUCCESS) { \ + std::cout << api_msg << " failed with error " << CUFILE_ERRSTR(status.err) \ + << std::endl; \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +// +// cuda driver error description +// +static inline const char* GetCuErrorString(CUresult curesult) +{ + const char* descp; + if (cuGetErrorName(curesult, &descp) != CUDA_SUCCESS) descp = "unknown cuda error"; + return descp; +} + +// +// cuFile APIs return both cuFile specific error codes as well as POSIX error codes +// for ease, the below template can be used for getting the error description depending +// on its type. + +// POSIX +template ::value, std::nullptr_t>::type = nullptr> +std::string cuFileGetErrorString(T status) +{ + status = std::abs(status); + return IS_CUFILE_ERR(status) ? std::string(CUFILE_ERRSTR(status)) + : std::string(std::strerror(status)); +} + +// CUfileError_t +template ::value, std::nullptr_t>::type = nullptr> +std::string cuFileGetErrorString(T status) +{ + std::string errStr = cuFileGetErrorString(static_cast(status.err)); + if (IS_CUDA_ERR(status)) errStr.append(".").append(GetCuErrorString(status.cu_err)); + return errStr; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp new file mode 100644 index 0000000000000000000000000000000000000000..4d9dc5445043229ede5ef7c9c2daf7f7521d7bd0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp @@ -0,0 +1,118 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* + GPUDirect Storage functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_py_gds_handle.h" +#include +#include "deepspeed_gds_op.h" + +using namespace std; + +int deepspeed_gds_handle_t::s_cuFile_init = 0; + +deepspeed_gds_handle_t::deepspeed_gds_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism) + : deepspeed_io_handle_t(block_size, queue_depth, single_submit, overlap_events, 1), + _intra_gds_op_parallelism(intra_op_parallelism) +{ + _init_cuFile(block_size, queue_depth); +} + +deepspeed_gds_handle_t::~deepspeed_gds_handle_t() { _close_cuFile(); } + +const int deepspeed_gds_handle_t::get_intra_op_parallelism() const +{ + return _intra_gds_op_parallelism; +} + +void deepspeed_gds_handle_t::_init_cuFile(const int block_size, const int queue_depth) +{ + if (deepspeed_gds_handle_t::s_cuFile_init == 0) { + std::string depthStr = std::to_string(queue_depth); + std::string threadsStr = std::to_string(_intra_gds_op_parallelism); + std::string json1 = R"({"execution": {"max_io_queue_depth": )" + depthStr + ", "; + std::string json2 = R"("max_request_parallelism": )" + threadsStr + ", "; + std::string json3 = R"("max_io_threads": )" + threadsStr + ", "; + std::string json4 = R"("parallel_io": true, "min_io_threshold_size_kb": 8192}})"; + std::ofstream outFile("local_cufile.json"); + if (outFile.is_open()) { + outFile << json1 + json2 + json3 + json4; + outFile.close(); + } else { + std::cerr << "Can't open local cufile" << std::endl; + exit(EXIT_FAILURE); + } + // TODO: Address the following issues with this code + // (1) Fix C++14 warning + // (2) Create file in a different location than PWD + // (3) Handle multi-GPU/multi-rank scenarios: should cufile be shared, is per-rank cufile + // safe? + putenv("CUFILE_ENV_PATH_JSON=$PWD/local_cufile.json"); + cuFileDriverOpen(); + cudaCheckError(); + size_t direct_io_size = (size_t)block_size / 1024; + CUfileError_t status = cuFileDriverSetMaxDirectIOSize(direct_io_size); + if (status.err != CU_FILE_SUCCESS) { + std::cerr << "file register error:" << cuFileGetErrorString(status) << std::endl; + exit(EXIT_FAILURE); + } + } + deepspeed_gds_handle_t::s_cuFile_init++; +} + +void deepspeed_gds_handle_t::_close_cuFile() +{ + deepspeed_gds_handle_t::s_cuFile_init--; + if (deepspeed_gds_handle_t::s_cuFile_init == 0) { cuFileDriverClose(); } +} + +torch::Tensor deepspeed_gds_handle_t::new_pinned_device_tensor(const size_t num_elem, + const torch::Tensor& example_tensor) +{ + auto options = torch::TensorOptions().dtype(example_tensor.scalar_type()).device(torch::kCUDA); + auto dev_tensor = torch::empty(num_elem, options); + pin_device_tensor(dev_tensor); + return dev_tensor; +} + +bool deepspeed_gds_handle_t::free_pinned_device_tensor(torch::Tensor& buffer) +{ + unpin_device_tensor(buffer); + return true; +} + +bool deepspeed_gds_handle_t::pin_device_tensor(const torch::Tensor& buffer) +{ + gds_op_desc_t::add_buffer_to_registry(buffer); + return true; +} + +bool deepspeed_gds_handle_t::unpin_device_tensor(const torch::Tensor& buffer) +{ + gds_op_desc_t::remove_buffer_from_registry(buffer); + return true; +} + +std::shared_ptr deepspeed_gds_handle_t::_create_io_op_desc( + const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const int64_t file_offset) +{ + if (buffer.is_cuda()) { + return std::make_shared( + read_op, buffer, fd, filename, _intra_op_parallelism, validate, file_offset); + } + return deepspeed_io_handle_t::_create_io_op_desc( + read_op, buffer, fd, filename, validate, file_offset); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h new file mode 100644 index 0000000000000000000000000000000000000000..f8090fc6fde44a14944a392bd7055e99c87e2bac --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include "deepspeed_py_io_handle.h" + +struct deepspeed_gds_handle_t : deepspeed_io_handle_t { + const int _intra_gds_op_parallelism; + + deepspeed_gds_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism); + + ~deepspeed_gds_handle_t(); + + torch::Tensor new_pinned_device_tensor(const size_t num_elem, + const torch::Tensor& example_tensor); + + bool free_pinned_device_tensor(torch::Tensor&); + + bool pin_device_tensor(const torch::Tensor& buffer); + + bool unpin_device_tensor(const torch::Tensor& buffer); + + void _init_cuFile(const int block_size, const int queue_depth); + + void _close_cuFile(); + + const int get_intra_op_parallelism() const; + + std::shared_ptr _create_io_op_desc(const bool read_op, + const torch::Tensor& buffer, + const int fd, + const char* filename, + const bool validate, + const int64_t file_offset); + + static int s_cuFile_init; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp new file mode 100644 index 0000000000000000000000000000000000000000..bc6327bb38068c9c9abbca9dd1830cbae7da6fcc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp @@ -0,0 +1,141 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include "deepspeed_py_gds_handle.h" +using namespace pybind11::literals; + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + py::class_(m, "gds_handle") + .def(py::init(), + "GDS handle constructor", + "block_size"_a = 1024 * 1024, + "queue_depth"_a = 128, + "single_submit"_a = false, + "overlap_events"_a = false, + "intra_op_parallelism"_a = 1) + + .def("get_block_size", &deepspeed_gds_handle_t::get_block_size) + .def("get_queue_depth", &deepspeed_gds_handle_t::get_queue_depth) + .def("get_single_submit", &deepspeed_gds_handle_t::get_single_submit) + .def("get_overlap_events", &deepspeed_gds_handle_t::get_overlap_events) + .def("get_intra_op_parallelism", &deepspeed_gds_handle_t::get_intra_op_parallelism) + .def("get_alignment", &deepspeed_gds_handle_t::get_alignment) + + .def("read", + &deepspeed_gds_handle_t::read, + "Synchronous and non-parallel file read. Returns count of completed read ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "file_offset"_a = 0) + + .def("write", + &deepspeed_gds_handle_t::write, + "Synchronous and non-parallel file write. Returns count of completed write ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "file_offset"_a = 0) + + .def("pread", + &deepspeed_gds_handle_t::pread, + "Parallel file read with option of parallelism. Returns count of completed read ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "async"_a, + "file_offset"_a = 0) + + .def("pwrite", + &deepspeed_gds_handle_t::pwrite, + "Parallel file write with option of parallelism. Returns count of completed write ops", + "buffer"_a, + "filename"_a, + "validate"_a, + "async"_a, + "file_offset"_a = 0) + + .def("sync_pread", + &deepspeed_gds_handle_t::sync_pread, + "Synchrononous parallel file read. Returns count of completed read ops", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def("sync_pwrite", + &deepspeed_gds_handle_t::sync_pwrite, + "Synchronous parallel file write. Returns count of completed write ops", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def("async_pread", + &deepspeed_gds_handle_t::async_pread, + "Asynchronous parallel file read. Returns 0 on success. Returns 0 on success, and " + "following wait() returns count of completed ops.", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def( + "async_pwrite", + py::overload_cast( + &deepspeed_gds_handle_t::async_pwrite), + "Asynchronous parallel file write. Returns 0 on success, and subsequent wait() returns " + "count of completed ops.", + "buffer"_a, + "filename"_a, + "file_offset"_a = 0) + + .def("async_pwrite", + py::overload_cast( + &deepspeed_gds_handle_t::async_pwrite), + "Asynchronous parallel file write using opened python file object.", + "buffer"_a, + "fd"_a, + "file_offset"_a = 0) + + .def("new_cpu_locked_tensor", + &deepspeed_gds_handle_t::new_cpu_locked_tensor, + "Allocate pinned CPU tensor.", + "num_elem"_a, + "example_tenosr"_a) + + .def("free_cpu_locked_tensor", + &deepspeed_gds_handle_t::free_cpu_locked_tensor, + "Free pinned CPU tensor.", + "tensor"_a) + + .def("new_pinned_device_tensor", + &deepspeed_gds_handle_t::new_pinned_device_tensor, + "Allocate pinned device tensor.", + "num_elem"_a, + "example_tenosr"_a) + + .def("free_pinned_device_tensor", + &deepspeed_gds_handle_t::free_pinned_device_tensor, + "Free pinned device tensor.", + "tensor"_a) + + .def("pin_device_tensor", + &deepspeed_gds_handle_t::pin_device_tensor, + "Pin device tensor.", + "tensor"_a) + + .def("unpin_device_tensor", + &deepspeed_gds_handle_t::unpin_device_tensor, + "Unpin device tensor.", + "tensor"_a) + + .def("wait", + &deepspeed_gds_handle_t::wait, + "Wait for (ongoing) asynchronous operations to complete"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/StopWatch.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/StopWatch.h new file mode 100644 index 0000000000000000000000000000000000000000..42e4eb886eaa086e9eff786601ea8cb784a5b4ed --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/StopWatch.h @@ -0,0 +1,103 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once +#ifdef _WIN32 +#include +#else +#include +#endif + +#ifdef _WIN32 + +class Stopwatch { +private: + double m_total_time; + LARGE_INTEGER m_start_time; + +public: + Stopwatch() { m_total_time = 0.0; } + + ~Stopwatch() {} + + void Reset() { m_total_time = 0.0; } + + void Start() { QueryPerformanceCounter(&m_start_time); } + + void Restart() + { + m_total_time = 0.0; + QueryPerformanceCounter(&m_start_time); + } + + void Stop() + { + LARGE_INTEGER frequency; + LARGE_INTEGER stop_time; + QueryPerformanceFrequency(&frequency); + QueryPerformanceCounter(&stop_time); + m_total_time += + ((double)(stop_time.QuadPart - m_start_time.QuadPart) / (double)frequency.QuadPart); + } + + double GetTimeInSeconds() { return m_total_time; } +}; + +#else + +class Stopwatch { +private: + double m_total_time; + struct timespec m_start_time; + bool m_is_started; + +public: + Stopwatch() + { + m_total_time = 0.0; + m_is_started = false; + } + + ~Stopwatch() {} + + void Reset() { m_total_time = 0.0; } + + void Start() + { + clock_gettime(CLOCK_MONOTONIC, &m_start_time); + m_is_started = true; + } + + void Restart() + { + m_total_time = 0.0; + clock_gettime(CLOCK_MONOTONIC, &m_start_time); + m_is_started = true; + } + + void Stop() + { + if (m_is_started) { + m_is_started = false; + + struct timespec end_time; + clock_gettime(CLOCK_MONOTONIC, &end_time); + + m_total_time += (double)(end_time.tv_sec - m_start_time.tv_sec) + + (double)(end_time.tv_nsec - m_start_time.tv_nsec) / 1e9; + } + } + + double GetTimeInSeconds() + { + if (m_is_started) { + Stop(); + Start(); + } + return m_total_time; + } +}; + +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/Timer.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/Timer.h new file mode 100644 index 0000000000000000000000000000000000000000..b67e04171dafb38151f2cfadaa1a800a84a21667 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/Timer.h @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#ifndef __TIMER_H__ +#define __TIMER_H__ + +#include +#include +#include "cuda.h" + +class GPUTimer { + cudaEvent_t start, stop; + +public: + GPUTimer() + { + cudaEventCreate(&start); + cudaEventCreate(&stop); + } + ~GPUTimer() + { + cudaEventDestroy(start); + cudaEventDestroy(stop); + } + inline void Record() { cudaEventRecord(start); } + inline void Elapsed(float& time_elapsed) + { + cudaEventRecord(stop); + cudaEventSynchronize(stop); + cudaEventElapsedTime(&time_elapsed, start, stop); + } +}; + +class CPUTimer { + std::chrono::high_resolution_clock::time_point start; + +public: + CPUTimer() : start(std::chrono::high_resolution_clock::now()) {} + inline void Reset() { start = std::chrono::high_resolution_clock::now(); } + inline float Elapsed() + { + auto temp = start; + start = std::chrono::high_resolution_clock::now(); + return (float)(std::chrono::duration_cast(start - temp).count() / + 1e3); + } +}; + +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/activation_type.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/activation_type.h new file mode 100644 index 0000000000000000000000000000000000000000..a44921d5d650d3fb4d68522e95b82cda08a1ad60 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/activation_type.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +enum ActivationType { + GELU = 0, + RELU = 1, + SILU = 2, + GEGLU = 3, + ReGLU = 4, + SiGLU = 5, + IDENTITY = 6, + InvalidType = -1 +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/compat.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/compat.h new file mode 100644 index 0000000000000000000000000000000000000000..6d54446d472e9dfa6141436f55e0e623be48acb4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/compat.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#ifndef TORCH_CHECK +#define TORCH_CHECK AT_CHECK +#endif + +#ifdef VERSION_GE_1_3 +#define DATA_PTR data_ptr +#else +#define DATA_PTR data +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/context.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/context.h new file mode 100644 index 0000000000000000000000000000000000000000..cd80f8fbeebe1450a52894b431e3229200b7ff52 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/context.h @@ -0,0 +1,180 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include +#include "cublas_v2.h" +#include "cuda.h" +#include "curand.h" +#include "gemm_test.h" + +#define WARP_SIZE 32 + +#define CUDA_CHECK(callstr) \ + { \ + cudaError_t error_code = callstr; \ + if (error_code != cudaSuccess) { \ + std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ + assert(0); \ + } \ + } + +#define CUDA_1D_KERNEL_LOOP(i, n) \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) + +#define CUDA_2D_KERNEL_LOOP(i, n, j, m) \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) \ + for (size_t j = blockIdx.y * blockDim.y + threadIdx.y; j < (m); j += blockDim.y * gridDim.y) + +#define DS_CUDA_NUM_THREADS 512 +#define DS_MAXIMUM_NUM_BLOCKS 262144 + +inline int DS_GET_BLOCKS(const int N) +{ + return (std::max)( + (std::min)((N + DS_CUDA_NUM_THREADS - 1) / DS_CUDA_NUM_THREADS, DS_MAXIMUM_NUM_BLOCKS), + // Use at least 1 block, since CUDA does not allow empty block + 1); +} + +class TrainingContext { +public: + TrainingContext() : _workspace(nullptr), _seed(42), _curr_offset(0) + { + curandCreateGenerator(&_gen, CURAND_RNG_PSEUDO_DEFAULT); + curandSetPseudoRandomGeneratorSeed(_gen, 123); + cublasStatus_t stat = cublasCreate(&_cublasHandle); + if (stat != CUBLAS_STATUS_SUCCESS) { + // It would be nice to use cublasGetStatusName and + // cublasGetStatusString, but they were only added in CUDA 11.4.2. + auto message = std::string("Failed to create cublas handle: cublasStatus_t was ") + + std::to_string(stat); + std::cerr << message << std::endl; + throw std::runtime_error(message); + } + } + + virtual ~TrainingContext() + { + cublasDestroy(_cublasHandle); + cudaFree(_workspace); + } + + static TrainingContext& Instance() + { + static TrainingContext _ctx; + return _ctx; + } + + void SetWorkSpace(void* workspace) + { + if (!workspace) { throw std::runtime_error("Workspace is null."); } + _workspace = workspace; + } + + void* GetWorkSpace() { return _workspace; } + + curandGenerator_t& GetRandGenerator() { return _gen; } + + cudaStream_t GetCurrentStream() + { + // get current pytorch stream. + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + return stream; + } + + cudaStream_t GetNewStream() { return at::cuda::getStreamFromPool(); } + + cublasHandle_t GetCublasHandle() { return _cublasHandle; } + + std::pair IncrementOffset(uint64_t offset_inc) + { + uint64_t offset = _curr_offset; + _curr_offset += offset_inc; + return std::pair(_seed, offset); + } + + void SetSeed(uint64_t new_seed) { _seed = new_seed; } + + void TestGemmFP16(bool test_gemm, int batch_size, int seq_len, int head_num, int size_per_head) + { + // avoid rerun. + if (_gemm_algos.size() > 0) return; + + if (test_gemm) { + cublasHandle_t handle = GetCublasHandle(); + + std::unique_ptr> test_qkv_fw( + new GemmTest<__half>(batch_size * seq_len, // M + head_num * size_per_head, // N + head_num * size_per_head, // K + CUBLAS_OP_T, + CUBLAS_OP_N, + handle)); + + std::unique_ptr> test_inter( + new GemmTest<__half>(batch_size * seq_len, // M + 4 * head_num * size_per_head, // N + head_num * size_per_head, // K + CUBLAS_OP_T, + CUBLAS_OP_N, + handle)); + + std::unique_ptr> test_output( + new GemmTest<__half>(batch_size * seq_len, // M + head_num * size_per_head, // N + 4 * head_num * size_per_head, // K + CUBLAS_OP_T, + CUBLAS_OP_N, + handle)); + + std::unique_ptr> test_attn_scores( + new StridedGemmTest<__half>(batch_size * head_num, // batch + seq_len, // M + seq_len, // N + size_per_head, // K + CUBLAS_OP_T, + CUBLAS_OP_N, + handle)); + + std::unique_ptr> test_attn_context( + new StridedGemmTest<__half>(batch_size * head_num, // batch + size_per_head, // M + seq_len, // N + seq_len, // K + CUBLAS_OP_N, + CUBLAS_OP_N, + handle)); + + _gemm_algos.push_back(test_qkv_fw->TestAlgo(100)); + _gemm_algos.push_back(test_inter->TestAlgo(100)); + _gemm_algos.push_back(test_output->TestAlgo(100)); + _gemm_algos.push_back(test_attn_scores->TestAlgo(100)); + _gemm_algos.push_back(test_attn_context->TestAlgo(100)); + } else { + // Use default algo. + _gemm_algos.push_back(std::array({99, 99, 99})); + _gemm_algos.push_back(std::array({99, 99, 99})); + _gemm_algos.push_back(std::array({99, 99, 99})); + _gemm_algos.push_back(std::array({99, 99, 99})); + _gemm_algos.push_back(std::array({99, 99, 99})); + } + } + + const std::vector>& GetGemmAlgos() const { return _gemm_algos; } + +private: + curandGenerator_t _gen; + cublasHandle_t _cublasHandle; + void* _workspace; + uint64_t _seed; + uint64_t _curr_offset; + std::vector> _gemm_algos; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/conversion_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/conversion_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..3a90a3e91ddf7d38d943e2b95a6f731796369e97 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/conversion_utils.h @@ -0,0 +1,640 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "ds_kernel_utils.h" + +#include + +#ifdef BF16_AVAILABLE +#include +#endif + +namespace conversion { + +// Basic primitive for constructing conversions +template +DS_D_INLINE TO to(FROM val) +{ + return to(val); +} + +// Specializations + +/********************* Identity Conversions *********************/ +/* +Identity conversions are useful in templated functions where we might have +a fixed destination type. For example, I might have a kernel that accepts +__half, __nv_bfloat16, and float but always want to do the core computation +at floating point: + +T mem_value = input[idx]; +float compute_value = conversion::to(mem_value); + +In practice, we should be able to elide the second template parameter: +float compute_val = conversion::to(mem_value); + +In this case, we need an implementation to handle the T = float case + +NOTE: The type inferencing system appears to be unable to handle inferring the first +template parameter, even in the trivial case. +*/ + +// Floating point types +template <> +DS_D_INLINE double to(double val) +{ + return val; +} +template <> +DS_D_INLINE float to(float val) +{ + return val; +} +template <> +DS_D_INLINE __half to(__half val) +{ + return val; +} +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE __nv_bfloat16 to(__nv_bfloat16 val) +{ + return val; +} +#endif + +// Integer types +template <> +DS_D_INLINE int8_t to(int8_t val) +{ + return val; +} +template <> +DS_D_INLINE uint8_t to(uint8_t val) +{ + return val; +} +template <> +DS_D_INLINE int16_t to(int16_t val) +{ + return val; +} +template <> +DS_D_INLINE uint16_t to(uint16_t val) +{ + return val; +} +template <> +DS_D_INLINE int32_t to(int32_t val) +{ + return val; +} +template <> +DS_D_INLINE uint32_t to(uint32_t val) +{ + return val; +} +template <> +DS_D_INLINE int64_t to(int64_t val) +{ + return val; +} +template <> +DS_D_INLINE uint64_t to(uint64_t val) +{ + return val; +} + +// TODO: evaluate if we want bools + +/********************* To Double Conversions *********************/ + +// * to double variants + +// Would normally like to not use C cast, but this is an important enough conversion +// to keep +template <> +DS_D_INLINE double to(float val) +{ +#ifdef PTX_AVAILABLE + double ret_val; + asm("ctv.rn.f64.f32 %0, %1;\n" : "=d"(ret_val) : "f"(val)); + return ret_val; +#else + return double(val); +#endif +} +// Note: there is a CVT instruction for __half -> double, but there's no inline interface +// for passing a single half value +template <> +DS_D_INLINE double to(__half val) +{ + return to(__half2float(val)); +} +template <> +DS_D_INLINE double to(int64_t val) +{ + return __ll2double_rn(val); +} +template <> +DS_D_INLINE double to(int32_t val) +{ + return __int2double_rn(val); +} +template <> +DS_D_INLINE double to(int16_t val) +{ + return __int2double_rn(val); +} +template <> +DS_D_INLINE double to(int8_t val) +{ + return __int2double_rn(val); +} +template <> +DS_D_INLINE double to(uint64_t val) +{ + return __ull2double_rn(val); +} +template <> +DS_D_INLINE double to(uint32_t val) +{ + return __uint2double_rn(val); +} +template <> +DS_D_INLINE double to(uint16_t val) +{ + return __uint2double_rn(val); +} +template <> +DS_D_INLINE double to(uint8_t val) +{ + return __uint2double_rn(val); +} + +// Same applies here +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE double to(__nv_bfloat16 val) +{ + return to(__bfloat162float(val)); +} +#endif + +/********************* To Float Conversions *********************/ + +template <> +DS_D_INLINE float to(double val) +{ + return __double2float_rn(val); +} +template <> +DS_D_INLINE float to(__half val) +{ + return __half2float(val); +} +template <> +DS_D_INLINE float to(int64_t val) +{ + return __ll2float_rn(val); +} +template <> +DS_D_INLINE float to(int32_t val) +{ + return __int2float_rn(val); +} +template <> +DS_D_INLINE float to(int16_t val) +{ + return __int2float_rn(val); +} +template <> +DS_D_INLINE float to(int8_t val) +{ + return __int2float_rn(val); +} +template <> +DS_D_INLINE float to(uint64_t val) +{ + return __ull2float_rn(val); +} +template <> +DS_D_INLINE float to(uint32_t val) +{ + return __uint2float_rn(val); +} +template <> +DS_D_INLINE float to(uint16_t val) +{ + return __uint2float_rn(val); +} +template <> +DS_D_INLINE float to(uint8_t val) +{ + return __uint2float_rn(val); +} + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE float to(__nv_bfloat16 val) +{ + return __bfloat162float(val); +} +#endif + +/********************* To Float2 Conversions *********************/ +template <> +DS_D_INLINE float2 to(__half2 val) +{ + return __half22float2(val); +} + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE float2 to(__nv_bfloat162 val) +{ + return __bfloat1622float2(val); +} +#endif + +/********************* To Half Conversions *********************/ +template <> +DS_D_INLINE __half to(double val) +{ +#ifdef __HIP_PLATFORM_AMD__ + float val_f = __double2float_rn(val); + return __float2half(val_f); +#else + return __double2half(val); +#endif +} +template <> +DS_D_INLINE __half to(float val) +{ + return __float2half(val); +} +template <> +DS_D_INLINE __half to(int64_t val) +{ + return __ll2half_rn(val); +} +template <> +DS_D_INLINE __half to(int32_t val) +{ + return __int2half_rn(val); +} +template <> +DS_D_INLINE __half to(int16_t val) +{ + return __short2half_rn(val); +} +template <> +DS_D_INLINE __half to(int8_t val) +{ + return __int2half_rn(val); +} +template <> +DS_D_INLINE __half to(uint64_t val) +{ + return __ull2half_rn(val); +} +template <> +DS_D_INLINE __half to(uint32_t val) +{ + return __uint2half_rn(val); +} +template <> +DS_D_INLINE __half to(uint16_t val) +{ + return __ushort2half_rn(val); +} +template <> +DS_D_INLINE __half to(uint8_t val) +{ + return __uint2half_rn(val); +} + +#ifdef BF16_AVAILABLE +// No direct conversion +template <> +DS_D_INLINE __half to(__nv_bfloat16 val) +{ + return to<__half>(to(val)); +} +#endif + +/********************* To Half2 Conversions *********************/ +template <> +DS_D_INLINE __half2 to(float2 val) +{ + return __float22half2_rn(val); +} +template <> +DS_D_INLINE __half2 to(float val) +{ + return __float2half2_rn(val); +} + +#ifdef BF16_AVAILABLE +// No direct conversion +template <> +DS_D_INLINE __half2 to(__nv_bfloat162 val) +{ + return to<__half2>(to(val)); +} +#endif + +/********************* To BF16 Conversions *********************/ +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE __nv_bfloat16 to(double val) +{ + return __double2bfloat16(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(float val) +{ + return __float2bfloat16(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(int64_t val) +{ + return __ll2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(int32_t val) +{ + return __int2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(int16_t val) +{ + return __short2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(int8_t val) +{ + return __int2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(uint64_t val) +{ + return __ull2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(uint32_t val) +{ + return __uint2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(uint16_t val) +{ + return __ushort2bfloat16_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat16 to(uint8_t val) +{ + return __uint2bfloat16_rn(val); +} +#endif + +/********************* To BF162 Conversions *********************/ +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE __nv_bfloat162 to(float2 val) +{ + return __float22bfloat162_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat162 to(float val) +{ + return __float2bfloat162_rn(val); +} +template <> +DS_D_INLINE __nv_bfloat162 to(__half2 val) +{ + return to<__nv_bfloat162>(to(val)); +} +#endif + +/********************* To INT64_T Conversions *********************/ +template <> +DS_D_INLINE int64_t to(double val) +{ + return __double2ll_rn(val); +} +template <> +DS_D_INLINE int64_t to(float val) +{ + return __float2ll_rn(val); +} +template <> +DS_D_INLINE int64_t to(__half val) +{ + return __half2ll_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE int64_t to(__nv_bfloat16 val) +{ + return __bfloat162ll_rn(val); +} +#endif + +/********************* To INT32_T Conversions *********************/ +template <> +DS_D_INLINE int32_t to(double val) +{ + return __double2int_rn(val); +} +template <> +DS_D_INLINE int32_t to(float val) +{ + return __float2int_rn(val); +} +template <> +DS_D_INLINE int32_t to(__half val) +{ + return __half2int_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE int32_t to(__nv_bfloat16 val) +{ + return __bfloat162int_rn(val); +} +#endif + +/********************* To INT16_T Conversions *********************/ +template <> +DS_D_INLINE int16_t to(double val) +{ + return __double2int_rn(val); +} +template <> +DS_D_INLINE int16_t to(float val) +{ + return __float2int_rn(val); +} +template <> +DS_D_INLINE int16_t to(__half val) +{ + return __half2int_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE int16_t to(__nv_bfloat16 val) +{ + return __bfloat162int_rn(val); +} +#endif + +/********************* To INT8_T Conversions *********************/ +template <> +DS_D_INLINE int8_t to(double val) +{ + return __double2int_rn(val); +} +template <> +DS_D_INLINE int8_t to(float val) +{ + return __float2int_rn(val); +} +template <> +DS_D_INLINE int8_t to(__half val) +{ + return __half2int_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE int8_t to(__nv_bfloat16 val) +{ + return __bfloat162int_rn(val); +} +#endif + +/********************* To UINT64_T Conversions *********************/ +template <> +DS_D_INLINE uint64_t to(double val) +{ + return __double2ull_rn(val); +} +template <> +DS_D_INLINE uint64_t to(float val) +{ + return __float2ull_rn(val); +} +template <> +DS_D_INLINE uint64_t to(__half val) +{ + return __half2ull_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE uint64_t to(__nv_bfloat16 val) +{ + return __bfloat162ull_rn(val); +} +#endif + +/********************* To UINT32_T Conversions *********************/ +template <> +DS_D_INLINE uint32_t to(double val) +{ + return __double2uint_rn(val); +} +template <> +DS_D_INLINE uint32_t to(float val) +{ + return __float2uint_rn(val); +} +template <> +DS_D_INLINE uint32_t to(__half val) +{ + return __half2uint_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE uint32_t to(__nv_bfloat16 val) +{ + return __bfloat162uint_rn(val); +} +#endif + +/********************* To UINT16_T Conversions *********************/ +template <> +DS_D_INLINE uint16_t to(double val) +{ + return __double2uint_rn(val); +} +template <> +DS_D_INLINE uint16_t to(float val) +{ + return __float2uint_rn(val); +} +template <> +DS_D_INLINE uint16_t to(__half val) +{ + return __half2uint_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE uint16_t to(__nv_bfloat16 val) +{ + return __bfloat162uint_rn(val); +} +#endif + +/********************* To UINT8_T Conversions *********************/ +template <> +DS_D_INLINE uint8_t to(double val) +{ + return __double2uint_rn(val); +} +template <> +DS_D_INLINE uint8_t to(float val) +{ + return __float2uint_rn(val); +} +template <> +DS_D_INLINE uint8_t to(__half val) +{ + return __half2uint_rn(val); +} +// No direct support for integer casts at the C++ level and I don't feel they're so important +// to demand an PTX at this time + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE uint8_t to(__nv_bfloat16 val) +{ + return __bfloat162uint_rn(val); +} +#endif + +} // namespace conversion diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adagrad.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adagrad.h new file mode 100644 index 0000000000000000000000000000000000000000..6f500250f033a1fa222bba0aaafef2f2852f1156 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adagrad.h @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include +#include "simd.h" + +#define STEP(SPAN) \ + template \ + void Step_##SPAN(ds_params_precision_t* _params, \ + ds_params_precision_t* grads, \ + ds_state_precision_t* _exp_avg_sq, \ + size_t _param_size); + +class Adagrad_Optimizer { +public: + Adagrad_Optimizer(float alpha = 1e-2, float eps = 1e-8, float weight_decay = 0) + : _alpha(alpha), _eps(eps), _weight_decay(weight_decay) + { + } + ~Adagrad_Optimizer() {} +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg_sq, + size_t param_size); +#endif + STEP(1) + STEP(4) + STEP(8) + inline void IncrementStep(size_t step) + { + _step++; + if (_step != step) { _step = step; } + } + inline void update_state(float lr, float epsilon, float weight_decay) + { + _alpha = lr; + _eps = epsilon; + _weight_decay = weight_decay; + } + +private: + float _alpha; + float _eps; + float _weight_decay; + + float _betta1_t; + float _betta2_t; + size_t _step; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, + ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ +#if !defined(__AVX512__) + if (std::is_same_v || + std::is_same_v) { + return; + } +#endif + size_t new_rounded_size = 0; + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + AVX_Data weight_decay4; + if (_weight_decay > 0) weight_decay4.data = SIMD_SET(_weight_decay); + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, grads + i); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i); + + if (_weight_decay > 0) { simd_fma(grad_4, param_4, weight_decay4, grad_4); } + + simd_fma(variance_4, grad_4, grad_4, variance_4); + simd_sqrt(grad_4, variance_4); + simd_add(grad_4, grad_4, eps_4); + simd_div(grad_4, momentum_4, grad_4); + simd_fma(param_4, grad_4, step_size_4, param_4); + + simd_store(_params + i, param_4); + simd_store(_exp_avg_sq + i, variance_4); + } + } + *rounded_size = new_rounded_size; +} +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adam.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adam.h new file mode 100644 index 0000000000000000000000000000000000000000..a7db6fda370543823fe9ac90b90d66f6e7a002ae --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adam.h @@ -0,0 +1,220 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include +#include "simd.h" + +#define STEP(SPAN) \ + template \ + void Step_##SPAN(ds_params_precision_t* _params, \ + ds_params_precision_t* grads, \ + ds_state_precision_t* _exp_avg, \ + ds_state_precision_t* _exp_avg_sq, \ + size_t _param_size); + +class Adam_Optimizer { +public: + Adam_Optimizer(float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float eps = 1e-8, + float weight_decay = 0, + bool adamw_mode = true) + : _alpha(alpha), + _betta1(betta1), + _betta2(betta2), + _eps(eps), + _weight_decay(weight_decay), + _betta1_t(1.0), + _betta2_t(1.0), + _step(0), + _adamw_mode(adamw_mode) + { + } + ~Adam_Optimizer() {} + +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + ds_state_precision_t* _exp_avg_sq, + size_t param_size); +#endif + STEP(1) + STEP(4) + STEP(8) + inline void IncrementStep(size_t step, float beta1, float beta2) + { + if (beta1 != _betta1 || beta2 != _betta2) { + _step = step; + _betta1 = beta1; + _betta2 = beta2; + _betta1_t = std::pow(_betta1, step); + _betta2_t = std::pow(_betta2, step); + } else { + _step++; + if (_step != step) { + _betta1_t = std::pow(_betta1, step); + _betta2_t = std::pow(_betta2, step); + _step = step; + } else { + _betta1_t *= _betta1; + _betta2_t *= _betta2; + } + } + } + inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) + { + _alpha = lr; + _eps = epsilon; + _weight_decay = weight_decay; + + _bias_correction1 = 1.0f; + _bias_correction2 = 1.0f; + if (bias_correction == 1) { + _bias_correction1 = 1 - _betta1_t; + _bias_correction2 = 1 / sqrt(1 - _betta2_t); + } + } + +private: + float _alpha; + float _betta1; + float _betta2; + float _eps; + float _weight_decay; + + float _betta1_t; + float _betta2_t; + size_t _step; + + float _bias_correction1; + float _bias_correction2; + + bool _adamw_mode; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Adam_Optimizer::Step_AVX(size_t* rounded_size, + ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + ds_state_precision_t* _exp_avg_sq, + size_t _param_size) +{ +#if !defined(__AVX512__) + if (std::is_same_v || + std::is_same_v) { + return; + } +#endif + size_t new_rounded_size = 0; + + AVX_Data betta1_4; + betta1_4.data = SIMD_SET(_betta1); + AVX_Data betta2_4; + betta2_4.data = SIMD_SET(_betta2); + + float betta1_minus1 = 1 - _betta1; + float betta2_minus1 = 1 - _betta2; + AVX_Data betta1_minus1_4; + betta1_minus1_4.data = SIMD_SET(betta1_minus1); + AVX_Data betta2_minus1_4; + betta2_minus1_4.data = SIMD_SET(betta2_minus1); + + AVX_Data bias2_sqrt; + bias2_sqrt.data = SIMD_SET(_bias_correction2); + + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha / _bias_correction1; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + float w_decay = -1 * _alpha * _weight_decay; + AVX_Data weight_decay4; + if (_weight_decay > 0) + weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, _exp_avg + i); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i); + + if (_weight_decay > 0 && !_adamw_mode) { + simd_fma(grad_4, param_4, weight_decay4, grad_4); + } + + simd_mul(momentum_4, momentum_4, betta1_4); + simd_fma(momentum_4, grad_4, betta1_minus1_4, momentum_4); + simd_mul(variance_4, variance_4, betta2_4); + simd_mul(grad_4, grad_4, grad_4); + simd_fma(variance_4, grad_4, betta2_minus1_4, variance_4); + simd_sqrt(grad_4, variance_4); + simd_fma(grad_4, grad_4, bias2_sqrt, eps_4); + simd_div(grad_4, momentum_4, grad_4); + + if (_weight_decay > 0 && _adamw_mode) { + simd_fma(param_4, param_4, weight_decay4, param_4); + } + + simd_fma(param_4, grad_4, step_size_4, param_4); + + simd_store(_params + i, param_4); + simd_store(_exp_avg + i, momentum_4); + simd_store(_exp_avg_sq + i, variance_4); + } + } + *rounded_size = new_rounded_size; +} +#endif + +int create_adam_optimizer(int optimizer_id, + float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float eps = 1e-8, + float weight_decay = 0, + bool adamw_mode = true, + bool should_log = false); + +int ds_adam_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + bool bias_correction, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq); + +int destroy_adam_optimizer(int optimizer_id); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_lion.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_lion.h new file mode 100644 index 0000000000000000000000000000000000000000..beaf357a3211773ecb9abe023a48ed39b1276362 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_lion.h @@ -0,0 +1,167 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include +#include "simd.h" + +#define STEP(SPAN) \ + template \ + void Step_##SPAN(ds_params_precision_t* _params, \ + ds_params_precision_t* grads, \ + ds_state_precision_t* _exp_avg, \ + size_t _param_size); + +class Lion_Optimizer { +public: + Lion_Optimizer(float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float weight_decay = 0) + : _alpha(alpha), _betta1(betta1), _betta2(betta2), _weight_decay(weight_decay), _step(0) + { + } + ~Lion_Optimizer() {} + +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + size_t param_size); +#endif + STEP(1) + STEP(4) + STEP(8) + + inline void IncrementStep(size_t step, float beta1, float beta2) + { + _step++; + if (_step != step || beta1 != _betta1 || beta2 != _betta2) { + _step = step; + _betta1 = beta1; + _betta2 = beta2; + } + } + inline void update_state(float lr, float weight_decay) + { + _alpha = lr; + _weight_decay = weight_decay; + } + +private: + float _alpha; + float _betta1; + float _betta2; + float _weight_decay; + size_t _step; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Lion_Optimizer::Step_AVX(size_t* rounded_size, + ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + size_t _param_size) +{ +#if !defined(__AVX512__) + if (std::is_same_v || + std::is_same_v) { + return; + } +#endif + size_t new_rounded_size = 0; + + constexpr float neg1 = -1.0f; + AVX_Data neg1_4; + neg1_4.data = SIMD_SET(neg1); + + AVX_Data betta1_4; + betta1_4.data = SIMD_SET(_betta1); + AVX_Data betta2_4; + betta2_4.data = SIMD_SET(_betta2); + + float betta1_minus1 = 1 - _betta1; + float betta2_minus1 = 1 - _betta2; + AVX_Data betta1_minus1_4; + betta1_minus1_4.data = SIMD_SET(betta1_minus1); + AVX_Data betta2_minus1_4; + betta2_minus1_4.data = SIMD_SET(betta2_minus1); + + float step_size = -_alpha; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + float after_decay = 1.0f - _alpha * _weight_decay; + AVX_Data after_decay_4; + if (_weight_decay > 0) after_decay_4.data = SIMD_SET(after_decay); + + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; + +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, _exp_avg + i); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i); + + AVX_Data tmp_4[span]; + + simd_mul(tmp_4, momentum_4, betta1_4); + simd_fma(tmp_4, grad_4, betta1_minus1_4, tmp_4); + // We already used intrinsics, so consider the machine representation fixed. + simd_and(tmp_4, tmp_4, neg1_4); + simd_xor(tmp_4, tmp_4, step_size_4); + if (_weight_decay > 0) { + simd_fma(param_4, param_4, after_decay_4, tmp_4); + } else { + simd_add(param_4, param_4, tmp_4); + } + + simd_mul(momentum_4, momentum_4, betta2_4); + simd_fma(momentum_4, grad_4, betta2_minus1_4, momentum_4); + + simd_store(_params + i, param_4); + simd_store(_exp_avg + i, momentum_4); + } + } + *rounded_size = new_rounded_size; +} +#endif + +int create_lion_optimizer(int optimizer_id, + float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float weight_decay = 0, + bool should_log = false); + +int ds_lion_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float weight_decay, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg); + +int destroy_lion_optimizer(int optimizer_id); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cublas_wrappers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cublas_wrappers.h new file mode 100644 index 0000000000000000000000000000000000000000..2721fb990c7e270e89dbba2dd636587612cb6465 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cublas_wrappers.h @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include +#ifndef __HIP_PLATFORM_AMD__ +#include +#endif +#ifdef __HIP_PLATFORM_AMD__ +#include +#endif +#include +#include + +int cublas_gemm_ex(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo algo = rocblas_gemm_algo_standard); +#else + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT); +#endif + +int cublas_gemm_ex(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const __half* A, + const __half* B, + __half* C, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo algo = rocblas_gemm_algo_standard); +#else + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + +int cublas_strided_batched_gemm(cublasHandle_t handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + cublasOperation_t op_A, + cublasOperation_t op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo algo = rocblas_gemm_algo_standard); +#else + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT); +#endif + +int cublas_strided_batched_gemm(cublasHandle_t handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const __half* A, + const __half* B, + __half* C, + cublasOperation_t op_A, + cublasOperation_t op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo algo = rocblas_gemm_algo_standard); +#else + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/custom_cuda_layers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/custom_cuda_layers.h new file mode 100644 index 0000000000000000000000000000000000000000..21f19749d4cf4741d2264e504d01e74379396e14 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/custom_cuda_layers.h @@ -0,0 +1,324 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "ds_kernel_utils.h" + +#include +#include +#include +#include +#include + +#include "context.h" +#include "cublas_wrappers.h" + +#define CUDA_CHECK(callstr) \ + { \ + cudaError_t error_code = callstr; \ + if (error_code != cudaSuccess) { \ + std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ + assert(0); \ + } \ + } + +#define MAX_THREADS 1024 +#define THREADS 256 + +#define MAX_THREAD_STRIDE 32 +#define TILE_DIM 32 + +// Maximum sequence-length support based on the number of threads (2048) allowed in each block and +// this MAX is 8K For higher sequence length we need to use higher Max, like for 64K : 32 +#define MAX_THREAD_ITERATIONS 8 // Maximum 8K +#define MAX_WARP_NUM 32 + +#define MAX_REGISTERS 256 + +#define MAX_REG 256 + +#define WARP_SIZE_BITS 5 + +// Fused bias add with gelu activation +template +void launch_bias_gelu(const T* input, + const T* bias, + T* output, + int intermediate_size, + int batch_size, + cudaStream_t stream); + +template +void launch_gelu(const T* input, + T* output, + int intermediate_size, + int batch_size, + cudaStream_t stream); + +template +void launch_d_gelu(T* d_output, + const T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream); + +// Custom fused bias add with layer normalization +template +void launch_bias_residual_layer_norm(T* vals, + const T* residual, + const T* gamma, + const T* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + T* vars, + T* means); + +template +void launch_bias_residual_layer_norm(T* vals, + const T* residual, + const T* gamma, + const T* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + T* vars); + +template +void launch_layerNorm_backward_fused_add(const T* out_grad1, + const T* out_grad2, + const T* X_data, + const T* vars, + const T* means, + const T* gamma, + T* gamma_grad, + T* betta_grad, + T* inp_grad, + int batch_size, + int hidden_dim, + cudaStream_t stream[2]); +template +void launch_layerNorm_backward_fused_add(const T* out_grad1, + const T* out_grad2, + const T* vals_hat, + const T* vars, + const T* gamma, + T* gamma_grad, + T* betta_grad, + T* inp_grad, + int batch_size, + int hidden_dim, + cudaStream_t stream[2], + bool invertible = false, + const T* betta = nullptr); + +template +void launch_layerNorm_backward(const T* out_grad, + const T* X_data, + const T* vars, + const T* means, + const T* gamma, + T* gamma_grad, + T* betta_grad, + T* inp_grad, + int batch_size, + int hidden_dim, + cudaStream_t stream[2]); + +template +void launch_layerNorm_backward(const T* out_grad, + const T* vals_hat, + const T* vars, + const T* gamma, + T* gamma_grad, + T* betta_grad, + T* inp_grad, + int batch_size, + int hidden_dim, + cudaStream_t stream[2], + bool invertible = false, + const T* betta = nullptr); + +template +void launch_layerNorm_backward_nreversible(const T* out_grad, + const T* vals, + const T* out_grad_trans, + const T* vals_trans, + const T* means, + const T* vars, + const T* gamma, + T* gamma_grad, + T* betta_grad, + T* inp_grad, + int batch_size, + int hidden_dim, + cudaStream_t stream[2]); + +template +void Transpose(const T* inp_mat, T* out_mat, int rows, int cols, cudaStream_t stream); + +template +void launch_attn_softmax_backward(T* out_grad, + const T* soft_inp, + int batch_size, + int heads, + int seq_length, + cudaStream_t stream); + +template +void launch_attn_softmax_backward_v2(T* out_grad, + const T* soft_inp, + int batch_size, + int heads, + int seq_length, + cudaStream_t stream); + +// Custom softmax with scaling and attention mask addition +template +void launch_attn_softmax(T* vals, + const T* attn_mask, + int batch_size, + int heads, + int sequence_length, + cudaStream_t stream); + +template +void launch_transform_0213(T* output, + const T* vals, + int batch_size, + int seq_length, + int hidden_dim, + int heads, + cudaStream_t stream); + +// Custom bias add +template +void launch_bias_add_transform_0213(T* outputs, + const T* vals, + const T* bias, + int batch_size, + int seq_length, + int hidden_dim, + int heads, + cudaStream_t stream, + int trans_count); + +// 4D transform [0, 1, 2, 3] -> [0, 2, 1, 3] +template +void launch_transform4d_0213(T* out, + const T* in, + int batch_size, + int heads, + int seq_length, + int hidden_dim, + cudaStream_t stream, + int trans_count); + +template +void launch_dropout(T* vals, + const T* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream); + +template +void launch_dropout(T* vals_out, + const T* vals, + uint8_t* mask, + int total_count, + int dim, + float ratio, + cudaStream_t stream, + bool bwd = false); + +template +void launch_dropout(T* out, + const T* vals, + const T* residual, + const T* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream); + +template +void launch_dropout_grad(T* vals, uint8_t* mask, int total_count, float ratio, cudaStream_t stream); + +template +void launch_dropout_grad(T* vals_out, + const T* vals, + uint8_t* mask, + int total_count, + float ratio, + cudaStream_t stream); + +template +void launch_fuse_transpose_bias_kernel(const T* inp, + T* out, + int rows, + int cols, + cudaStream_t stream); + +void launch_token_sort(int32_t* indices, + int layers, + int batch_size, + int reserved_size, + int original_tokens, + cudaStream_t stream); + +template +void launch_gather_tokens(T* retained_tokens, + T* activations, + int32_t* gather_indices, + int32_t batch_size, + int32_t sampled_tokens, + int32_t channels, + int32_t read_batch_stride, + int32_t read_seq_stride, + int32_t write_batch_stride, + int32_t write_seq_stride, + cudaStream_t stream); + +template +void launch_scatter_tokens(T* all_activations, + T* layer_activations, + int32_t* gather_indices, + int32_t batch_size, + int32_t sampled_tokens, + int32_t channels, + int32_t read_batch_stride, + int32_t read_seq_stride, + int32_t write_batch_stride, + int32_t write_seq_stride, + cudaStream_t stream); + +template +void launch_slice_gpt_mask(T* output_mask, + const T* input_mask, + int batch_size, + int truncated_seq_len, + int orig_seq_len, + cudaStream_t stream); + +template +void launch_slice_bert_mask(T* output_mask, + const T* input_mask, + const int32_t* retained_indices, + int32_t layers, + int32_t batch_size, + int32_t truncated_seq_len, + int32_t orig_seq_len, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/deepcompile.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/deepcompile.h new file mode 100644 index 0000000000000000000000000000000000000000..810664267d3127321be309b7a6a3ed575a1c7e19 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/deepcompile.h @@ -0,0 +1,615 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#define USE_C10D_NCCL + +#include +#include + +#include +#include +#include +#include +#include +#include + +#if __has_include() +#include +#else +#include +#endif + +namespace dc { + +template +static bool hasKey(const std::unordered_map& map, const K& key) +{ + return map.find(key) != map.end(); +} + +template +inline std::string to_string(const T& v) +{ + std::stringstream ss; + ss << v; + return ss.str(); +} + +template +size_t productDim(const L& dim) +{ + size_t prod = 1; + for (auto d : dim) { prod *= d; } + return prod; +} + +template +std::string join_as_str(const T& v, const char* delim = ",", const size_t maxlen = 0) +{ + std::stringstream ss; + + if (!v.empty()) { + auto it = v.begin(); + ss << to_string(*it); + it++; + for (; it != v.end(); ++it) { + if (delim) ss << delim; + ss << to_string(*it); + } + } + + std::string s = ss.str(); + if (maxlen > 0 && s.length() > maxlen) { s = s.substr(0, maxlen) + " ..."; } + + return "[" + s + "]"; +} + +template +std::string tensorPtrToString(T* ptr, size_t size, size_t str_len = 100) +{ + std::vector vals; + for (size_t i = 0; i < size; i++) { + vals.push_back(*ptr); + ptr++; + } + return join_as_str(vals, ",", str_len); +} + +std::string tensorPtrToString(void* ptr, + size_t size, + c10::ScalarType datatype, + size_t max_elem = 20, + size_t max_str_len = 100); + +std::string tensorToString(const at::Tensor& t, size_t max_elem = 20, size_t max_str_len = 100); + +std::string tensorDimToString(const at::Tensor& t); + +at::Tensor test_call(at::Tensor param); + +extern c10::intrusive_ptr process_group; +extern c10::intrusive_ptr symm_mem; +extern ncclComm_t nccl_comm; +extern bool use_symm_mem; +extern bool clone_custom_op_output; +extern bool profile; +extern bool pre_div_reduce; + +extern bool sync_before_reduce; // for debugging +extern bool sync_after_reduce; // for debugging +extern bool sync_before_allgather; // for debugging +extern bool sync_after_allgather; // for debugging + +std::vector sizes_to_int_vector(at::IntArrayRef sizes); +void enable_profiling(bool enable); +bool is_profiling(); + +c10::intrusive_ptr getSymmMemWorkspace(int64_t size); +void lazy_init_symm_memory(); +ncclDataType_t get_nccl_data_type(at::ScalarType scalar_type); +void cleanup(); + +class ReduceTask { +public: + ReduceTask(long ds_id, at::Tensor grad, at::Tensor send_buf) + : ds_id_(ds_id), grad_(std::move(grad)), send_buf_(std::move(send_buf)) + { + } + + long getDSId() const { return ds_id_; } + at::Tensor getSendBuf() const { return send_buf_; } + +private: + long ds_id_; + at::Tensor grad_; + at::Tensor send_buf_; +}; + +class ReduceBucket { +public: + ReduceBucket(int64_t size, at::ScalarType scalar_type) : size_(size), scalar_type_(scalar_type) + { + buffer_ = torch::empty({size}, at::TensorOptions().dtype(scalar_type).device(at::kCUDA)); + offset_ = 0; + } + + int64_t getSize() const { return size_; } + int64_t getOffset() const { return offset_; } + at::Tensor getBuffer() const { return buffer_; } + at::ScalarType getScalarType() const { return scalar_type_; } + + void reserve(int64_t size) + { + if (size > size_) { + buffer_ = + torch::empty({size}, at::TensorOptions().dtype(scalar_type_).device(at::kCUDA)); + size_ = size; + } + } + + at::Tensor allocate(int64_t numel) + { + if (offset_ + numel > size_) { + throw std::runtime_error("Buffer size exceeds the reduce bucket size"); + } + + at::Tensor result = buffer_.index({torch::indexing::Slice(offset_, offset_ + numel)}); + offset_ += numel; + return result; + } + + bool shouldFlush(int64_t numel) { return offset_ > 0 && offset_ + numel > size_; } + + void reset() { offset_ = 0; } + +private: + int64_t size_; + int64_t offset_; + at::Tensor buffer_; + at::ScalarType scalar_type_; +}; + +class DoubleBufferedReduceBucket { +public: + DoubleBufferedReduceBucket(int64_t initial_bucket_size, bool enable_double_buffer) + : initial_bucket_size_(initial_bucket_size), enable_double_buffer_(enable_double_buffer) + { + } + + void swap(at::ScalarType scalar_type, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream) + { + assert(hasKey(current_buffer_, scalar_type)); + assert(hasKey(current_buffer_events_, scalar_type)); + + current_buffer_.at(scalar_type)->reset(); + current_buffer_events_.at(scalar_type)->record(rs_stream); + + if (enable_double_buffer_) { + assert(hasKey(shadow_buffer_, scalar_type)); + assert(hasKey(shadow_buffer_events_, scalar_type)); + + auto tmp = current_buffer_.at(scalar_type); + current_buffer_[scalar_type] = shadow_buffer_.at(scalar_type); + shadow_buffer_[scalar_type] = tmp; + + auto tmp_event = current_buffer_events_.at(scalar_type); + current_buffer_events_[scalar_type] = shadow_buffer_events_.at(scalar_type); + shadow_buffer_events_[scalar_type] = tmp_event; + } + } + + std::shared_ptr getBuffer(at::ScalarType scalar_type) + { + if (!hasKey(current_buffer_, scalar_type)) { + current_buffer_[scalar_type] = + std::make_shared(initial_bucket_size_, scalar_type); + current_buffer_events_[scalar_type] = + std::make_shared(cudaEventDisableTiming); + + if (enable_double_buffer_) { + shadow_buffer_[scalar_type] = + std::make_shared(initial_bucket_size_, scalar_type); + shadow_buffer_events_[scalar_type] = + std::make_shared(cudaEventDisableTiming); + } + } + + return current_buffer_.at(scalar_type); + } + + std::shared_ptr getEvent(at::ScalarType scalar_type) + { + assert(hasKey(current_buffer_events_, scalar_type)); + return current_buffer_events_.at(scalar_type); + } + + void clear() + { + current_buffer_.clear(); + shadow_buffer_.clear(); + current_buffer_events_.clear(); + shadow_buffer_events_.clear(); + } + +private: + int64_t initial_bucket_size_; + bool enable_double_buffer_; + std::unordered_map> current_buffer_; + std::unordered_map> shadow_buffer_; + std::unordered_map> current_buffer_events_; + std::unordered_map> shadow_buffer_events_; +}; + +class DSParam { +public: + DSParam(long id, + std::vector ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + bool partitioned, + int64_t offset, // for Z1 + bool persistent // for Z3 + ) + : id_(id), + shape_(std::move(ds_shape)), + ds_tensor_(ds_tensor), + grad_buffer_(grad_buffer), + partitioned_(partitioned), + offset_(offset), + persistent_(persistent), + offload_stream_(at::cuda::getStreamFromPool()), + reload_stream_(at::cuda::getStreamFromPool()) + { + } + + long getId() const { return id_; } + std::vector getShape() const { return shape_; } + at::Tensor getDSTensor() const + { + // If the reload event exists and is complete, return the reloaded tensor (if defined) + if (reload_done_event_) { + if (!reload_done_event_->query()) { + reload_done_event_->block(at::cuda::getCurrentCUDAStream()); + } + if (ds_reload_tensor_.defined()) { return ds_reload_tensor_; } + } + // Otherwise, if an offload event exists, wait for it to complete + if (offload_done_event_) { + if (!offload_done_event_->query()) { + offload_done_event_->block(at::cuda::getCurrentCUDAStream()); + } + } + return ds_tensor_; + } + at::Tensor getGradBuffer() const { return grad_buffer_; } + bool isPartitioned() const { return partitioned_; } + int64_t getOffset() const { return offset_; } + void setPersistent(bool persistent) { persistent_ = persistent; } + bool isPersistent() const { return persistent_; } + + void offload() + { + // If a reloaded tensor exists, offload its data back to ds_tensor_ + if (ds_reload_tensor_.defined()) { + auto comp_stream = at::cuda::getCurrentCUDAStream(); + comp_done_event_ = std::make_shared(cudaEventDisableTiming); + // Record completion and wait on the offload stream + comp_done_event_->record(comp_stream); + comp_done_event_->block(offload_stream_); + offload_done_event_ = std::make_shared(cudaEventDisableTiming); + + { + at::cuda::CUDAStreamGuard guard(offload_stream_); + ds_tensor_.copy_(ds_reload_tensor_, /*non_blocking=*/true); + ds_reload_tensor_.reset(); // Clear the reloaded tensor + offload_done_event_->record(offload_stream_); + } + // Reset the reload event to indicate that no valid reload is present. + if (reload_done_event_) { reload_done_event_.reset(); } + } + } + + void reload() + { + // Reload only if the current ds_tensor_ is on CPU + if (ds_tensor_.device().is_cpu()) { + auto comp_stream = at::cuda::getCurrentCUDAStream(); + comp_done_event_ = std::make_shared(cudaEventDisableTiming); + // Record and wait on the reload stream + comp_done_event_->record(comp_stream); + comp_done_event_->block(reload_stream_); + reload_done_event_ = std::make_shared(cudaEventDisableTiming); + + { + at::cuda::CUDAStreamGuard guard(reload_stream_); + ds_reload_tensor_ = + at::empty_like(ds_tensor_, ds_tensor_.options().device(torch::kCUDA)); + ds_reload_tensor_.copy_(ds_tensor_, /*non_blocking=*/true); + reload_done_event_->record(reload_stream_); + } + // Reset offload_done_event if it exists to clear any stale offload state. + if (offload_done_event_) { offload_done_event_.reset(); } + } + } + +private: + long id_; + std::vector shape_; + at::Tensor ds_tensor_; + at::Tensor ds_reload_tensor_; + at::Tensor grad_buffer_; + bool partitioned_; + int64_t offset_; // for Z1 + bool persistent_; // for Z3 + mutable bool is_reloaded = false; + + at::cuda::CUDAStream offload_stream_; + at::cuda::CUDAStream reload_stream_; + std::shared_ptr comp_done_event_; + std::shared_ptr offload_done_event_; + std::shared_ptr reload_done_event_; +}; + +class DSParamRegistry { +public: + DSParamRegistry() {} + ~DSParamRegistry() {} + + void registerParam(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + bool partitioned, + int64_t offset, // for Z1 + bool persistent // for Z3 + ) + { + grad_buffer.zero_(); + params_.emplace( + ds_id, + DSParam(ds_id, ds_shape, ds_tensor, grad_buffer, partitioned, offset, persistent)); + valid_[ds_id] = false; + } + + void registerGatheredParam(long ds_id, at::Tensor ds_tensor) + { + gathered_params_.emplace(ds_id, ds_tensor); + } + + void unregisterGatheredParam(long ds_id) + { + assert(hasKey(gathered_params_, ds_id)); + gathered_params_.erase(ds_id); + valid_[ds_id] = false; + } + + const std::unordered_map& getParams() const { return params_; } + + const DSParam& getParam(long ds_id) const { return params_.at(ds_id); } + const size_t getNumParams() const { return params_.size(); } + const at::Tensor& getGatheredParam(long ds_id) const + { + assert(hasKey(gathered_params_, ds_id)); + return gathered_params_.at(ds_id); + } + bool hasGatheredParam(long ds_id) const { return hasKey(gathered_params_, ds_id); } + void setPersistent(long ds_id, bool persistent) { params_.at(ds_id).setPersistent(persistent); } + void offload(long ds_id) { params_.at(ds_id).offload(); } + void reload(long ds_id) { params_.at(ds_id).reload(); } + + void setValid(long ds_id, bool valid) { valid_[ds_id] = valid; } + bool isValid(long ds_id) const + { + assert(hasKey(valid_, ds_id)); + return valid_.at(ds_id); + } + +private: + std::unordered_map params_; + std::unordered_map gathered_params_; + std::unordered_map valid_; +}; + +class CustomOpExecutor { +public: + CustomOpExecutor(c10::intrusive_ptr process_group, + std::shared_ptr param_registry, + std::shared_ptr reduce_buckets, + std::vector ds_ids, + ncclComm_t nccl_comm, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream, + bool pre_div_reduce) + : process_group_(process_group), + param_registry_(std::move(param_registry)), + reduce_buckets_(std::move(reduce_buckets)), + ds_ids_(std::move(ds_ids)), + nccl_comm_(nccl_comm), + rs_stream_(rs_stream), + copy_stream_(copy_stream), + pre_div_reduce_(pre_div_reduce) + { + for (long ds_id : ds_ids_) { + has_acc_grad_[ds_id] = false; + + rs_comp_done_events_[ds_id] = + std::make_shared(cudaEventDisableTiming); + rs_copy_done_events_[ds_id] = + std::make_shared(cudaEventDisableTiming); + } + reduce_counter_ = ds_ids_.size(); + } + ~CustomOpExecutor() {} + + virtual void startForward() {} + + virtual void endForward() {} + + virtual void startBackward(bool update) { param_updated_ = update; } + + virtual void endBackward() {} + + virtual at::Tensor reduceGrad(at::Tensor grad_tensor, long ds_id) + { + int world_size = process_group_->getSize(); + const DSParam& param = param_registry_->getParam(ds_id); + const auto scalar_type = grad_tensor.scalar_type(); + std::shared_ptr reduce_bucket = reduce_buckets_->getBuffer(scalar_type); + + auto comp_stream = at::cuda::getCurrentCUDAStream(); + + if (reduce_bucket->shouldFlush(grad_tensor.numel())) { + int rank = process_group_->getRank(); + + flushReduceBucket(scalar_type); + + // reduce_bucket is swapped in flushReduceBucket if double buffering is enabled + reduce_bucket = reduce_buckets_->getBuffer(scalar_type); + } + + if (grad_tensor.numel() > reduce_bucket->getSize()) { + // extend buckets + at::cuda::stream_synchronize(rs_stream_); + reduce_bucket->reserve(grad_tensor.numel()); + } + + at::Tensor reduce_in_buffer = reduce_bucket->allocate(grad_tensor.numel()); + + // This ensures the order of reduce_scatter -> copy + // Without this block, copy may start while reduce_scatter is still running + reduce_buckets_->getEvent(scalar_type)->block(comp_stream); + auto copy_src = grad_tensor.contiguous().view({-1}).detach(); + // keep references to copy src + reduce_tasks_[scalar_type].emplace_back(ds_id, copy_src, reduce_in_buffer); + + // computation must be done before copy + rs_comp_done_events_[ds_id]->record(comp_stream); + rs_comp_done_events_[ds_id]->block(copy_stream_); + { + at::cuda::CUDAStreamGuard guard(copy_stream_); + reduce_in_buffer.copy_(copy_src, true); + rs_copy_done_events_[ds_id]->record(copy_stream_); + } + + reduce_counter_--; + + if (reduce_counter_ == 0) { + flushAllReduceBuckets(); + + reduce_counter_ = ds_ids_.size(); + + // This synchronization ensures all of reduce calls are done before optimizer's step. + at::cuda::stream_synchronize(rs_stream_); + + endBackward(); + } + + return at::Tensor(); + } + + bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); } + +protected: + c10::intrusive_ptr process_group_; + std::shared_ptr param_registry_; + std::shared_ptr reduce_buckets_; + std::vector ds_ids_; + ncclComm_t nccl_comm_; + at::cuda::CUDAStream rs_stream_; + at::cuda::CUDAStream copy_stream_; + + std::unordered_map> rs_comp_done_events_; + std::unordered_map> rs_copy_done_events_; + + size_t reduce_counter_ = 0; + bool param_updated_ = false; + std::unordered_map> reduce_tasks_; + std::unordered_map has_acc_grad_; + bool pre_div_reduce_; + + virtual void flushReduceBucket(at::ScalarType scalar_type) = 0; + + void flushAllReduceBuckets() + { + for (const auto& it : reduce_tasks_) { flushReduceBucket(it.first); } + } + + // Common helper methods for flushReduceBucket implementations + void blockCopyEvents(at::ScalarType scalar_type) + { + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); + copy_done_event->block(rs_stream_); + } + } + + void applyPreDivision(at::ScalarType scalar_type) + { + if (pre_div_reduce_) { + at::cuda::CUDAStreamGuard guard(rs_stream_); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + t.getSendBuf().div_(process_group_->getSize()); + } + } + } + + ncclRedOp_t getReductionOp() const { return pre_div_reduce_ ? ncclSum : ncclAvg; } + + void performCleanup(at::ScalarType scalar_type) + { + reduce_buckets_->swap(scalar_type, rs_stream_, copy_stream_); + + // Prevent grad tensor from being released before the copy is done + auto comp_stream = at::cuda::getCurrentCUDAStream(); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); + copy_done_event->block(comp_stream); + } + reduce_tasks_[scalar_type].clear(); + } +}; + +template +std::shared_ptr getExecutor(long graph_id, + const std::unordered_map>& executors) +{ + assert(hasKey(executors, graph_id)); + if (auto executor = std::dynamic_pointer_cast(executors.at(graph_id))) { return executor; } + throw std::runtime_error("Invalid executor type"); +} + +extern std::shared_ptr param_registry; +extern std::unordered_map> executors; +extern std::shared_ptr reduce_buckets; + +at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id); +at::Tensor reduce_grad_meta(at::Tensor grad_tensor, long graph_id, long ds_id); +void free_tensors(std::vector tensors); +void free_tensors_meta(std::vector tensors); + +void init(c10::intrusive_ptr pg, + int64_t initial_reduce_bucket_size, + bool enable_double_buffer, + bool _use_symm_mem, + bool _clone_custom_op_output, + bool _sync_before_reduce, + bool _sync_after_reduce, + bool _sync_before_allgather, + bool _sync_after_allgather); +void reset(); +void cleanup(); + +void start_forward(); +void end_forward(); +void start_backward(bool update); + +} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dequantization_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dequantization_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..37a7cdf27a014a4eae42b9400ae0c2a12132c449 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dequantization_utils.h @@ -0,0 +1,177 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "quantization.h" +#include "quantization_utils.h" + +namespace cg = cooperative_groups; + +#pragma once + +namespace dequantize { +using Type = quantize::Type; + +template +using Params = quantize::Params; + +constexpr int granularity = quantize::granularity; +using PackedInt4 = quantize::PackedInt4; + +constexpr int h_per_chunk = granularity / sizeof(__half); +constexpr int h2_per_chunk = granularity / sizeof(__half2); + +/* +Device function that reads quantized data from global memory, dequantizes +it, and stores it to global memory. +Template Arguments : + numBits - Number of bits in quantized element. int: 4, 8 + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric + unroll - Number of load steps to internally unroll int + threads - Number of threads to perform dequant int +Function arguments: + global_output - __half pointer in global memory + data - Quantized data in global memory + global_params - Quantization parameters in global memory + elems_per_group - Number of elements in each quantization group + total_elems - Tensor size (note, does not need to be multiple of elems_per_group) +*/ +template +DS_D_INLINE void to_global(__half* global_output, + const int8_t* data, + const float* global_params, + const int elems_per_group, + const int total_elems); + +/* +Device function that quantizes 16 bytes of __half type input data. +Template Arguments : + numBits - Number of bits in quantized element. int : 8 or 4 + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric +Function Arguments : + local_output - Local array to store dequantized data __half* or __half2* + data - Pointer to quantized input data. int8_t* + Params - Parameters for quantization. Params +*/ +template +DS_D_INLINE void chunk(__half2* local_output, const int8_t* data, Params q_params); + +template +DS_D_INLINE void chunk(T* local_output, const int8_t* data, Params q_params); + +/**************** Implementations ******************/ + +template +DS_D_INLINE void chunk(T* local_output, const int8_t* data, Params q_params) +{ + constexpr int32_t num_elems_packed = 8 / numBits; + constexpr int32_t iters = h_per_chunk / num_elems_packed; + +#pragma unroll + for (int i = 0; i < iters; i++) { + if constexpr (num_elems_packed == 1) { + local_output[i] = q_params.template dequantize(data[i]); + } else { + auto accessible_data = *(PackedInt4*)(&data[i]); + local_output[2 * i] = q_params.template dequantize(accessible_data.low); + local_output[2 * i + 1] = q_params.template dequantize(accessible_data.high); + } + } +} + +template +DS_D_INLINE void chunk(__half2* local_output, const int8_t* data, Params q_params) +{ + __half* local_output_cast = reinterpret_cast<__half*>(local_output); + chunk<__half, numBits>(local_output_cast, data, q_params); +} + +template +DS_D_INLINE void _to_global(T* global_output, + const int8_t* data, + const float* global_params, + const int elems_per_group, + const int total_elems) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Load constants + // TODO(cmikeh2): Refactor into functions? + constexpr int load_granularity = (granularity / (sizeof(T))) / (numBits == 8 ? 1 : 2); + constexpr int load_step_stride = load_granularity * threads; + constexpr int load_block_stride = load_step_stride * unroll; + + // Store constants + constexpr int T_per_chunk = granularity / sizeof(T); + constexpr int store_step_stride = T_per_chunk * threads; + constexpr int store_block_stride = store_step_stride * unroll; + + // Load offsets + const int load_block_offset = tb.group_index().x * load_block_stride; + // Note: we can use `load_granularity` since the dtype is `int8_t`. + const int load_thread_offset = tb.thread_index().x * load_granularity; + const int8_t* load_base = data + load_block_offset + load_thread_offset; + + // Store offsets + const int store_block_offset = tb.group_index().x * store_block_stride; + const int store_thread_offset = tb.thread_index().x * T_per_chunk; + const int elem_id_base = store_block_offset + store_thread_offset; + + int8_t local_load_buffer[load_granularity * unroll]; + T local_dequant_buffer[T_per_chunk * unroll]; + + /* + Note: Splitting this loop in half gave about 3-5% performance increase for reasons that aren't + totally clear to me, so this is a deliberately weird code structure. + */ +#pragma unroll + for (int i = 0; i < unroll; i++) { + const int elem_id_iter = elem_id_base + i * store_step_stride; + + if (elem_id_iter < total_elems) { + mem_access::load_global(local_load_buffer + i * load_granularity, + load_base + i * load_step_stride); + } + } + +#pragma unroll + for (int i = 0; i < unroll; i++) { + const int elem_id_iter = elem_id_base + i * store_step_stride; + if (elem_id_iter < total_elems) { + // TODO(cmikeh2): Can we amortize this division? Perform once on the first iteration and + // use indexing math to do division free interpolation of the successive groups? + const int group_index = elem_id_iter / elems_per_group; + Params q_params(global_params, group_index); + + chunk(local_dequant_buffer + i * T_per_chunk, + local_load_buffer + i * load_granularity, + q_params); + mem_access::store_global(global_output + elem_id_iter, + local_dequant_buffer + i * T_per_chunk); + } + } +} + +template +DS_D_INLINE void to_global(T* global_output, + const int8_t* data, + const float* global_params, + const int elems_per_group, + const int total_elems) +{ + if constexpr (numBits == 4 || numBits == 8) { + _to_global( + global_output, data, global_params, elems_per_group, total_elems); + } else if constexpr (numBits == 3) { + // TODO(cmikeh2): Need this implementation + assert(false); + } else { + assert(false); + } +} + +} // namespace dequantize diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dropout.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dropout.h new file mode 100644 index 0000000000000000000000000000000000000000..c00943a97687c3514ff96f63b4b0540805e9b545 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dropout.h @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include + +template +class Dropout { +public: + struct Config { + float ratio; + uint32_t dim; + bool training; + + Config(float r, uint32_t d) : ratio(r), dim(d), training(true) {} + + float RATIO() const { return training ? ratio : 0.0; } + inline void SetDim(uint32_t d) { dim = d; } + }; + + Dropout(const Config& config) : _config(config), _mask(nullptr) {} + + virtual ~Dropout() {} + + void Forward(int bsz, T* out, const T* vals, cudaStream_t stream, bool bwd = false) + { + launch_dropout( + out, vals, _mask, bsz * _config.dim, _config.dim, _config.RATIO(), stream, bwd); + } + + void ForwardWithBias(int bsz, T* vals, const T* bias, cudaStream_t stream) + { + launch_dropout(vals, bias, _mask, bsz, _config.dim, _config.RATIO(), stream); + } + + void ForwardWithBias(int bsz, + T* out, + const T* vals, + const T* residual, + const T* bias, + cudaStream_t stream) + { + launch_dropout( + out, vals, residual, bias, _mask, bsz, _config.dim, _config.RATIO(), stream); + } + + void Backward(int bsz, T* d_vals, cudaStream_t stream) + { + launch_dropout_grad(d_vals, _mask, bsz * _config.dim, _config.RATIO(), stream); + } + + void Backward(int bsz, T* d_vals_out, const T* d_vals, cudaStream_t stream) + { + launch_dropout_grad( + d_vals_out, d_vals, _mask, bsz * _config.dim, _config.RATIO(), stream); + } + + bool HasDropout() const { return _config.RATIO() > 0.0; } + + void SetTrainingMode(bool training) { _config.training = training; } + + void SetMask(uint8_t* mask) + { + if (!mask) { throw std::runtime_error("Dropout mask is null."); } + + _mask = mask; + } + + Config GetConfig() const { return _config; } + + inline void SetDimension(uint32_t dim) { _config.SetDim(dim); } + +private: + uint8_t* _mask; + Config _config; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_kernel_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_kernel_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..f8b16ee6a315f4c2ff5ca72fab077c0534c070e3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_kernel_utils.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Centralized header file for preprocessor macros and constants +used throughout the codebase. +*/ + +#pragma once + +#include +#include + +#ifdef BF16_AVAILABLE +#include +#endif + +#define DS_HD_INLINE __host__ __device__ __forceinline__ +#define DS_D_INLINE __device__ __forceinline__ + +#ifdef __HIP_PLATFORM_AMD__ + +// constexpr variant of warpSize for templating +constexpr int hw_warp_size = ROCM_WAVEFRONT_SIZE; +#define HALF_PRECISION_AVAILABLE = 1 +#include +#include + +#else // !__HIP_PLATFORM_AMD__ + +// constexpr variant of warpSize for templating +constexpr int hw_warp_size = 32; + +#if __CUDA_ARCH__ >= 530 +#define HALF_PRECISION_AVAILABLE = 1 +#define PTX_AVAILABLE +#endif // __CUDA_ARCH__ >= 530 + +#if __CUDA_ARCH__ >= 800 +#define ASYNC_COPY_AVAILABLE +#endif // __CUDA_ARCH__ >= 800 + +#include +#include + +#endif //__HIP_PLATFORM_AMD__ + +inline int next_pow2(const int val) +{ + int rounded_val = val - 1; + rounded_val |= rounded_val >> 1; + rounded_val |= rounded_val >> 2; + rounded_val |= rounded_val >> 4; + rounded_val |= rounded_val >> 8; + return rounded_val + 1; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_transformer_cuda.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_transformer_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..4dbd1d4773385db8e40a9e2f0ac8d6ebf2433c7a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_transformer_cuda.h @@ -0,0 +1,189 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include "cublas_v2.h" +#include "cuda.h" +#include "dropout.h" +#include "feed_forward.h" +#include "gelu.h" +#include "general_kernels.h" +#include "normalize_layer.h" +#include "softmax.h" +#include "strided_batch_gemm.h" + +struct BertGemmAlgos { + int m_gemm_qkv_algo; + int m_gemm_inter_algo; + int m_gemm_output_algo; + int m_gemm_batch1_algo; + int m_gemm_batch2_algo; + + BertGemmAlgos() + : m_gemm_qkv_algo(-1), + m_gemm_inter_algo(-1), + m_gemm_output_algo(-1), + m_gemm_batch1_algo(-1), + m_gemm_batch2_algo(-1) + { + } +}; + +template +class BertTransformerLayer { +public: + BertTransformerLayer(unsigned layer_id, + unsigned batch_size, + unsigned hidden_size, + unsigned num_heads, + unsigned intermediate_size, + unsigned seq_length, + float attn_dropout_ratio, + float hidden_output_dropout_ratio, + float layer_norm_eps, + bool pre_or_postLayerNorm, + const std::vector>& gemm_algos, + bool attn_dropout_checkpoint, + bool normalize_invertible, + bool gelu_checkpoint, + bool stochastic_mode); + + virtual ~BertTransformerLayer(); + + void Forward(unsigned bsz, + const T* input_ptr, + const T* input_mask_ptr, + const T* attn_qkvw_ptr, + const T* attn_qkvb_ptr, + const T* attn_ow_ptr, + const T* attn_ob_ptr, + const T* attn_nw_ptr, + const T* attn_nb_ptr, + const T* inter_w_ptr, + const T* inter_b_ptr, + const T* output_w_ptr, + const T* output_b_ptr, + const T* norm_w_ptr, + const T* norm_b_ptr, + T* out_ptr, + T* inp_norm_ptr, + T* q_tf_ptr, + T* k_tf_ptr, + T* v_tf_ptr, + T* softmax_output_ptr, + T* ctx_bufB_ptr, + T* attn_o_inp_ptr, + T* add_res_ptr, + T* ff1_inp_ptr, + T* gelu_inp_ptr, + T* ff2_inp_ptr); + + void Backward(unsigned bsz, + const T* grad_output_ptr, + const T* input_ptr, + const T* output_ptr, + const T* inp_norm_ptr, + const T* q_tf_ptr, + const T* k_tf_ptr, + const T* v_tf_ptr, + const T* softmax_output_ptr, + const T* ctx_bufB_ptr, + const T* attn_o_inp_ptr, + const T* add_res_ptr, + const T* ff1_inp_ptr, + const T* gelu_inp_ptr, + const T* ff2_inp_ptr, + const T* input_mask_ptr, + const T* attn_qkvw_ptr, + const T* attn_ow_ptr, + const T* attn_nw_ptr, + const T* attn_nb_ptr, + const T* inter_w_ptr, + const T* inter_b_ptr, + const T* output_w_ptr, + const T* norm_w_ptr, + const T* norm_b_ptr, + + T* grad_input_ptr, + T* grad_attn_qkvw_ptr, + T* grad_attn_qkvb_ptr, + T* grad_attn_ow_ptr, + T* grad_attn_ob_ptr, + T* grad_attn_nw_ptr, + T* grad_attn_nb_ptr, + T* grad_inter_w_ptr, + T* grad_inter_b_ptr, + T* grad_output_w_ptr, + T* grad_output_b_ptr, + T* grad_norm_w_ptr, + T* grad_norm_b_ptr); + + void SetIntermediateBuffers(uint8_t* attn_prob_dropout_mask_ptr, + uint8_t* attn_output_dropout_mask_ptr, + uint8_t* layer_output_dropout_mask_ptr, + T* layer_norm_var, + T* layer_norm_mean, + T* attn_layer_norm_var, + T* attn_layer_norm_mean); + + inline unsigned GetBatchSize() const { return _batch_size; } + inline unsigned GetNumHeads() const { return _heads; } + inline unsigned GetSeqLength() const { return _seq_length; } + inline unsigned GetIntermediateSize() const { return _intermediate_size; } + + void SetSeqLength(unsigned seq_len); + inline unsigned GetHiddenSize() const { return _hidden_size; } + void SetTrainingMode(bool training); + inline bool IsTrainingMode() const { return _training; } + inline bool GeluCheckpoint() const { return _gelu_checkpoint; } + +private: + void Initialize(); + size_t getWorkspaceSize(int maxBatchSize) const; + + // Params + unsigned _layer_id; + unsigned _batch_size; + unsigned _hidden_size; + unsigned _heads; + unsigned _size_per_head; + unsigned _intermediate_size; + unsigned _seq_length; + + bool _pre_or_postLayerNorm; + + cublasHandle_t _cublasHandle; + cudaStream_t _stream; + + // layers + FeedForward _qkv_linear; + FeedForward _attn_out_linear; + Normalize_Layer _attn_layer_norm; + Normalize_Layer _layer_norm; + Normalize_Layer* _last_normalize; + FeedForward _ff1, _ff2; + Softmax _softmax; + Gelu _gelu; + Dropout _attn_prob_dropout; + Dropout _attn_output_dropout; + Dropout _layer_output_dropout; + StridedBatchGemm _attn_scores; + StridedBatchGemm _attn_context; + + bool _training; + + // Memory saving flags + bool _attn_dropout_checkpoint; + bool _normalize_invertible; + bool _gelu_checkpoint; + + // High Performance flags + bool _stochastic_mode; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/feed_forward.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/feed_forward.h new file mode 100644 index 0000000000000000000000000000000000000000..d2056403d265e4083efd9da8f7ee9c887e0d9afc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/feed_forward.h @@ -0,0 +1,114 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#ifndef __FEEDFORWARD_H__ +#define __FEEDFORWARD_H__ + +#include +#include +#include +#include "custom_cuda_layers.h" + +template +class FeedForward { +public: + struct Config { + int batchSize, outputSize; + int inputSize; + std::array gemm_algos; + Config(int batch, int outputs, int inputs, const std::array& algos) + : batchSize(batch), outputSize(outputs), inputSize(inputs), gemm_algos(algos) + { + } + }; + + FeedForward(Config config) : config_(config) {} + + ~FeedForward() {} + + void Forward(int bsz, + const T* input_ptr, + const T* weights, + T* out, + cublasHandle_t& _cublasHandle) + { + float alpha = T(1.); + float beta = T(0.); + + cublas_gemm_ex(_cublasHandle, + CUBLAS_OP_T, + CUBLAS_OP_N, + config_.outputSize, + bsz, + config_.inputSize, + &alpha, + &beta, + weights, + input_ptr, + out, +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(config_.gemm_algos[0])); +#else + cublasGemmAlgo_t(config_.gemm_algos[0])); +#endif + } + void Backward(int bsz, + const T* out_grad, + const T* input_ptr, + const T* weights, + T* weights_grad, + T* bias_grad, + cublasHandle_t& _cublasHandle, + cudaStream_t& stream, + T* inp_grad_out = nullptr, + T* out_grad_trans_out = nullptr) + { + float alpha = (T)1.0, beta = (T)0.0; + cublas_gemm_ex(_cublasHandle, + CUBLAS_OP_N, + CUBLAS_OP_T, + config_.inputSize, + config_.outputSize, + bsz, + &alpha, + &beta, + input_ptr, + out_grad, + weights_grad, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(config_.gemm_algos[1])); +#else + cublasGemmAlgo_t(config_.gemm_algos[1])); +#endif + + cublas_gemm_ex(_cublasHandle, + CUBLAS_OP_N, + CUBLAS_OP_N, + config_.inputSize, + bsz, + config_.outputSize, + &alpha, + &beta, + weights, + out_grad, + inp_grad_out, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(config_.gemm_algos[2])); +#else + cublasGemmAlgo_t(config_.gemm_algos[2])); +#endif + + launch_fuse_transpose_bias_kernel(out_grad, bias_grad, bsz, config_.outputSize, stream); + } + +private: + Config config_; +}; + +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gelu.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gelu.h new file mode 100644 index 0000000000000000000000000000000000000000..b75c1510bad65101604dff2aaebfb79d5d435531 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gelu.h @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include "custom_cuda_layers.h" + +template +class Gelu { +public: + struct Config { + uint32_t intermediate_size; + Config(uint32_t inter_size) : intermediate_size(inter_size) {} + }; + + Gelu(const Config& config) : _config(config) {} + + virtual ~Gelu() {} + + void ForwardWithBiasAdd(int bsz, + const T* input_buf, + const T* bias, + T* output, + cudaStream_t stream) + { + launch_bias_gelu(input_buf, bias, output, _config.intermediate_size, bsz, stream); + } + + void Backward(int bsz, T* d_output, const T* input_buf, const T* bias, cudaStream_t stream) + { + launch_d_gelu(d_output, input_buf, bias, _config.intermediate_size, bsz, stream); + } + +private: + Config _config; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gemm_test.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gemm_test.h new file mode 100644 index 0000000000000000000000000000000000000000..de5b55cd3df16f211af75fde895a49b3871e0331 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gemm_test.h @@ -0,0 +1,350 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#ifndef __HIP_PLATFORM_AMD__ +#include +#endif +#ifdef __HIP_PLATFORM_AMD__ +#include +#endif +#include +#include +#include +#include +#include +#include +#include "StopWatch.h" +#include "cublas_wrappers.h" + +template +void check(T result, char const* const func, const char* const file, int const line) +{ + if (result) { + std::cout << (std::string("CUDA runtime error: ") + +file + ":" + std::to_string(line) + + " \n"); + } +} + +#define check_cuda_error(val) check((val), #val, __FILE__, __LINE__) + +template +class GemmTest { +public: + GemmTest(int m, int n, int k, cublasOperation_t ta, cublasOperation_t tb, cublasHandle_t h) + : M(m), N(n), K(k), transa(ta), transb(tb), handle(h) + { + check_cuda_error(cudaMalloc((void**)&A, sizeof(T) * M * K)); + check_cuda_error(cudaMalloc((void**)&B, sizeof(T) * K * N)); + check_cuda_error(cudaMalloc((void**)&C, sizeof(T) * M * N)); + } + + ~GemmTest() + { + check_cuda_error(cudaFree(A)); + check_cuda_error(cudaFree(B)); + check_cuda_error(cudaFree(C)); + } + + std::array TestAlgo(int loops) + { + float alpha = (T)1.0f; + float beta = (T)0.0f; + + int algo_fw = Run(loops, [=](int algo) { + cublas_gemm_ex(handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + N, + M, + K, + &alpha, + &beta, + B, + A, + C, +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + static_cast(algo)); +#else + static_cast(algo)); +#endif + }); + + int algo_bw1 = Run(loops, [=](int algo) { + cublas_gemm_ex(handle, + CUBLAS_OP_N, + CUBLAS_OP_T, + K, + N, + M, + &alpha, + &beta, + A, + C, + B, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + static_cast(algo)); +#else + static_cast(algo)); +#endif + }); + + int algo_bw2 = Run(loops, [=](int algo) { + cublas_gemm_ex(handle, + CUBLAS_OP_N, + CUBLAS_OP_N, + K, + M, + N, + &alpha, + &beta, + B, + C, + A, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + static_cast(algo)); +#else + static_cast(algo)); +#endif + }); + + return std::array({algo_fw, algo_bw1, algo_bw2}); + } + + template + int Run(int loops, Func f) + { + float fast_latency = (std::numeric_limits::max)(); + int fast_algo = 0; + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + for (int algo = (int)rocblas_gemm_algo_standard; algo <= (int)rocblas_gemm_algo_standard; +#elif defined(__HIP_PLATFORM_AMD__) + for (int algo = (int)HIPBLAS_GEMM_DEFAULT; algo <= (int)HIPBLAS_GEMM_DEFAULT; +#else + for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; + algo <= (int)CUBLAS_GEMM_ALGO15_TENSOR_OP; +#endif + algo++) { + int warm_up = 5; + for (int i = 0; i < warm_up; ++i) f(algo); + + cudaDeviceSynchronize(); + Stopwatch timer; + timer.Restart(); + + for (int i = 0; i < loops; ++i) f(algo); + + cudaDeviceSynchronize(); + timer.Stop(); + + float avg_latency = (float)timer.GetTimeInSeconds() * 1000 / loops; + + printf("algo-%d: %.3fms\n", algo, avg_latency); + + if (avg_latency < fast_latency) { + fast_latency = avg_latency; + fast_algo = algo; + } + } + + printf("fast_algo %d: %.3f ms\n", fast_algo, fast_latency); + + return fast_algo; + } + +private: + int M, N, K; + cublasHandle_t handle; + cublasOperation_t transa, transb; + T *A, *B, *C; +}; + +template +class StridedGemmTest { +public: + StridedGemmTest(int b, + int m, + int n, + int k, + cublasOperation_t ta, + cublasOperation_t tb, + cublasHandle_t h) + : bsz(b), M(m), N(n), K(k), transa(ta), transb(tb), handle(h) + { + check_cuda_error(cudaMalloc((void**)&A, sizeof(T) * M * K * bsz)); + check_cuda_error(cudaMalloc((void**)&B, sizeof(T) * K * N * bsz)); + check_cuda_error(cudaMalloc((void**)&C, sizeof(T) * M * N * bsz)); + } + + ~StridedGemmTest() + { + check_cuda_error(cudaFree(A)); + check_cuda_error(cudaFree(B)); + check_cuda_error(cudaFree(C)); + } + + std::array TestAlgo(int loops) + { + float alpha = (T)1.0f; + float beta = (T)0.0f; + + int algo_fw = Run(loops, [=](int algo) { + int stride_a = M * K; + int stride_b = N * K; + int stride_c = M * N; + + cublas_strided_batched_gemm(handle, + M, + N, + K, + &alpha, + &beta, + A, + B, + C, + transa, + transb, + stride_a, + stride_b, + stride_c, + bsz, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + static_cast(algo)); +#else + static_cast(algo)); +#endif + }); + + int algo_bw1 = Run(loops, [=](int algo) { + int mb = (transa == CUBLAS_OP_T ? K : M); + int kb = (transa == CUBLAS_OP_T ? M : K); + + int stride_a = mb * N; + int stride_b = N * kb; + int stride_c = M * K; + + // B need to transpose. + cublasOperation_t op_b = (transb == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); + + // Calculate d_A. + cublas_strided_batched_gemm(handle, + mb, + kb, + N, + &alpha, + &beta, + (transa == CUBLAS_OP_T ? B : C), + (transa == CUBLAS_OP_T ? C : B), + A, + CUBLAS_OP_N, + op_b, + stride_a, + stride_b, + stride_c, + bsz, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + static_cast(algo)); +#else + static_cast(algo)); +#endif + }); + + int algo_bw2 = Run(loops, [=](int algo) { + // A need to transpose. + cublasOperation_t op_a = (transa == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); + + int stride_a = M * K; + int stride_b = M * N; + int stride_c = N * K; + + // Calculate d_B. + cublas_strided_batched_gemm(handle, + K, + N, + M, + &alpha, + &beta, + A, + C, + B, + op_a, + CUBLAS_OP_N, + stride_a, + stride_b, + stride_c, + bsz, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + static_cast(algo)); +#else + static_cast(algo)); +#endif + }); + + return std::array({algo_fw, algo_bw1, algo_bw2}); + } + + template + int Run(int loops, Func f) + { + float fast_latency = (std::numeric_limits::max)(); + int fast_algo = 0; + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + for (int algo = (int)rocblas_gemm_algo_standard; algo <= (int)rocblas_gemm_algo_standard; +#else +#ifdef __HIP_PLATFORM_AMD__ + for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; + algo <= (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; +#else + for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; + algo <= (int)CUBLAS_GEMM_ALGO15_TENSOR_OP; +#endif +#endif + algo++) { + int warm_up = 5; + for (int i = 0; i < warm_up; ++i) f(algo); + + cudaDeviceSynchronize(); + Stopwatch timer; + timer.Restart(); + + for (int i = 0; i < loops; ++i) f(algo); + + cudaDeviceSynchronize(); + timer.Stop(); + + float avg_latency = (float)timer.GetTimeInSeconds() * 1000 / loops; + + printf("algo-%d: %.3fms\n", algo, avg_latency); + + if (avg_latency < fast_latency) { + fast_latency = avg_latency; + fast_algo = algo; + } + } + + printf("fast_algo %d: %.3f ms\n", fast_algo, fast_latency); + + return fast_algo; + } + +private: + int bsz, M, N, K; + cublasHandle_t handle; + cublasOperation_t transa, transb; + T *A, *B, *C; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/general_kernels.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/general_kernels.h new file mode 100644 index 0000000000000000000000000000000000000000..bd621d3c4329c1e12ba5aac25404faaa9d1bca3e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/general_kernels.h @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include + +#ifdef __HIP_PLATFORM_AMD__ +#include +#else +#include +#endif +#include + +#include "context.h" +#include "cublas_wrappers.h" + +#define THREADS 256 +#define TILE_DIM 32 + +#define minus_infinity -1 * std::numeric_limits::infinity() + +#define FINAL_MASK 0xffffffff + +template +void launch_fused_add2(T* out, + const T* inp1, + const T* inp2, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream); + +template +void launch_fused_add4(T* out, + const T* inp1, + const T* inp2, + const T* inp3, + const T* inp4, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream); + +template +void launch_fused_add3(T* out, + const T* inp1, + const T* inp2, + const T* inp3, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/memory_access_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/memory_access_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..cb990b75bbe8460176a431c4ea0c2573b742dffe --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/memory_access_utils.h @@ -0,0 +1,1144 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include "ds_kernel_utils.h" + +/////////////////////////////// Memory Access Utils /////////////////////////////// +namespace mem_access { + +enum class LoadPolicy { + CacheAll, // Cache at all levels + CacheGlobal, // Cache at L2 only + CacheStreaming // Cache with evict first policy +}; + +enum class StorePolicy { + Writeback, // Cache in L1, write-back on eviction + CacheGlobal, // Bypass L1, write-back on eviction + CacheStreaming // Allocate cache line with evict first policy +}; + +template +__device__ __forceinline__ void load_global(void* dst, const void* src); + +template +__device__ __forceinline__ void load_global(void* dst, const void* src, bool do_access); + +// Shared accesses have no cache policy +template +__device__ __forceinline__ void load_shared(void* dst, const void* src); + +template +__device__ __forceinline__ void load_shared(void* dst, const void* src, bool do_access); + +template +__device__ __forceinline__ void store_global(void* dst, const void* src); + +// Shared accesses have no cache policy +template +__device__ __forceinline__ void store_shared(void* dst, const void* src); + +#ifdef ASYNC_COPY_AVAILABLE +template +__device__ __forceinline__ void memcpy_async(void* shr, const void* gbl); + +template +__device__ __forceinline__ void memcpy_async_nop(void* shr, const void* gbl, bool predicate); + +template +__device__ __forceinline__ void memcpy_async_zero(void* shr, const void* gbl, bool predicate); + +__device__ __forceinline__ void memcpy_async_fence(); + +template +__device__ __forceinline__ void memcpy_async_wait(); + +template +__device__ __forceinline__ void tail_complete_wait(int remaining_stages); +#endif + +// Util for tracking pipeline buffers +// TODO: Evaluate whether this should also be guarded by ASYNC_COPY_AVAILABLE +template +class BufferTracker { +public: + int current_state; + + __device__ __forceinline__ BufferTracker() : current_state(0) {} + + __device__ __forceinline__ int get() + { + int return_val = current_state++; + current_state = (current_state == max ? 0 : current_state); + return return_val; + } +}; + +__device__ __forceinline__ uint32_t lane_id() +{ +#ifdef PTX_AVAILABLE + unsigned int lane_id; + asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane_id)); + return lane_id; +#else + return threadIdx.x & (warpSize - 1); // Portable +#endif +} + +/////////// Load Global /////////// +template <> +__device__ __forceinline__ void load_global<16>(void* dst, const void* src) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.ca.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "l"(src)); +#else + const uint4* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<16>(void* dst, const void* src, bool do_access) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %5, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\tmov.b32 %2, 0;\n" + "\tmov.b32 %3, 0;\n" + "\t@p ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "l"(src), "r"((int)do_access)); +#else + const uint4* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + data[0].z = 0; + data[0].w = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<16, LoadPolicy::CacheGlobal>(void* dst, const void* src) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "l"(src)); +#else + const uint4* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<16, LoadPolicy::CacheGlobal>(void* dst, + const void* src, + bool do_access) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %5, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\tmov.b32 %2, 0;\n" + "\tmov.b32 %3, 0;\n" + "\t@p ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "l"(src), "r"((int)do_access)); +#else + const uint4* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + data[0].z = 0; + data[0].w = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<16, LoadPolicy::CacheStreaming>(void* dst, + const void* src) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cs.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "l"(src)); +#else + const uint4* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<16, LoadPolicy::CacheStreaming>(void* dst, + const void* src, + bool do_access) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %5, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\tmov.b32 %2, 0;\n" + "\tmov.b32 %3, 0;\n" + "\t@p ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "l"(src), "r"((int)do_access)); +#else + const uint4* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + data[0].z = 0; + data[0].w = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<8>(void* dst, const void* src) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.ca.v2.u32 {%0, %1}, [%2];\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "l"(src)); +#else + const uint2* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<8>(void* dst, const void* src, bool do_access) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %3, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\t@p ld.global.v2.u32 {%0, %1}, [%2];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "l"(src), "r"((int)do_access)); +#else + const uint2* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<8, LoadPolicy::CacheGlobal>(void* dst, const void* src) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cg.v2.u32 {%0, %1}, [%2];\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "l"(src)); +#else + const uint2* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<8, LoadPolicy::CacheGlobal>(void* dst, + const void* src, + bool do_access) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %3, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\t@p ld.global.cg.v2.u32 {%0, %1}, [%2];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "l"(src), "r"((int)do_access)); +#else + const uint2* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<8, LoadPolicy::CacheStreaming>(void* dst, + const void* src) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cs.v2.u32 {%0, %1}, [%2];\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "l"(src)); +#else + const uint2* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<8, LoadPolicy::CacheStreaming>(void* dst, + const void* src, + bool do_access) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %3, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\t@p ld.global.cs.v2.u32 {%0, %1}, [%2];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "l"(src), "r"((int)do_access)); +#else + const uint2* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<4>(void* dst, const void* src) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.ca.u32 {%0}, [%1];\n" : "=r"(*data) : "l"(src)); +#else + const int32_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<4>(void* dst, const void* src, bool do_access) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.b32 %0, 0;\n" + "\t@p ld.global.u32 {%0}, [%1];\n" + "}\n" + : "=r"(data[0]) + : "l"(src), "r"((int)do_access)); +#else + const int32_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<4, LoadPolicy::CacheGlobal>(void* dst, const void* src) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cg.u32 {%0}, [%1];\n" : "=r"(*data) : "l"(src)); +#else + const int32_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<4, LoadPolicy::CacheGlobal>(void* dst, + const void* src, + bool do_access) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.b32 %0, 0;\n" + "\t@p ld.global.cg.u32 {%0}, [%1];\n" + "}\n" + : "=r"(data[0]) + : "l"(src), "r"((int)do_access)); +#else + const int32_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<4, LoadPolicy::CacheStreaming>(void* dst, + const void* src) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cs.u32 {%0}, [%1];\n" : "=r"(*data) : "l"(src)); +#else + const int32_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<4, LoadPolicy::CacheStreaming>(void* dst, + const void* src, + bool do_access) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.b32 %0, 0;\n" + "\t@p ld.global.cs.u32 {%0}, [%1];\n" + "}\n" + : "=r"(data[0]) + : "l"(src), "r"((int)do_access)); +#else + const int32_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<2>(void* dst, const void* src) +{ + int16_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.ca.u16 {%0}, [%1];\n" : "=h"(*data) : "l"(src)); +#else + const int16_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<2>(void* dst, const void* src, bool do_access) +{ + int16_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.u16 %0, 0;\n" + "\t@p ld.global.u16 {%0}, [%1];\n" + "}\n" + : "=h"(*data) + : "l"(src), "r"((int)do_access)); +#else + const int16_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<2, LoadPolicy::CacheGlobal>(void* dst, const void* src) +{ + int16_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cg.u16 {%0}, [%1];\n" : "=h"(*data) : "l"(src)); +#else + const int16_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<2, LoadPolicy::CacheGlobal>(void* dst, + const void* src, + bool do_access) +{ + int16_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.u16 %0, 0;\n" + "\t@p ld.global.cg.u16 {%0}, [%1];\n" + "}\n" + : "=h"(*data) + : "l"(src), "r"((int)do_access)); +#else + const int16_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_global<2, LoadPolicy::CacheStreaming>(void* dst, + const void* src) +{ + int16_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile("ld.global.cs.u16 {%0}, [%1];\n" : "=h"(*data) : "l"(src)); +#else + const int16_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_global<2, LoadPolicy::CacheStreaming>(void* dst, + const void* src, + bool do_access) +{ + int16_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.u16 %0, 0;\n" + "\t@p ld.global.cs.u16 {%0}, [%1];\n" + "}\n" + : "=h"(*data) + : "l"(src), "r"((int)do_access)); +#else + const int16_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +/////////// Load Shared /////////// +namespace internal { + +#ifdef PTX_AVAILABLE +__device__ __forceinline__ unsigned convert_to_shared(const void* ptr) +{ +#if __CUDACC_VER_MAJOR__ >= 11 + // In CUDA 11 we have a builtin intrinsic + return __cvta_generic_to_shared(ptr); +#else + unsigned ret_val; + asm volatile( + "{\n" + "\t.reg .u64 p1;\n" + "\tcvta.to.shared.u64 p1, %1\n" + "\tcvt.u32.u64 %0, p1;\n" + "}\n" + : "=r"(ret_val) + : "l"(ptr)); + return ret_val; +#endif +} +#endif + +} // namespace internal + +template <> +__device__ __forceinline__ void load_shared<16>(void* dst, const void* src) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + unsigned src_shr = internal::convert_to_shared(src); + + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "r"(src_shr)); +#else + const uint4* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_shared<16>(void* dst, const void* src, bool do_access) +{ + uint4* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + unsigned src_shr = internal::convert_to_shared(src); + + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %5, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\tmov.b32 %2, 0;\n" + "\tmov.b32 %3, 0;\n" + "\t@p ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) + : "r"(src_shr), "r"((int)do_access)); +#else + const uint4* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + data[0].z = 0; + data[0].w = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_shared<8>(void* dst, const void* src) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + unsigned src_shr = internal::convert_to_shared(src); + + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "r"(src_shr)); +#else + const uint2* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_shared<8>(void* dst, const void* src, bool do_access) +{ + uint2* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + unsigned src_shr = internal::convert_to_shared(src); + + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %3, 0;\n" + "\tmov.b32 %0, 0;\n" + "\tmov.b32 %1, 0;\n" + "\t@p ld.shared.v2.u32 {%0, %1}, [%2];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y) + : "r"(src_shr), "r"((int)do_access)); +#else + const uint2* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0].x = 0; + data[0].y = 0; + } +#endif +} + +template <> +__device__ __forceinline__ void load_shared<4>(void* dst, const void* src) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + unsigned src_shr = internal::convert_to_shared(src); + + asm volatile("ld.shared.u32 {%0}, [%1];\n" : "=r"(*data) : "r"(src_shr)); +#else + const int32_t* src_cast = reinterpret_cast(src); + data[0] = src_cast[0]; +#endif +} + +template <> +__device__ __forceinline__ void load_shared<4>(void* dst, const void* src, bool do_access) +{ + int32_t* data = reinterpret_cast(dst); +#ifdef PTX_AVAILABLE + unsigned src_shr = internal::convert_to_shared(src); + + asm volatile( + "{\n" + "\t.reg .pred p;\n" + "\tsetp.ne.b32 p, %2, 0;\n" + "\tmov.b32 %0, 0;\n" + "\t@p ld.shared.u32 %0, [%1];\n" + "}\n" + : "=r"(data[0]) + : "r"(src_shr), "r"((int)do_access)); +#else + const int32_t* src_cast = reinterpret_cast(src); + if (do_access) { + data[0] = src_cast[0]; + } else { + data[0] = 0; + } +#endif +} + +/////////// Store Global /////////// + +template <> +__device__ __forceinline__ void store_global<16>(void* dst, const void* src) +{ + const uint4* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.wb.v4.u32 [%0], {%1, %2, %3, %4};\n" + : + : "l"(dst), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w) + : "memory"); +#else + uint4* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<16, StorePolicy::CacheGlobal>(void* dst, + const void* src) +{ + const uint4* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.cg.v4.u32 [%0], {%1, %2, %3, %4};\n" + : + : "l"(dst), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w) + : "memory"); +#else + uint4* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<16, StorePolicy::CacheStreaming>(void* dst, + const void* src) +{ + const uint4* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.cs.v4.u32 [%0], {%1, %2, %3, %4};\n" + : + : "l"(dst), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w) + : "memory"); +#else + uint4* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<8>(void* dst, const void* src) +{ + const uint2* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.wb.v2.u32 [%0], {%1, %2};\n" + : + : "l"(dst), "r"(data[0].x), "r"(data[0].y)); +#else + uint2* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<8, StorePolicy::CacheGlobal>(void* dst, + const void* src) +{ + const uint2* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.cg.v2.u32 [%0], {%1, %2};\n" + : + : "l"(dst), "r"(data[0].x), "r"(data[0].y)); +#else + uint2* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<8, StorePolicy::CacheStreaming>(void* dst, + const void* src) +{ + const uint2* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.cs.v2.u32 [%0], {%1, %2};\n" + : + : "l"(dst), "r"(data[0].x), "r"(data[0].y)); +#else + uint2* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<4>(void* dst, const void* src) +{ + const int32_t* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.wb.u32 [%0], %1;\n" : : "l"(dst), "r"(*data)); +#else + int32_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<4, StorePolicy::CacheGlobal>(void* dst, + const void* src) +{ + const int32_t* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.cg.u32 [%0], %1;\n" : : "l"(dst), "r"(*data)); +#else + int32_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<4, StorePolicy::CacheStreaming>(void* dst, + const void* src) +{ + const int32_t* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + asm volatile("st.global.cs.u32 [%0], %1;\n" : : "l"(dst), "r"(*data)); +#else + int32_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_global<2>(void* dst, const void* src) +{ + const int16_t* data = reinterpret_cast(src); + + int16_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +} + +template <> +__device__ __forceinline__ void store_global<2, StorePolicy::CacheGlobal>(void* dst, + const void* src) +{ + const int16_t* data = reinterpret_cast(src); + + int16_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +} + +template <> +__device__ __forceinline__ void store_global<2, StorePolicy::CacheStreaming>(void* dst, + const void* src) +{ + const int16_t* data = reinterpret_cast(src); + + int16_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +} + +/////////// Store Shared /////////// + +template <> +__device__ __forceinline__ void store_shared<16>(void* dst, const void* src) +{ + const uint4* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + unsigned dst_int = internal::convert_to_shared(dst); + + asm volatile("st.shared.v4.u32 [%0], {%1, %2, %3, %4};\n" + : + : "r"(dst_int), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w)); +#else + uint4* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_shared<8>(void* dst, const void* src) +{ + const uint2* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + unsigned dst_int = internal::convert_to_shared(dst); + + asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" + : + : "r"(dst_int), "r"(data[0].x), "r"(data[0].y)); +#else + uint2* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +template <> +__device__ __forceinline__ void store_shared<4>(void* dst, const void* src) +{ + const int32_t* data = reinterpret_cast(src); +#ifdef PTX_AVAILABLE + unsigned dst_int = internal::convert_to_shared(dst); + + asm volatile("st.shared.u32 [%0], %1;\n" : : "r"(dst_int), "r"(*data)); +#else + int32_t* dst_cast = reinterpret_cast(dst); + dst_cast[0] = data[0]; +#endif +} + +/////////// Asynchronous Memory Copy /////////// + +#ifdef ASYNC_COPY_AVAILABLE +template +__device__ __forceinline__ void memcpy_async(void* shr, const void* gbl) +{ + static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); + unsigned shr_int = internal::convert_to_shared(shr); + + asm volatile("cp.async.ca.shared.global [%0], [%1], %2;\n" + : + : "r"(shr_int), "l"(gbl), "n"(AccessSize)); +} + +template +__device__ __forceinline__ void memcpy_async_nop(void* shr, const void* gbl, bool predicate) +{ + static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); + unsigned shr_int = internal::convert_to_shared(shr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" + : + : "r"((int)predicate), "r"(shr_int), "l"(gbl), "n"(AccessSize)); +} + +template +__device__ __forceinline__ void memcpy_async_zero(void* shr, const void* gbl, bool predicate) +{ + static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); + unsigned shr_int = internal::convert_to_shared(shr); + int bytes_to_copy = (predicate ? AccessSize : 0); + + asm volatile("cp.async.ca.shared.global [%0], [%1], %2, %3;\n" + : + : "r"(shr_int), "l"(gbl), "n"(AccessSize), "r"(bytes_to_copy)); +} + +template +__device__ __forceinline__ void memcpy_async_zero_nop(void* shr, + const void* gbl, + bool zero_predicate, + bool nop_predicate) +{ + static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); + unsigned shr_int = internal::convert_to_shared(shr); + int bytes_to_copy = (zero_predicate ? AccessSize : 0); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3, %4;\n" + "}\n" + : + : "r"((int)nop_predicate), "r"(shr_int), "l"(gbl), "n"(AccessSize), "r"(bytes_to_copy)); +} + +// Cache global variants. Separate interface to require deliberate use of them. +__device__ __forceinline__ void memcpy_async_cg(void* shr, const void* gbl) +{ + unsigned shr_int = internal::convert_to_shared(shr); + + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(shr_int), "l"(gbl)); +} + +__device__ __forceinline__ void memcpy_async_nop_cg(void* shr, const void* gbl, bool predicate) +{ + unsigned shr_int = internal::convert_to_shared(shr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], 16;\n" + "}\n" + : + : "r"((int)predicate), "r"(shr_int), "l"(gbl)); +} + +__device__ __forceinline__ void memcpy_async_zero_cg(void* shr, const void* gbl, bool predicate) +{ + unsigned shr_int = internal::convert_to_shared(shr); + int bytes_to_copy = (predicate ? 16 : 0); + + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" + : + : "r"(shr_int), "l"(gbl), "r"(bytes_to_copy)); +} + +__device__ __forceinline__ void memcpy_async_zero_nop_cg(void* shr, + const void* gbl, + bool zero_predicate, + bool nop_predicate) +{ + unsigned shr_int = internal::convert_to_shared(shr); + int bytes_to_copy = (zero_predicate ? 16 : 0); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], 16, %3;\n" + "}\n" + : + : "r"((int)nop_predicate), "r"(shr_int), "l"(gbl), "r"(bytes_to_copy)); +} + +__device__ __forceinline__ void memcpy_async_fence() { asm volatile("cp.async.commit_group;\n"); } + +template +__device__ __forceinline__ void memcpy_async_wait() +{ + static_assert(stages <= 8); + + asm volatile("cp.async.wait_group %0;\n" : : "n"(stages)); +} + +// TODO: The tail complete should be a known compile time artifact, should try and induce this +// without all of the branches from the call-site. This is a hacky solution. +template <> +__device__ __forceinline__ void tail_complete_wait<1>(int remaining_stages) +{ + if (remaining_stages == 0) memcpy_async_wait<0>(); +} + +template <> +__device__ __forceinline__ void tail_complete_wait<2>(int remaining_stages) +{ + if (remaining_stages == 1) + memcpy_async_wait<1>(); + else if (remaining_stages == 0) + memcpy_async_wait<0>(); +} + +template <> +__device__ __forceinline__ void tail_complete_wait<3>(int remaining_stages) +{ + if (remaining_stages == 2) + memcpy_async_wait<2>(); + else if (remaining_stages == 1) + memcpy_async_wait<1>(); + else if (remaining_stages == 0) + memcpy_async_wait<0>(); +} + +template <> +__device__ __forceinline__ void tail_complete_wait<4>(int remaining_stages) +{ + if (remaining_stages == 3) + memcpy_async_wait<3>(); + else if (remaining_stages == 2) + memcpy_async_wait<2>(); + else if (remaining_stages == 1) + memcpy_async_wait<1>(); + else if (remaining_stages == 0) + memcpy_async_wait<0>(); +} + +template <> +__device__ __forceinline__ void tail_complete_wait<5>(int remaining_stages) +{ + if (remaining_stages == 4) + memcpy_async_wait<4>(); + else if (remaining_stages == 3) + memcpy_async_wait<3>(); + else if (remaining_stages == 2) + memcpy_async_wait<2>(); + else if (remaining_stages == 1) + memcpy_async_wait<1>(); + else if (remaining_stages == 0) + memcpy_async_wait<0>(); +} + +template <> +__device__ __forceinline__ void tail_complete_wait<6>(int remaining_stages) +{ + if (remaining_stages == 5) + memcpy_async_wait<5>(); + else if (remaining_stages == 4) + memcpy_async_wait<4>(); + else if (remaining_stages == 3) + memcpy_async_wait<3>(); + else if (remaining_stages == 2) + memcpy_async_wait<2>(); + else if (remaining_stages == 1) + memcpy_async_wait<1>(); + else if (remaining_stages == 0) + memcpy_async_wait<0>(); +} +#endif + +} // namespace mem_access diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/normalize_layer.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/normalize_layer.h new file mode 100644 index 0000000000000000000000000000000000000000..b9c719087a6ba0af693589f3c36da452b82e21d7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/normalize_layer.h @@ -0,0 +1,207 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include "custom_cuda_layers.h" + +using namespace std; + +template +class Normalize_Layer { +public: + struct Config { + uint32_t batchSize; + uint32_t seqLength; + uint32_t hiddenDim; + float epsilon; + bool training; + bool useMean; + Config(uint32_t batch, + uint32_t seq, + uint32_t h, + float epsilon = 1e-12, + bool training = true, + bool useMean = true) + : batchSize(batch), + seqLength(seq), + hiddenDim(h), + epsilon(epsilon), + training(training), + useMean(useMean) + { + } + }; + + Normalize_Layer(Config config) + : config_(config), vars(nullptr), means(nullptr), vals_hat(nullptr) + { + } + + ~Normalize_Layer() {} + + void ForwardCheckpoint(int bsz, // batch * seq + T* vals, + const T* residual, + const T* gamma, + const T* betta, + cudaStream_t& stream, + bool preLayerNorm = false) + { + launch_bias_residual_layer_norm(vals, + residual, + gamma, + betta, + config_.epsilon, + bsz, + config_.hiddenDim, + stream, + preLayerNorm, + config_.training, + vars, + means); + } + + void Forward(int bsz, + T* vals, + const T* residual, + const T* gamma, + const T* betta, + cudaStream_t& stream, + bool preLayerNorm = false) + { + launch_bias_residual_layer_norm(vals, + residual, + gamma, + betta, + config_.epsilon, + bsz, + config_.hiddenDim, + stream, + preLayerNorm, + config_.training, + vars); + } + + void Backward(int bsz, + const T* out_grad, + const T* gamma, + T* gamma_grad, + T* betta_grad, + cudaStream_t stream[2], + T* inp_grad_out, + const T* norm_in = nullptr) + { + launch_layerNorm_backward(out_grad, + norm_in, + vars, + means, + gamma, + gamma_grad, + betta_grad, + inp_grad_out, + bsz, + config_.hiddenDim, + stream); + } + + void Backward(int bsz, + const T* out_grad, + const T* gamma, + const T* betta, + T* gamma_grad, + T* betta_grad, + cudaStream_t stream[2], + T* inp_grad_out, + const T* norm_out) + { + launch_layerNorm_backward(out_grad, + norm_out, + vars, + gamma, + gamma_grad, + betta_grad, + inp_grad_out, + bsz, + config_.hiddenDim, + stream, + !config_.useMean, + betta); + } + + void BackwardFusedAdd(int bsz, + const T* out_grad1, + const T* out_grad2, + const T* gamma, + T* gamma_grad, + T* betta_grad, + cudaStream_t stream[2], + T* inp_grad_out, + const T* norm_in = nullptr) + { + launch_layerNorm_backward_fused_add(out_grad1, + out_grad2, + norm_in, + vars, + means, + gamma, + gamma_grad, + betta_grad, + inp_grad_out, + bsz, + config_.hiddenDim, + stream); + } + + void BackwardFusedAdd(int bsz, + const T* out_grad1, + const T* out_grad2, + const T* gamma, + const T* betta, + T* gamma_grad, + T* betta_grad, + cudaStream_t stream[2], + T* inp_grad_out, + const T* norm_out) + { + launch_layerNorm_backward_fused_add(out_grad1, + out_grad2, + norm_out, + vars, + gamma, + gamma_grad, + betta_grad, + inp_grad_out, + bsz, + config_.hiddenDim, + stream, + !config_.useMean, + betta); + } + + inline bool UseMean() const { return config_.useMean; } + + inline void SetVar(T* variance) + { + if (!variance) { throw std::runtime_error("Normalize variance is null."); } + vars = variance; + } + + inline void SetMean(T* mean) + { + if (!mean) { throw std::runtime_error("Normalize mean is null."); } + means = mean; + } + +private: + Config config_; + T* vars; + T* means; + T* vals_hat; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization.h new file mode 100644 index 0000000000000000000000000000000000000000..5bdc96061a311d78b94ec9051e7f13919454934b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization.h @@ -0,0 +1,138 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include "ds_kernel_utils.h" + +namespace quantize { + +enum class Type { Symmetric, Asymmetric }; + +struct PackedInt4 { + int8_t high : 4; + int8_t low : 4; +}; + +DS_HD_INLINE bool requires_offset(Type qType) { return qType == Type::Asymmetric; } + +} // namespace quantize + +void launch_quant(int8_t* output_data, + float* params, + const __half* input_data, + const int groups, + const int elems_per_group, + const int num_bits, + const quantize::Type quant_type, + cudaStream_t stream); + +template +void launch_dequantize_kernel(T* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream); + +void launch_swizzled_quant(int8_t* q_data, + float* q_scales, + const __half* input_data, + int num_bits, + quantize::Type q_type, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream); + +void launch_loco_swizzled_quant(int8_t* quantized_data, + float* quantized_scales, + const __half* uncompressed_data, + __half* error_feedback, + const float err_beta, + int num_bits, + quantize::Type quant_type, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream); + +void launch_loco_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int num_gpus, + int num_bits, + quantize::Type quant_type, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + __half2* error_feedback, + const float err_beta, + cudaStream_t stream); + +void launch_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int num_gpus, + int num_bits, + quantize::Type quant_type, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + cudaStream_t stream); + +template +void launch_fake_quantize_kernel(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template +void launch_sr_fake_quantize_kernel(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template +void launch_fake_quantize_kernel_asym(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template +void launch_sr_fake_quantize_kernel_asym(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +void launch_dequantize_int4_to_half_experimental(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size, + cudaStream_t stream); + +void launch_dequantize_int8_to_half_experimental(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..94958fb455c6f632713aea4f27b7bc1c901bfa45 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization_utils.h @@ -0,0 +1,469 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" +#include "quantization.h" +#include "reduction_utils.h" + +#pragma once + +using rop = reduce::ROpType; + +namespace quantize { +constexpr int granularity = 16; +constexpr int h_per_load = granularity / sizeof(__half); +constexpr int h2_per_load = granularity / sizeof(__half2); +constexpr int max_threads = 1024; + +/* +Class to hold the quantization parameters for a given tensor. +Holds the implementation of the quantization operation. +*/ + +template +class Params { +public: + /* + Quantization implementation, supports + 1) 4 Bit + 2) 8 Bit + 3) Symmetric + 4) Asymmetric + Function Arguments : + val : The __half value to quantize. + */ + DS_D_INLINE int8_t quantize(__half val); + + template + DS_D_INLINE T dequantize(int8_t val); + + DS_D_INLINE void store(float* params, int group_index); + + // Initialize from memory + DS_D_INLINE Params(const float* params, int group_index); +}; + +template +class Params { +public: + float scale; + + DS_D_INLINE Params(float max) + { + if (max == 0) { + scale = 1.0; + } else { + scale = (1 << numBits) / (2 * max); + } + } + + DS_D_INLINE int8_t quantize(__half val) + { + constexpr int32_t q_min = -(1 << (numBits - 1)); + constexpr int32_t q_max = (1 << (numBits - 1)) - 1; + + float val_f = conversion::to(val) * scale; + int32_t data_i32 = conversion::to(val_f); + data_i32 = min(max(data_i32, q_min), q_max); + return (int8_t)data_i32; + } + + template + DS_D_INLINE T dequantize(int8_t val) + { + const float val_deq_f = conversion::to(val) * scale; + return conversion::to(val_deq_f); + } + + DS_D_INLINE void store(float* params, int group_index) + { + const float store_scale = 1 / scale; + mem_access::store_global(params + group_index, &store_scale); + } + + DS_D_INLINE Params(const float* params, int group_index) + { + mem_access::load_global(&scale, params + group_index); + } +}; + +template +class Params { +public: + float scale; + float offset; + + DS_D_INLINE Params(float max, float min) + { + if (max == min) { + scale = 1.0; + } else { + scale = ((1 << numBits)) / (max - min); + } + offset = (max + min) / 2; + } + + DS_D_INLINE int8_t quantize(__half val) + { + constexpr int32_t q_min = -(1 << (numBits - 1)); + constexpr int32_t q_max = (1 << (numBits - 1)) - 1; + + float val_f = (conversion::to(val) - offset) * scale; + int32_t data_i32 = conversion::to(val_f); + data_i32 = min(max(data_i32, q_min), q_max); + return (int8_t)data_i32; + } + + template + DS_D_INLINE T dequantize(int8_t val) + { + const float val_deq_f = ((conversion::to(val)) * scale) + offset; + return conversion::to(val_deq_f); + } + + DS_D_INLINE void store(float* params, int group_index) + { + // Codegen should turn this into stg.64 + const float store_scale = 1 / scale; + mem_access::store_global(params + 2 * group_index, &store_scale); + mem_access::store_global(params + 2 * group_index + 1, &offset); + } + + DS_D_INLINE Params(const float* params, int group_index) + { + // Codegen should turn this into ldg.64 + mem_access::load_global(&scale, params + 2 * group_index); + mem_access::load_global(&offset, params + 2 * group_index + 1); + } +}; + +/* +Group stats tracks the necessary statistics about the quantized group +to abstract the particulars for the main loop. +*/ +template +class GroupStats { +public: + DS_D_INLINE void update(__half2 val); + + DS_D_INLINE void reduce(cg::thread_block& tb, cg::thread_block_tile& warp); +}; + +template <> +class GroupStats { +public: + // Symmetric quantization only tracks the maximum absolute value + __half2 cur_max; + float max; + + /* + Technically, this would give bad results if there + are 0 values to process since the reduction would + give -inf instead of 0. We do not consider this + to be a reasonable edge case. + */ + DS_D_INLINE GroupStats() { cur_max = reduce::init(); } + + /* + Updated the running absmax used to calculate params. + Function Arguments : + val : The __half2 value to update the running min and max with. + */ + DS_D_INLINE void update(__half2 val) + { + cur_max = reduce::element(cur_max, __habs2(val)); + } + + /* + Function to return calculated quantization params. + Template Arguments : + numBits - Number of bits in quantized element. int : 8 or 4 + Function Arguments : + tb - Threadblock object. cg::thread_block + warp - Warp object. cg::thread_block_tile + */ + template + DS_D_INLINE Params get_params( + cg::thread_block& tb, + cg::thread_block_tile& warp) + { + const float2 partial_max = conversion::to(cur_max); + float max = reduce::element(partial_max.x, partial_max.y); + + reduce::partitioned_block(tb, warp, max); + Params params(max); + + return params; + } +}; + +template <> +class GroupStats { +public: + __half2 cur_max; + __half2 cur_min; + + /* + Initialize cur_max to -inf, cur_min to inf since + we are doing a true range analysis. + */ + DS_D_INLINE GroupStats() + { + cur_max = reduce::init(); + cur_min = reduce::init(); + } + + /* + Updated the running min and max used to calculate params. + Function Arguments : + val : The __half2 value to update the running min and max with. + */ + DS_D_INLINE void update(__half2 val) + { + cur_max = reduce::element(cur_max, val); + cur_min = reduce::element(cur_min, val); + } + + /* + Function to return calculated quantization params. + Template Arguments : + numBits - Number of bits in quantized element. int : 8 or 4 + Function Arguments : + tb - Threadblock object. cg::thread_block + warp - Warp object. cg::thread_block_tile + */ + template + DS_D_INLINE Params get_params( + cg::thread_block& tb, + cg::thread_block_tile& warp) + { + const float2 partial_max = conversion::to(cur_max); + float max = reduce::element(partial_max.x, partial_max.y); + + const float2 partial_min = conversion::to(cur_min); + float min = reduce::element(partial_min.x, partial_min.y); + + reduce::partitioned_block(tb, warp, max, min); + + Params params(max, min); + + return params; + } +}; + +/* +Device function that quantizes 16 bytes of __half type input data. +Template Arguments : + numBits - Number of bits in quantized element. int : 8 or 4 + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric +Function Arguments : + local_output - Pointer to local memory to store quantized data. int8_t* + data - Pointer to input data. __half* + Params - Parameters for quantization. Params +*/ +template +DS_D_INLINE void _chunk(int8_t* local_output, const __half* data, Params q_params); + +/* +Device function that quantizes 16 bytes of __half2 type input data. +Template Arguments : + numBits - Number of bits in quantized element. int : 8 or 4 + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric +Function Arguments : + local_output - Pointer to local memory to store quantized data. int8_t* + data - Pointer to input data. __half2* + Params - Parameters for quantization. Params +*/ +template +DS_D_INLINE void _chunk(int8_t* local_output, const __half2* data, Params q_params); + +/* +Helper function to do serial reduction on register-file arrays. +Template Arguments : + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric + numChunks - Number of bits in quantized element. int : 8 or 4 +Function Arguments : + local_buffer - Pointer memory with input half2 data to be quantized. +*/ +template +DS_D_INLINE GroupStats _local_serial_reduce(__half2* local_buffer); + +/* +The main loop of the kernel that quantizes array in local memory of __half2 type input data, when +Quantization parameters are pre-computed. +Template Arguments : + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric + numBits - Number of bits in quantized element. int : 8 or 4 + numChunks - Number of chunks(16 bytes of Input data). int : 8 or 4 +Function Arguments : + local_buffer - Pointer memory with input half2 data to be quantized. + scales - Pointer to output scales. + offsets - Pointer to output offsets. + output_data - Pointer to output data. + elems_per_group - Number of elements to quantize in a group. + q_params - Quantization parameters. +*/ +template +DS_D_INLINE void local_array(cg::thread_block& tb, + cg::thread_block_tile& warp, + __half2* local_buffer, + float* __restrict__ scales, + float* __restrict__ offsets, + int8_t* __restrict__ output_data, + const int& elems_per_group, + const int& groups, + Params q_params); + +/* +The main loop of the kernel that quantizes array in local memory of __half2 type input data. +This function computes quantization parameters for each group. +Template Arguments : + qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric + numBits - Number of bits in quantized element. int : 8 or 4 + numChunks - Number of chunks(16 bytes of Input data). int : 8 or 4 +Function Arguments : + local_buffer - Pointer memory with input half2 data to be quantized. + scales - Pointer to output scales. + offsets - Pointer to output offsets. + output_data - Pointer to output data. + elems_per_group - Number of elements to quantize in a group. +*/ +template +__device__ void local_array(__half2* local_buffer, + float* __restrict__ scales, + float* __restrict__ offsets, + int8_t* __restrict__ output_data, + const int& elems_per_group, + const int& groups); + +template +DS_D_INLINE void _chunk(int8_t* local_output, const __half* data, Params q_params) +{ + constexpr int32_t elems = 16 / sizeof(__half); + constexpr int32_t num_elems_packed = 8 / numBits; + +#pragma unroll + for (int i = 0, oi = 0; i < elems; i += num_elems_packed, oi++) { + if (num_elems_packed == 1) { + // TODO(cmikeh2): refactor to use conversion utils + local_output[i] = q_params.quantize(data[i]); + } else if (num_elems_packed == 2) { + int8_t data_i8_1 = q_params.quantize(data[i]); + int8_t data_i8_2 = q_params.quantize(data[i + 1]); + auto data_i8 = PackedInt4{data_i8_2, data_i8_1}; + local_output[oi] = *((int8_t*)(&data_i8)); + } + } +} + +template +DS_D_INLINE void _chunk(int8_t* local_output, const __half2* data, Params q_params) +{ + const __half* data_cast = reinterpret_cast(data); + _chunk(local_output, data_cast, q_params); +} + +template +DS_D_INLINE GroupStats _local_serial_reduce(__half2* local_buffer) +{ + GroupStats stats; +#pragma unroll + for (int i = 0; i < numChunks * h2_per_load; i++) { stats.update(local_buffer[i]); } + + return stats; +} + +template +DS_D_INLINE void local_array(cg::thread_block& tb, + cg::thread_block_tile& warp, + __half2* local_buffer, + float* __restrict__ global_params, + int8_t* __restrict__ output_data, + const int& elems_per_group, + const int& groups, + Params q_params) +{ + constexpr int num_ele_int8 = 8 / numBits; + constexpr int num_int8_out = quantize::h_per_load / num_ele_int8; + + // Indexing offsets + const int block_num = + (tb.group_index().x * max_threads / threads_per_group) + tb.thread_index().y; + const int block_offset = block_num * elems_per_group; + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset = (block_offset + elem_offset) / num_ele_int8; + const int stride = tb.size() * quantize::h_per_load / num_ele_int8; + + int8_t local_output[num_int8_out]; + + if (tb.thread_index().x == 0 && block_num < groups) { + q_params.store( + global_params, + (tb.group_index().x * max_threads / threads_per_group) + tb.thread_index().y); + } +#pragma unroll + for (int i = 0; i < numChunks; i++) { + if (elem_offset + i * stride * num_ele_int8 < elems_per_group && block_num < groups) { + quantize::_chunk( + local_output, local_buffer + i * quantize::h2_per_load, q_params); + mem_access::store_global(output_data + (base_offset + i * stride), + local_output); + } + } +} + +template +DS_D_INLINE void local_array(cg::thread_block& tb, + cg::thread_block_tile& warp, + __half* local_buffer, + float* __restrict__ global_params, + int8_t* __restrict__ output_data, + const int& elems_per_group, + const int& groups, + Params q_params) +{ + __half2* local_buffer_h2 = reinterpret_cast<__half2*>(local_buffer); + + quantize::local_array( + tb, warp, local_buffer, global_params, output_data, elems_per_group, groups, q_params); +} + +template +__device__ void local_array(__half2* local_buffer, + float* __restrict__ global_params, + int8_t* __restrict__ output_data, + const int& elems_per_group, + const int& groups) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + auto group_stats = _local_serial_reduce(local_buffer); + auto params = group_stats.template get_params(tb, warp); + + quantize::local_array( + tb, warp, local_buffer, global_params, output_data, elems_per_group, groups, params); +} + +template +__device__ void local_array(__half* local_buffer, + float* __restrict__ global_params, + int8_t* __restrict__ output_data, + const int& elems_per_group, + const int& groups) +{ + __half2* local_buffer_h2 = reinterpret_cast<__half2*>(local_buffer); + quantize::local_array( + local_buffer_h2, global_params, output_data, elems_per_group, groups); +} + +} // namespace quantize diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantizer.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantizer.h new file mode 100644 index 0000000000000000000000000000000000000000..f4f63160d79b23aaa94ae183a176d6c05024496a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantizer.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#ifdef __HIP_PLATFORM_AMD__ +#include +#else +#include +#endif + +#include +#include +#include +#include +#include +#include diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/reduction_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/reduction_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..eb9afb66a894a63889c3b07c33a634ff60f54f36 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/reduction_utils.h @@ -0,0 +1,826 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +namespace reduce { + +enum class ROpType { + // Addition + Add, + + // Maximum reduction + Max, + + // Minimum reduction + Min, +}; + +constexpr int max_threads = 1024; +constexpr int max_warps = max_threads / hw_warp_size; + +/* +High level API. The API takes in a set of operations and variables +and performs that reduction operation on that variable. The reductions +of each of the arguments are completely independent of each other ( +i.e., the val1-op1 combination has no impact on val2-op2). + +Example usage: +``` cpp +float max_val; +float min_val; +reduce::block(tb, warp, max_val, min_val); +``` + +TODO(cmikeh2): In theory, we might be able to do this sequentially with +device functions and rely on the assembler correctly behaving. My initial +instinct is this won't work, but if it does it would reduce implementation +cost significantly. + +TODO(cmikeh2): We need to support sub-block reductions. The warp intrinsic +currently supports this (more incidentally than anything else). It is not +uncommon in something like softmax or a fused attention kernel to map multiple +reductions to a thread block, but each reduction itself is only scoped +to part of the threads (i.e block size = 512, 128 threads per reduction). +*/ +template +DS_D_INLINE void block(cg::thread_block& tb, cg::thread_block_tile& warp, float& val); + +template +DS_D_INLINE void block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2); + +template +DS_D_INLINE void block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3); + +template +DS_D_INLINE void block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3, + float& val4); + +/* +The partitioned block is a special case of the above where in the warps of a threadblock are +partitioned into separate independent reductions. For example, I might have an 8 warp thread block +in which each pair of warps is processing an independent piece of data. I would then reduce that +data with the something like the following: +``` cpp +float max_val; +reduce::partitioned_block(tb, warp, max_val); +``` +After which, each pair of warps would have coherent data with each other. Note, this API will not +provide correct results if the number of warps per partition is not a power of 2. +*/ +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val); + +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2); + +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3); + +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3, + float& val4); + +/* +Single element reduction primitives. Used inside serial collection +loops. + +Example usage: +using rop = reduce::OpType; +float min = init(); +for (int i = 0; i < 4; i++) { + min = reduce::element(min, data[i]); +} +*/ + +template +DS_D_INLINE T element(const T lhs, const T rhs); + +template +DS_D_INLINE T init(); + +/********************** Internal reduction APIs **********************/ + +/* +Single element "reductions". TODO(cmikeh2): this sort of "op" concept +should be refactored into its own implementation at some point. This interface +may be easily expanded for new types/operations, but the typical reductions +we need are covered with min/max/add on float. + +NOTE: there is no mean reduction because that relies on knowledge of how +many values were already reduced into each scalar. Implementing this on top +of reduce should be straightforward (can just wrap the sum reduction) and +would be a good extension of the header. +*/ + +DS_D_INLINE int _warp_rank() +{ + const int thread_rank = + threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y; + return thread_rank / hw_warp_size; +} + +/* Float element reduce implementations */ +template <> +DS_D_INLINE float element(const float lhs, const float rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE double element(const double lhs, const double rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE float element(const float lhs, const float rhs) +{ + return fmaxf(lhs, rhs); +} + +template <> +DS_D_INLINE float element(const float lhs, const float rhs) +{ + return fminf(lhs, rhs); +} + +/* __half element reduce implementation */ +template <> +DS_D_INLINE __half element(const __half lhs, const __half rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE __half element(const __half lhs, const __half rhs) +{ +#if __CUDA_ARCH__ >= 800 + // Intrinsic limited to Ampere + newer + return __hmax(lhs, rhs); +#else + return (lhs > rhs) ? lhs : rhs; +#endif +} + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE __nv_bfloat16 element(const __nv_bfloat16 lhs, const __nv_bfloat16 rhs) +{ +#if __CUDA_ARCH__ >= 800 + // Intrinsic limited to Ampere + newer + return __hmax(lhs, rhs); +#else + return (lhs > rhs) ? lhs : rhs; +#endif +} +#endif + +template <> +DS_D_INLINE __half element(const __half lhs, const __half rhs) +{ +#if __CUDA_ARCH__ >= 800 + // Intrinsic limited to Ampere + newer + return __hmin(lhs, rhs); +#else + return (lhs < rhs) ? lhs : rhs; +#endif +} + +/* __half2 element reduce implementation */ +template <> +DS_D_INLINE __half2 element(const __half2 lhs, const __half2 rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE __half2 element(const __half2 lhs, const __half2 rhs) +{ +#if __CUDA_ARCH__ >= 800 + return __hmax2(lhs, rhs); +#else + __half2 ret_val; + ret_val.x = (lhs.x > rhs.x) ? lhs.x : rhs.x; + ret_val.y = (lhs.y > rhs.y) ? lhs.y : rhs.y; + return ret_val; +#endif +} + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE __nv_bfloat162 element(const __nv_bfloat162 lhs, const __nv_bfloat162 rhs) +{ +#if __CUDA_ARCH__ >= 800 + return __hmax2(lhs, rhs); +#else + __nv_bfloat162 ret_val; + ret_val.x = (lhs.x > rhs.x) ? lhs.x : rhs.x; + ret_val.y = (lhs.y > rhs.y) ? lhs.y : rhs.y; + return ret_val; +#endif +} +#endif + +template <> +DS_D_INLINE __half2 element(const __half2 lhs, const __half2 rhs) +{ +#if __CUDA_ARCH__ >= 800 + return __hmin2(lhs, rhs); +#else + __half2 ret_val; + ret_val.x = (lhs.x < rhs.x) ? lhs.x : rhs.x; + ret_val.y = (lhs.y < rhs.y) ? lhs.y : rhs.y; + return ret_val; +#endif +} + +template <> +DS_D_INLINE int32_t element(const int32_t lhs, const int32_t rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE int32_t element(const int32_t lhs, const int32_t rhs) +{ + return (lhs > rhs) ? lhs : rhs; +} + +template <> +DS_D_INLINE int32_t element(const int32_t lhs, const int32_t rhs) +{ + return (lhs < rhs) ? lhs : rhs; +} + +template <> +DS_D_INLINE uint32_t element(const uint32_t lhs, const uint32_t rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE uint32_t element(const uint32_t lhs, const uint32_t rhs) +{ + return (lhs > rhs) ? lhs : rhs; +} + +template <> +DS_D_INLINE uint32_t element(const uint32_t lhs, const uint32_t rhs) +{ + return (lhs < rhs) ? lhs : rhs; +} + +template <> +DS_D_INLINE int64_t element(const int64_t lhs, const int64_t rhs) +{ + return lhs + rhs; +} + +template <> +DS_D_INLINE int64_t element(const int64_t lhs, const int64_t rhs) +{ + return (lhs > rhs) ? lhs : rhs; +} + +template <> +DS_D_INLINE int64_t element(const int64_t lhs, const int64_t rhs) +{ + return (lhs < rhs) ? lhs : rhs; +} + +/* +Reduction initialization primitives +*/ +template <> +DS_D_INLINE float init() +{ + return 0.0f; +} +template <> +DS_D_INLINE double init() +{ + return (double)0.0f; +} + +template <> +DS_D_INLINE float init() +{ + // Positive infinity + return INFINITY; +} + +template <> +DS_D_INLINE float init() +{ + // Negative infinity + return -INFINITY; +} + +template <> +DS_D_INLINE __half init() +{ + constexpr __half_raw zero = {0x0000}; + return __half(zero); +} + +template <> +DS_D_INLINE __half init() +{ + constexpr __half_raw inf = {0x7C00}; + return __half(inf); +} + +template <> +DS_D_INLINE __half init() +{ + constexpr __half_raw neg_inf = {0xFC00}; + return __half(neg_inf); +} + +#ifdef BF16_AVAILABLE +template <> +DS_D_INLINE __nv_bfloat16 init() +{ + constexpr __nv_bfloat16_raw neg_inf = {0xFF80}; + return __nv_bfloat16(neg_inf); +} +#endif + +template <> +DS_D_INLINE __half2 init() +{ +#ifdef __HIP_PLATFORM_AMD__ + return __half2{_Float16_2{0x0000, 0x0000}}; +#else + constexpr __half2_raw zero = {0x0000, 0x0000}; + return __half2(zero); +#endif +} + +template <> +DS_D_INLINE __half2 init() +{ +#ifdef __HIP_PLATFORM_AMD__ + return __half2{_Float16_2{0x7C00, 0x7C00}}; +#else + constexpr __half2_raw inf = {0x7C00, 0x7C00}; + return __half2(inf); +#endif +} + +template <> +DS_D_INLINE __half2 init() +{ +#ifdef __HIP_PLATFORM_AMD__ + return __half2{_Float16_2{0xFC00, 0xFC00}}; +#else + constexpr __half2_raw neg_inf = {0xFC00, 0xFC00}; + return __half2(neg_inf); +#endif +} + +template <> +DS_D_INLINE int32_t init() +{ + return 0; +} + +template <> +DS_D_INLINE int32_t init() +{ + return 0x7FFFFFFF; +} + +template <> +DS_D_INLINE int32_t init() +{ + return 0x80000000; +} + +template <> +DS_D_INLINE uint32_t init() +{ + return 0; +} + +template <> +DS_D_INLINE uint32_t init() +{ + return 0xFFFFFFFF; +} + +template <> +DS_D_INLINE uint32_t init() +{ + return 0; +} + +template <> +DS_D_INLINE int64_t init() +{ + return 0; +} + +template <> +DS_D_INLINE int64_t init() +{ + return 0x7FFFFFFFFFFFFFFF; +} + +template <> +DS_D_INLINE int64_t init() +{ + return 0x8000000000000000; +} + +template <> +DS_D_INLINE uint64_t init() +{ + return 0; +} + +template <> +DS_D_INLINE uint64_t init() +{ + return 0xFFFFFFFFFFFFFFFF; +} + +template <> +DS_D_INLINE uint64_t init() +{ + return 0; +} + +template +DS_D_INLINE void init(T* data) +{ + data[0] = init(); +} + +template +DS_D_INLINE void init(T* data) +{ + data[0] = init(); + data[1] = init(); +} + +template +DS_D_INLINE void init(T* data) +{ + data[0] = init(); + data[1] = init(); + data[2] = init(); +} + +template +DS_D_INLINE void init(T* data) +{ + data[0] = init(); + data[1] = init(); + data[2] = init(); + data[3] = init(); +} + +/* +Warp reduction primitives + +`reduction_width` is an unsafe template parameter, that is that +when using `reduction_width` < hw_warp_size the warp is partitioned +into `hw_warp_size` / `reduction_width` groups of partial sums. + +If someone can figure out how to use variadic templates in a reasonable way +here (fold is C++17 only and I don't think helps and recursion feels like +huge overkill that harms readability) that would be wonderful. +*/ + +template +DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) +{ +#pragma unroll + for (int i = 1; i < reduce_width; i *= 2) { + data[0] = element(data[0], warp.shfl_xor(data[0], i)); + } +} + +template +DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) +{ +#pragma unroll + for (int i = 1; i < reduce_width; i *= 2) { + data[0] = element(data[0], warp.shfl_xor(data[0], i)); + data[1] = element(data[1], warp.shfl_xor(data[1], i)); + } +} + +template +DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) +{ +#pragma unroll + for (int i = 1; i < reduce_width; i *= 2) { + data[0] = element(data[0], warp.shfl_xor(data[0], i)); + data[1] = element(data[1], warp.shfl_xor(data[1], i)); + data[2] = element(data[2], warp.shfl_xor(data[2], i)); + } +} + +template +DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) +{ +#pragma unroll + for (int i = 1; i < reduce_width; i *= 2) { + data[0] = element(data[0], warp.shfl_xor(data[0], i)); + data[1] = element(data[1], warp.shfl_xor(data[1], i)); + data[2] = element(data[2], warp.shfl_xor(data[2], i)); + data[3] = element(data[3], warp.shfl_xor(data[3], i)); + } +} + +/* +Implementation for primary block reduction that serves both `block` and +`partitioned_block`. + +Total warps refers to the reduction width of the reduction, not +the number of warps in the block (which may exceed that +if the block is partitioned or if we do a conservative bound at +compile time). +*/ +template +DS_D_INLINE void _block(cg::thread_block& tb, + cg::thread_block_tile& warp_arg, + T* data) +{ + constexpr int elems = sizeof...(Ops); + constexpr int bytes = sizeof(T); + // Unused when `partition_size == 1` or total_warps == 1 + __shared__ T reduce_buffer[max_warps * elems]; + +#ifdef __HIP_PLATFORM_AMD__ + const int total_threads = blockDim.x * blockDim.y * blockDim.z; + const int running_warps = total_threads / hw_warp_size; +#else + const int running_warps = warp_arg.meta_group_size(); +#endif + + // Always perform warp-scope reduction + _warp(warp_arg, data); + + // If max_warps == 1 let's skip the runtime check + if (total_warps != 1) { + if (warp_arg.thread_rank() == 0) { +#pragma unroll + for (int i = 0; i < elems; i++) { + mem_access::store_shared(reduce_buffer + elems * _warp_rank() + i, data + i); + } + } + + // Synchronization inside block-uniform conditional is safe + tb.sync(); + + if (_warp_rank() == 0) { + if (warp_arg.thread_rank() < running_warps) { +#pragma unroll + for (int i = 0; i < elems; i++) { + mem_access::load_shared( + data + i, reduce_buffer + elems * warp_arg.thread_rank() + i); + } + } else { + init(data); + } + + _warp(warp_arg, data); + +#pragma unroll + for (int i = 0; i < elems; i++) { + mem_access::store_shared(reduce_buffer + elems * warp_arg.thread_rank() + i, + data + i); + } + } + + // Synchronization inside block-uniform conditional is safe + tb.sync(); + +#pragma unroll + for (int i = 0; i < elems; i++) { + mem_access::load_shared(data + i, reduce_buffer + _warp_rank() * elems + i); + } + } +} + +/* +Main API implementations. For the most part, they just convert the individual +variables into arrays, which makes working with them easier with a single +implementation. In theory, we could use the `_block` implementation as another +option, but the nature of using a pointer is a little less safe and this allows +us to obfuscate the details of the partitioned implementation. +*/ +template +DS_D_INLINE void block(cg::thread_block& tb, cg::thread_block_tile& warp, float& val) +{ + _block(tb, warp, &val); +} + +template +DS_D_INLINE void block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2) +{ + float data[2] = {val1, val2}; + _block(tb, warp, data); + val1 = data[0]; + val2 = data[1]; +} + +template +DS_D_INLINE void block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3) +{ + float data[3] = {val1, val2, val3}; + _block(tb, warp, data); + val1 = data[0]; + val2 = data[1]; + val3 = data[2]; +} + +template +DS_D_INLINE void block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3, + float& val4) +{ + float data[4] = {val1, val2, val3, val4}; + _block(tb, warp, data); + val1 = data[0]; + val2 = data[1]; + val3 = data[2]; + val4 = data[3]; +} + +/* +Note: for the partitioned blocks, the implementation does not support non-power of 2 blocks in order +to shorten block scale reduction length. +*/ +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val) +{ + if (num_threads <= hw_warp_size) { + _warp(warp, &val); + } else { + constexpr int num_warps = num_threads / hw_warp_size; + _block(tb, warp, &val); + } +} + +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2) +{ + float data[2] = {val1, val2}; + + if (num_threads <= hw_warp_size) { + _warp(warp, data); + } else { + constexpr int num_warps = num_threads / hw_warp_size; + _block(tb, warp, data); + } + + val1 = data[0]; + val2 = data[1]; +} + +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3) +{ + float data[3] = {val1, val2, val3}; + + if (num_threads <= hw_warp_size) { + _warp(warp, data); + } else { + constexpr int num_warps = num_threads / hw_warp_size; + _block(tb, warp, data); + } + + val1 = data[0]; + val2 = data[1]; + val3 = data[2]; +} + +template +DS_D_INLINE void partitioned_block(cg::thread_block& tb, + cg::thread_block_tile& warp, + float& val1, + float& val2, + float& val3, + float& val4) +{ + float data[4] = {val1, val2, val3, val4}; + + if (num_threads <= hw_warp_size) { + _warp(warp, data); + } else { + constexpr int num_warps = num_threads / hw_warp_size; + _block(tb, warp, data); + } + + val1 = data[0]; + val2 = data[1]; + val3 = data[2]; + val4 = data[3]; +} + +/* +Arg-reduce is a specialization of the above. We only support this with a single reduction +parameter. This only works for max/min reductions. +*/ + +__align__(8) struct IdxReduceResult { + /* + NOTE: ORDERING MATTERS HERE! The idx is the least significant set of bits + and the val is the most significant. Changing the order of this declaration + will break the code. + */ + int idx; + float val; +}; + +template +DS_D_INLINE IdxReduceResult +idx_reduce(cg::thread_block& tb, cg::thread_block_tile& warp, float val, int idx) +{ + IdxReduceResult res = {idx, val}; + + // Clear out the nan. This shouldn't be an issue for our initial applications + if (isnan(val)) res.val = init(); + + // Can do float compares as integers. By packing the index into the lower bits + // we can just do a single int64 rather than a branch, compare, and select. + // One side benefit of this is that it is by nature a stable algorithm and + // will always bias ties to the higher index. + int64_t* res_as_int = reinterpret_cast(&res); + + // The way floating point compare works is normally to perform a sign comparison + // and if they match, then do a comparison of the rest of the bits as unsigned + // integers. Since we are bundling these, that means for negative values we need + // to reverse the sort order, which we can do with an XOR. + if (val < 0) { *res_as_int ^= 0x7fffffff00000000; } + + _block(tb, warp, res_as_int); + + // Sign bit is preserved, so we can check if we need to invert the mantissa back + if (res.val < 0) { *res_as_int ^= 0x7fffffff00000000; } + + return res; +} + +} // namespace reduce diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/simd.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/simd.h new file mode 100644 index 0000000000000000000000000000000000000000..a205026ec7c1155faf565739e7a407fcfab08c41 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/simd.h @@ -0,0 +1,299 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#if (__x86_64__ || __i386__) +#include +#include +#endif + +#define TILE (128 * 1024 * 1024) +#if defined(__AVX512__) or defined(__AVX256__) +#include + +template +inline T readAs(const void* src) +{ + T res; + std::memcpy(&res, src, sizeof(T)); + return res; +} +template +inline void writeAs(void* dst, const T& val) +{ + std::memcpy(dst, &val, sizeof(T)); +} + +#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) + +#if defined(__AVX512__) +#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d) +#define SIMD_LOAD(x) _mm512_loadu_ps(x) +#define SIMD_SET(x) _mm512_set1_ps(x) +#define SIMD_ADD(x, y) _mm512_add_ps(x, y) +#define SIMD_MUL(x, y) _mm512_mul_ps(x, y) +#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) +#define SIMD_SQRT(x) _mm512_sqrt_ps(x) +#define SIMD_DIV(x, y) _mm512_div_ps(x, y) +#define SIMD_AND(x, y) _mm512_and_ps(x, y) +#define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) +#define SIMD_OR(x, y) _mm512_or_ps(x, y) +#define SIMD_XOR(x, y) _mm512_xor_ps(x, y) +#define SIMD_WIDTH 16 + +static __m512 load_16_bf16_as_f32(const void* data) +{ + __m256i a = readAs<__m256i>(data); // use memcpy to avoid aliasing + __m512i b = _mm512_cvtepu16_epi32(a); // convert 8 u16 to 8 u32 + __m512i c = _mm512_slli_epi32(b, 16); // logical shift left of all u32 by + // 16 bits (representing bf16->f32) + return readAs<__m512>(&c); // use memcpy to avoid aliasing +} + +static void store_16_f32_as_bf16_nearest(__m512 v, void* data) +{ + __m512i u32 = readAs<__m512i>(&v); + + // flow assuming non-nan: + + // uint32_t rounding_bias = ((U32 >> 16) & 1) + UINT32_C(0x7FFF); + __m512i b = _mm512_srli_epi32(u32, 16); + __m512i lsb_mask = _mm512_set1_epi32(0x00000001); + __m512i c = _mm512_and_si512(b, lsb_mask); + __m512i bias_constant = _mm512_set1_epi32(0x00007fff); + __m512i rounding_bias = _mm512_add_epi32(c, bias_constant); + + // uint16_t res = static_cast((U32 + rounding_bias) >> 16); + __m512i d = _mm512_add_epi32(u32, rounding_bias); + __m512i e = _mm512_srli_epi32(d, 16); + __m256i non_nan_res = _mm512_cvtusepi32_epi16(e); + + // handle nan (exp is all 1s and mantissa != 0) + // if ((x & 0x7fffffffU) > 0x7f800000U) + __m512i mask_out_sign = _mm512_set1_epi32(0x7fffffff); + __m512i non_sign_bits = _mm512_and_si512(u32, mask_out_sign); + __m512i nan_threshold = _mm512_set1_epi32(0x7f800000); + __mmask16 nan_mask = _mm512_cmp_epi32_mask(non_sign_bits, nan_threshold, _MM_CMPINT_GT); + + // mix in results with nans as needed + __m256i nans = _mm256_set1_epi16(0x7fc0); + __m256i res = _mm256_mask_mov_epi16(non_nan_res, nan_mask, nans); + + writeAs(data, res); +} +#define SIMD_LOAD_BF16(x) load_16_bf16_as_f32(x) +#define SIMD_STORE_BF16(x, d) store_16_f32_as_bf16_nearest(d, x) + +#define SIMD_LOAD_FP16(x) _mm512_cvtph_ps(_mm256_castps_si256(_mm256_loadu_ps(x))) +#define SIMD_STORE_FP16(x, d) \ + _mm256_store_ps(x, _mm256_castsi256_ps(_mm512_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) + +#define INTV __m256i +#elif defined(__AVX256__) +#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d) +#define SIMD_LOAD(x) _mm256_loadu_ps(x) +#define SIMD_SET(x) _mm256_set1_ps(x) +#define SIMD_ADD(x, y) _mm256_add_ps(x, y) +#define SIMD_MUL(x, y) _mm256_mul_ps(x, y) +#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) +#define SIMD_SQRT(x) _mm256_sqrt_ps(x) +#define SIMD_DIV(x, y) _mm256_div_ps(x, y) +#define SIMD_AND(x, y) _mm256_and_ps(x, y) +#define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) +#define SIMD_OR(x, y) _mm256_or_ps(x, y) +#define SIMD_XOR(x, y) _mm256_xor_ps(x, y) +#define SIMD_WIDTH 8 + +#define SIMD_LOAD_BF16(x) static_assert(false && "AVX256 does not support BFloat16") +#define SIMD_STORE_BF16(x, d) static_assert(false && "AVX256 does not support BFloat16") +#define SIMD_LOAD_FP16(x) _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*)x)) +#define SIMD_STORE_FP16(x, d) \ + _mm_store_ps(x, _mm_castsi128_ps(_mm256_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) + +#define INTV __m128i +#endif + +union AVX_Data { +#if defined(__AVX512__) + __m512 data; +#elif defined(__AVX256__) + __m256 data; +#endif + // float data_f[16]; +}; + +template +inline typename std::enable_if_t, void> simd_store(T* dst, + AVX_Data* src) +{ + size_t width = SIMD_WIDTH; +#pragma unroll + for (size_t i = 0; i < span; ++i) { SIMD_STORE_FP16((float*)(dst + width * i), src[i].data); } +} + +template +inline typename std::enable_if_t, void> simd_store(T* dst, + AVX_Data* src) +{ +#ifdef __AVX512__ + size_t width = SIMD_WIDTH; +#pragma unroll + for (size_t i = 0; i < span; ++i) { SIMD_STORE_BF16((float*)(dst + width * i), src[i].data); } +#else + throw std::runtime_error("AVX512 required for BFloat16"); +#endif +} + +template +inline typename std::enable_if_t, void> simd_store(T* dst, AVX_Data* src) +{ + size_t width = SIMD_WIDTH; +#pragma unroll + for (size_t i = 0; i < span; ++i) { SIMD_STORE(dst + width * i, src[i].data); } +} + +template +inline typename std::enable_if_t, void> simd_load(AVX_Data* dst, + T* src) +{ + size_t width = SIMD_WIDTH; +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD_FP16((float*)(src + width * i)); } +} + +template +inline typename std::enable_if_t, void> simd_load(AVX_Data* dst, + T* src) +{ +#ifdef __AVX512__ + size_t width = SIMD_WIDTH; +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD_BF16((float*)(src + width * i)); } +#else + throw std::runtime_error("AVX512 required for BFloat16"); +#endif +} + +template +inline typename std::enable_if_t, void> simd_load(AVX_Data* dst, T* src) +{ + size_t width = SIMD_WIDTH; +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD(src + width * i); } +} + +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data* src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a[i].data); + } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a.data); + } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Data* src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r[i].data, src_a[i].data); + } +} +template +inline void simd_sqrt(AVX_Data* dst, AVX_Data* src) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_SQRT(src[i].data); } +} +template +inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r[i].data); + } +} +template +inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r[i].data); } +} + +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/softmax.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/softmax.h new file mode 100644 index 0000000000000000000000000000000000000000..b6320a3f91bec40b63d7f96390a6031a61a9d2cc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/softmax.h @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include "custom_cuda_layers.h" + +#include + +using namespace std; + +template +class Softmax { +public: + struct Config { + size_t batchSize; + size_t heads; + size_t seq_length; + size_t prob_depth; + float temperature; + bool mem_alloc; + Config(size_t batch, size_t h, size_t seq, int prob_size = 0, bool mem_alloc = false) + : batchSize(batch), + heads(h), + seq_length(seq), + prob_depth(prob_size), + temperature(1.0), + mem_alloc(mem_alloc) + { + } + }; + + Softmax(Config config) : config_(config) {} + + ~Softmax() {} + + void Forward(int bsz, T* vals, const T* attn_mask, cudaStream_t& stream) + { + launch_attn_softmax(vals, attn_mask, bsz, config_.heads, config_.seq_length, stream); + } + + void Backward(int bsz, T* out_grad, const T* soft_out, cudaStream_t stream) + { + launch_attn_softmax_backward_v2( + out_grad, soft_out, bsz, config_.heads, config_.seq_length, stream); + } + + inline size_t GetProbDepth() const { return config_.prob_depth; } + + inline size_t GetBatchSize() const { return config_.batchSize; } + + inline size_t GetNumHeads() const { return config_.heads; } + + inline size_t GetSeqLength() const { return config_.seq_length; } + + inline void SetSeqLength(size_t seq_len) { config_.seq_length = seq_len; } + +private: + Config config_; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/strided_batch_gemm.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/strided_batch_gemm.h new file mode 100644 index 0000000000000000000000000000000000000000..9767fcf589b82356230016abc2d183b74ab621e2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/strided_batch_gemm.h @@ -0,0 +1,205 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include "context.h" + +template +class StridedBatchGemm { +public: + struct Config { + int batch_size; + int m; + int n; + int k; + float alpha; + float beta; + cublasOperation_t op_A; + cublasOperation_t op_B; + std::array gemm_algos; + + Config(int batch, + int mm, + int nn, + int kk, + float param_alpha, + float param_beta, + cublasOperation_t opA, + cublasOperation_t opB, + const std::array& algos) + : batch_size(batch), + m(mm), + n(nn), + k(kk), + alpha(param_alpha), + beta(param_beta), + op_A(opA), + op_B(opB), + gemm_algos(algos) + { + } + void SetConfig(int mm, int nn, int kk) + { + m = mm; + n = nn; + k = kk; + } + }; + + StridedBatchGemm(const Config& config) : _config(config) {} + + virtual ~StridedBatchGemm() {} + + void Forward(int bsz, T* output, const T* _buffer_a, const T* _buffer_b, cublasHandle_t handle) + { + int stride_a = _config.m * _config.k; + int stride_b = _config.n * _config.k; + int stride_c = _config.m * _config.n; + + cublas_strided_batched_gemm(handle, + _config.m, + _config.n, + _config.k, + &_config.alpha, + &_config.beta, + _buffer_a, + _buffer_b, + output, + _config.op_A, + _config.op_B, + stride_a, + stride_b, + stride_c, + bsz, +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(_config.gemm_algos[0])); +#else + cublasGemmAlgo_t(_config.gemm_algos[0])); +#endif + } + + void ForwardPlusSave(T* output, const T* _buffer_a, const T* _buffer_b, cublasHandle_t handle) + { + int stride_a = _config.m * _config.k; + int stride_b = _config.n * _config.k; + int stride_c = _config.m * _config.n; + + cublas_strided_batched_gemm(handle, + _config.m, + _config.n, + _config.k, + &_config.alpha, + &_config.beta, + _buffer_a, + _buffer_b, + output, + _config.op_A, + _config.op_B, + stride_a, + stride_b, + stride_c, + _config.batch_size, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(_config.gemm_algos[0])); +#else + cublasGemmAlgo_t(_config.gemm_algos[0])); +#endif + + k_buf = _buffer_a; + q_buf = _buffer_b; + } + + void Backward(int bsz, + const T* d_output, + const T* _buffer_a, + const T* _buffer_b, + cublasHandle_t handle, + T* inpGradA = nullptr, + T* inpGradB = nullptr) + { + int mb = (_config.op_A == CUBLAS_OP_T ? _config.k : _config.m); + int kb = (_config.op_A == CUBLAS_OP_T ? _config.m : _config.k); + + int stride_a = mb * _config.n; + int stride_b = _config.n * kb; + int stride_c = _config.m * _config.k; + + // B need to transpose. + cublasOperation_t op_b = (_config.op_B == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); + + // Calculate d_A. + cublas_strided_batched_gemm(handle, + mb, + kb, + _config.n, + &_config.alpha, + &_config.beta, + (_config.op_A == CUBLAS_OP_T ? _buffer_b : d_output), + (_config.op_A == CUBLAS_OP_T ? d_output : _buffer_b), + inpGradA, + CUBLAS_OP_N, + op_b, + stride_a, + stride_b, + stride_c, + bsz, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(_config.gemm_algos[1])); +#else + cublasGemmAlgo_t(_config.gemm_algos[1])); +#endif + + // A need to transpose. + cublasOperation_t op_a = (_config.op_A == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); + + stride_a = _config.m * _config.k; + stride_b = _config.m * _config.n; + stride_c = _config.n * _config.k; + + // Calculate d_B. + cublas_strided_batched_gemm(handle, + _config.k, + _config.n, + _config.m, + &_config.alpha, + &_config.beta, + _buffer_a, + d_output, + inpGradB, + op_a, + CUBLAS_OP_N, + stride_a, + stride_b, + stride_c, + bsz, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo(_config.gemm_algos[2])); +#else + cublasGemmAlgo_t(_config.gemm_algos[2])); +#endif + } + + inline int GetN() const { return _config.k; } + + inline const T* GetBufferA() const { return k_buf; } + + inline const T* GetBufferB() const { return q_buf; } + + inline void SetConfig(int m, int n, int k) { _config.SetConfig(m, n, k); } + +private: + Config _config; + const T* q_buf; + const T* k_buf; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/type_shim.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/type_shim.h new file mode 100644 index 0000000000000000000000000000000000000000..78ade5a06daa12cf628e9966ce2ba8424a6e96ba --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/type_shim.h @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* Taken from NVIDIA/apex commit 855808f3fc268e9715d613f3c2e56469d8c986d8 */ +#include + +// Forward/backward compatibility hack around +// https://github.com/pytorch/pytorch/commit/3aeb78079bcd68282fe9117088e138b77318e288 +// pending more future-proof guidance from upstream. +// struct TypeShim +// { +// const at::Type& payload; +// TypeShim(const at::Type& type) : payload(type) {} +// // Enable trivial conversion to a const at::Type& for pre-3aeb78 +// operator const at::Type&(){ return payload; }; +// // Enable dispatch switch statements to take *this directly for post-3aeb78 +// //operator at::ScalarType(){ return payload.; }; +// }; + +#define DISPATCH_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t_##LEVEL = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t_##LEVEL = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +#define DISPATCH_DOUBLE_AND_FLOAT(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +template +__device__ __forceinline__ T +reduce_block_into_lanes(T* x, + T val, + int lanes = 1, + bool share_result = false) // lanes is intended to be <= 32. +{ + int tid = threadIdx.x + threadIdx.y * blockDim.x; + int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32. + + if (blockSize >= 64) { + x[tid] = val; + __syncthreads(); + } + +#pragma unroll + for (int i = (blockSize >> 1); i >= 64; i >>= 1) { + if (tid < i) x[tid] = x[tid] + x[tid + i]; + __syncthreads(); + } + + T final; + + if (tid < 32) { + if (blockSize >= 64) + final = x[tid] + x[tid + 32]; + else + final = val; + // __SYNCWARP(); + +#pragma unroll + for (int i = 16; i >= lanes; i >>= 1) + final = final + __shfl_down_sync(0xffffffff, final, i); + } + + if (share_result) { + if (tid < lanes) x[tid] = final; // EpilogueOp + // Make sure the smem result is visible to all warps. + __syncthreads(); + } + + return final; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c629b93517d278f586613dd1d32cd4ac4dc2867a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +// CUDA forward declaration +void fused_lamb_cuda(at::Tensor& p, + at::Tensor& p_copy, + at::Tensor& m, + at::Tensor& v, + at::Tensor& g, + float lr, + float beta1, + float beta2, + float max_coeff, + float min_coeff, + float eps, + float grad_scale, + int step, + int mode, + int bias_correction, + float decay, + at::Tensor& w_l2_i, + at::Tensor& u_l2_i, + at::Tensor& lamb_coeff_val); + +#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) + +// C++ interface +at::Tensor lamb(at::Tensor& p, + at::Tensor& p_copy, + at::Tensor& m, + at::Tensor& v, + at::Tensor& g, + float lr, + float beta1, + float beta2, + float max_coeff, + float min_coeff, + float eps, + float grad_scale, + int step, + int mode, + int bias_correction, + float decay) +{ + CHECK_INPUT(p); + if (p_copy.numel() > 0) CHECK_INPUT(p_copy); + CHECK_INPUT(m); + CHECK_INPUT(v); + CHECK_INPUT(g); + int64_t num_elem = p.numel(); + AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal"); + AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal"); + AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal"); + AT_ASSERTM( + p_copy.numel() == num_elem || p_copy.numel() == 0, + "number of elements in p_copy and p tensors should be equal, or p_copy should be empty"); + + // intermediate for weight L2 reduction + // make sure that the threads per block is at least 512 during the kernel launch otherwise the + // behaviour is unexpected + at::Tensor w_l2_i = at::empty( + {512}, + p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float + : p.type().scalarType())); + + // intermediate for update L2 reduction + // make sure that the threads per block is at least 512 during the kernel launch otherwise the + // behaviour is unexpected + at::Tensor u_l2_i = at::empty( + {512}, + p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float + : p.type().scalarType())); + + at::Tensor lamb_coeff_val = at::empty( + {1}, + p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float + : p.type().scalarType())); + + fused_lamb_cuda(p, + p_copy, + m, + v, + g, + lr, + beta1, + beta2, + max_coeff, + min_coeff, + eps, + grad_scale, + step, + mode, + bias_correction, + decay, + w_l2_i, + u_l2_i, + lamb_coeff_val); + + return lamb_coeff_val; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("lamb", &lamb, "Adam optimized CUDA implementation with LAMB."); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..d9bacae73457fe1e757f1f6410d64d18ebcf31bb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu @@ -0,0 +1,478 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "ATen/ATen.h" +#include "ATen/TensorUtils.h" +#include "ATen/cuda/CUDAContext.h" +#include "ATen/cuda/detail/IndexUtils.cuh" +// #include "ATen/Type.h" +#include "ATen/AccumulateType.h" + +#include + +// #include +#if defined(__HIP_PLATFORM_AMD__) && HIP_VERSION > 305 +#include +#else +#include +#endif +#include +#include + +namespace cg = cooperative_groups; + +// Utility class used to avoid linker errors with extern +// unsized shared memory arrays with templated type +namespace { +// This is the un-specialized struct. Note that we prevent instantiation of this +// struct by putting an undefined symbol in the function body so it won't compile. +template +struct SharedMemory { + // Ensure that we won't compile any un-specialized types + __device__ inline operator T*() + { +#ifndef _WIN32 + extern __device__ void error(void); + error(); +#endif + return NULL; + } +}; + +template <> +struct SharedMemory { + __device__ inline operator float*() + { + extern __shared__ float s_float[]; + return s_float; + } +}; + +template <> +struct SharedMemory { + __device__ inline operator double*() + { + extern __shared__ double s_double[]; + return s_double; + } +}; +} // namespace + +#include "type_shim.h" + +typedef enum { + ADAM_MODE_0 = 0, // eps under square root + ADAM_MODE_1 = 1 // eps outside square root +} adamMode_t; + +// s_a and s_b are in shared memory +// g_a and g_b are in shared memory +template +__device__ void reduce_block_in_shared_memory(T* s_a, T* s_b, T* g_a, T* g_b) +{ + // Handle to thread block group + cg::thread_block cta = cg::this_thread_block(); + + // perform block reduction in shared memory, + unsigned int tid = cta.thread_rank(); + + T a_sum = s_a[tid]; + T b_sum = s_b[tid]; + + cg::sync(cta); + + // do reduction in shared mem + if ((blockSize >= 512) && (tid < 256)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 256]; + s_b[tid] = b_sum = b_sum + s_b[tid + 256]; + } + + cg::sync(cta); + + if ((blockSize >= 256) && (tid < 128)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 128]; + s_b[tid] = b_sum = b_sum + s_b[tid + 128]; + } + + cg::sync(cta); + + if ((blockSize >= 128) && (tid < 64)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 64]; + s_b[tid] = b_sum = b_sum + s_b[tid + 64]; + } + + cg::sync(cta); + +#if (__CUDA_ARCH__ >= 300) || (defined(__HIP_PLATFORM_AMD__) && HIP_VERSION >= 502) + if (tid < 32) { + cg::coalesced_group active = cg::coalesced_threads(); + + // Fetch final intermediate sum from 2nd warp + if (blockSize >= 64) { + a_sum = a_sum + s_a[tid + 32]; + b_sum = b_sum + s_b[tid + 32]; + } + + // Reduce final warp using shuffle + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + a_sum += active.shfl_down(a_sum, offset); + b_sum += active.shfl_down(b_sum, offset); + } + } +#else + if ((blockSize >= 64) && (tid < 32)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 32]; + s_b[tid] = b_sum = b_sum + s_b[tid + 32]; + } + + cg::sync(cta); + + if ((blockSize >= 32) && (tid < 16)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 16]; + s_b[tid] = b_sum = b_sum + s_b[tid + 16]; + } + + cg::sync(cta); + + if ((blockSize >= 16) && (tid < 8)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 8]; + s_b[tid] = b_sum = b_sum + s_b[tid + 8]; + } + + cg::sync(cta); + + if ((blockSize >= 8) && (tid < 4)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 4]; + s_b[tid] = b_sum = b_sum + s_b[tid + 4]; + } + + cg::sync(cta); + + if ((blockSize >= 4) && (tid < 2)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 2]; + s_b[tid] = b_sum = b_sum + s_b[tid + 2]; + } + + cg::sync(cta); + + if ((blockSize >= 2) && (tid < 1)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 1]; + s_b[tid] = b_sum = b_sum + s_b[tid + 1]; + } + + cg::sync(cta); + +#endif + + // write result for this block to global mem + if (tid == 0) { + g_a[blockIdx.x] = (T)a_sum; + g_b[blockIdx.x] = (T)b_sum; + } +} + +template +__device__ void reduce_two_vectors_in_register(T a, T b, T* g_a, T* g_b) +{ + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + + T* s_a = SharedMemory(); + T* s_b = SharedMemory() + cg::this_thread_block().size(); + + s_a[threadIdInBlock] = a; + s_b[threadIdInBlock] = b; + + reduce_block_in_shared_memory(s_a, s_b, g_a, g_b); +} + +template +__global__ void lamb_cuda_kernel_part1( + T* __restrict__ p, + GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed + T* __restrict__ m, + T* __restrict__ v, + const GRAD_T* __restrict__ g, + const float b1, + const float b2, + const float eps, + const float grad_scale, + const float step_size, + const size_t tsize, + adamMode_t mode, + const float decay, + T* __restrict__ w_l2_i, + T* __restrict__ u_l2_i) +{ + // Assuming 2D grids and 2D blocks + const int blockId = gridDim.x * blockIdx.y + blockIdx.x; + const int threadsPerBlock = blockDim.x * blockDim.y; + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + const int i = (blockId * threadsPerBlock + threadIdInBlock); + const int totThreads = gridDim.x * gridDim.y * threadsPerBlock; + + T reg_w = 0; + T reg_u = 0; + + for (int j = i; j < tsize; j += totThreads) { + T scaled_grad = g[j] / grad_scale; + T pj = p[j]; + m[j] = b1 * m[j] + (1 - b1) * scaled_grad; + v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad; + float denom; + if (mode == ADAM_MODE_0) + denom = sqrtf(v[j] + eps); + else // Mode 1 + denom = sqrtf(v[j]) + eps; + T update = (m[j] / denom) + (decay * p[j]); + + reg_u += update * update; + reg_w += pj * pj; + } + + reduce_two_vectors_in_register(reg_w, reg_u, w_l2_i, u_l2_i); +} + +template +__global__ void lamb_cuda_kernel_part2(const size_t tsize, T* __restrict__ g_a, T* __restrict__ g_b) +{ + T* s_a = SharedMemory(); + T* s_b = SharedMemory() + cg::this_thread_block().size(); + + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + + s_a[threadIdInBlock] = g_a[threadIdInBlock]; + s_b[threadIdInBlock] = g_b[threadIdInBlock]; + + if (threadIdInBlock >= tsize) { + s_a[threadIdInBlock] = 0.0; + s_b[threadIdInBlock] = 0.0; + } + + reduce_block_in_shared_memory(s_a, s_b, g_a, g_b); +} + +template +__global__ void lamb_cuda_kernel_part3( + T* __restrict__ p, + GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed + T* __restrict__ m, + T* __restrict__ v, + const GRAD_T* __restrict__ g, + const float b1, + const float b2, + const float max_coeff, + const float min_coeff, + const float eps, + const float grad_scale, + const float step_size, + const size_t tsize, + adamMode_t mode, + const float decay, + T* __restrict__ w_l2_i, + T* __restrict__ u_l2_i, + T* __restrict__ lamb_coeff_val) +{ + // Assuming 2D grids and 2D blocks + const int blockId = gridDim.x * blockIdx.y + blockIdx.x; + const int threadsPerBlock = blockDim.x * blockDim.y; + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + const int i = (blockId * threadsPerBlock + threadIdInBlock); + const int totThreads = gridDim.x * gridDim.y * threadsPerBlock; + + T reg_w = sqrtf(w_l2_i[0]); + T reg_u = sqrtf(u_l2_i[0]); + + float lamb_coeff = 1.0; + + if (reg_w != 0 && reg_u != 0) { + lamb_coeff = reg_w / reg_u; + if (lamb_coeff > max_coeff) { lamb_coeff = max_coeff; } + if (lamb_coeff < min_coeff) { lamb_coeff = min_coeff; } + } + + if (blockId == 0 && threadIdInBlock == 0) { + lamb_coeff_val[0] = lamb_coeff; + // printf("Cuda Lamb Coeff is %.6f \n",lamb_coeff); + } + + for (int j = i; j < tsize; j += totThreads) { + T pj = (float)p[j]; + T mj = m[j]; + T vj = v[j]; + float denom; + if (mode == ADAM_MODE_0) + denom = sqrtf(vj + eps); + else // Mode 1 + denom = sqrtf(vj) + eps; + T update = (mj / denom) + (decay * pj); + + pj = pj - (step_size * lamb_coeff * update); + p[j] = pj; + if (p_copy != NULL) p_copy[j] = (GRAD_T)pj; + } +} + +void fused_lamb_cuda(at::Tensor& p, + at::Tensor& p_copy, + at::Tensor& m, + at::Tensor& v, + at::Tensor& g, + float lr, + float beta1, + float beta2, + float max_coeff, + float min_coeff, + float eps, + float grad_scale, + int step, + int mode, + int bias_correction, + float decay, + at::Tensor& w_l2_i, + at::Tensor& u_l2_i, + at::Tensor& lamb_coeff) +{ + // using namespace at; + + // Get tensor size + int tsize = p.numel(); + // Determine #threads and #blocks + const int threadsPerBlock = 512; + int num_blocks = (tsize + threadsPerBlock - 1) / threadsPerBlock; + if (num_blocks > 512) num_blocks = 512; + + int smemsize = 0; + + if (p.type().scalarType() == at::ScalarType::Double) + smemsize = 2 * threadsPerBlock * sizeof(double); + else + smemsize = 2 * threadsPerBlock * sizeof(float); + + const dim3 blocks(num_blocks); + const dim3 threads(threadsPerBlock); + + AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p), + "parameter tensor is too large to be indexed with int32"); + // Constants + float step_size = 0; + if (bias_correction == 1) { + const float bias_correction1 = 1 - std::pow(beta1, step); + const float bias_correction2 = 1 - std::pow(beta2, step); + step_size = lr * std::sqrt(bias_correction2) / bias_correction1; + } else { + step_size = lr; + } + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + if (g.type().scalarType() == at::ScalarType::Half) { + // all other values should be fp32 for half gradients + AT_ASSERTM(p.type().scalarType() == at::ScalarType::Float, + "expected parameter to be of float type"); + // dispatch is done on the gradient type + using namespace at; // prevents "toString is undefined" errors + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + g.scalar_type(), "lamb_cuda_kernel", ([&] { + using accscalar_t = at::acc_type; + + lamb_cuda_kernel_part1 + <<>>( + p.data(), + p_copy.numel() ? p_copy.data() : NULL, + m.data(), + v.data(), + g.data(), + beta1, + beta2, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data()); + + lamb_cuda_kernel_part2 + <<<1, threadsPerBlock, smemsize, stream>>>( + num_blocks, w_l2_i.data(), u_l2_i.data()); + + lamb_cuda_kernel_part3 + <<>>( + p.data(), + p_copy.numel() ? p_copy.data() : NULL, + m.data(), + v.data(), + g.data(), + beta1, + beta2, + max_coeff, + min_coeff, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data(), + lamb_coeff.data()); + })); + } else { + using namespace at; + AT_DISPATCH_FLOATING_TYPES( + g.scalar_type(), "lamb_cuda_kernel", ([&] { + lamb_cuda_kernel_part1 + <<>>( + p.data(), + NULL, // don't output p_copy for fp32, it's wasted write + m.data(), + v.data(), + g.data(), + beta1, + beta2, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data()); + + lamb_cuda_kernel_part2 + <<<1, threadsPerBlock, smemsize, stream>>>( + num_blocks, w_l2_i.data(), u_l2_i.data()); + + lamb_cuda_kernel_part3 + <<>>( + p.data(), + NULL, // don't output p_copy for fp32, it's wasted write + m.data(), + v.data(), + g.data(), + beta1, + beta2, + max_coeff, + min_coeff, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data(), + lamb_coeff.data()); + })); + } + C10_CUDA_CHECK(cudaGetLastError()); +} + +// template __device__ void reduce_two_vectors_in_register(float a, float b, float* g_a, +// float* g_b, cg::grid_group &cgg); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c5cf3e9e9235662a84a4be2c634e3e6d5ebb735b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion.cpp @@ -0,0 +1,13 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_lion.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("lion_update", &ds_lion_step, "DeepSpeed CPU Lion update (C++)"); + m.def("create_lion", &create_lion_optimizer, "DeepSpeed CPU Lion (C++)"); + m.def("destroy_lion", &destroy_lion_optimizer, "DeepSpeed CPU Lion destroy (C++)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion_impl.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion_impl.cpp new file mode 100644 index 0000000000000000000000000000000000000000..6a98162314f9f396f2f58d872405856a18011fc8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion_impl.cpp @@ -0,0 +1,221 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "cpu_lion.h" + +using namespace std::string_literals; +static std::unordered_map> s_optimizers; + +// C++ interface + +template +void Lion_Optimizer::Step_1(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<1>(&rounded_size, _params, grads, _exp_avg, _param_size); +#endif + if (_param_size > rounded_size) { + float betta1_minus1 = 1 - _betta1; + float betta2_minus1 = 1 - _betta2; + + float alpha = _alpha; + float after_decay = 1 - alpha * _weight_decay; + + for (size_t t = rounded_size; t < _param_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > _param_size) copy_size = _param_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t k = t; k < offset; k++) { + float grad = (float)grads[k]; + float param = (float)_params[k]; + float momentum = _exp_avg[k]; + float tmp = momentum * _betta1; + tmp = grad * betta1_minus1 + tmp; + // Rely on portable C++ methods to manipulate the sign bit of a floating-point + // number. + tmp = -std::copysignf(alpha, tmp); + if (_weight_decay > 0) { + param = param * after_decay + tmp; + } else { + param = param + tmp; + } + momentum = momentum * _betta2; + momentum = grad * betta2_minus1 + momentum; + _params[k] = param; + _exp_avg[k] = momentum; + } + } + } +} + +template +void Lion_Optimizer::Step_4(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<4>(&rounded_size, _params, grads, _exp_avg, _param_size); +#endif + if (_param_size > rounded_size) + Step_1((_params + rounded_size), + (grads + rounded_size), + (_exp_avg + rounded_size), + (_param_size - rounded_size)); +} + +int create_lion_optimizer(int optimizer_id, + float alpha, + float betta1, + float betta2, + float weight_decay, + bool should_log) +{ + auto opt = std::make_shared(alpha, betta1, betta2, weight_decay); + + s_optimizers[optimizer_id] = opt; + + if (should_log) { + std::string avx_type = ""; +#if defined(__AVX512__) + avx_type = "AVX512"; +#else +#if defined(__AVX256__) + avx_type = "AVX2"; +#else + avx_type = "scalar"; +#endif +#endif + + printf("Lion Optimizer #%d is created with %s arithmetic capability.\n", + optimizer_id, + avx_type.c_str()); + printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f\n", + alpha, + betta1, + betta2, + weight_decay); + } + + return 0; +} + +template +void Lion_Optimizer::Step_8(ds_params_precision_t* _params, + ds_params_precision_t* grads, + ds_state_precision_t* _exp_avg, + size_t _param_size) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<8>(&rounded_size, _params, grads, _exp_avg, _param_size); +#endif + if (_param_size > rounded_size) + Step_4((_params + rounded_size), + (grads + rounded_size), + (_exp_avg + rounded_size), + (_param_size - rounded_size)); +} + +template +void step_invoker(std::shared_ptr opt, + void* _params, + void* grads, + void* _exp_avg, + size_t _param_size) +{ + opt->Step_8((ds_params_precision_t*)(_params), + (ds_params_precision_t*)(grads), + (ds_state_precision_t*)(_exp_avg), + _param_size); +} + +std::map, + std::function, void*, void*, void*, size_t)>> + invokers; + +// Fill map with template functions for each type +template +void create_invoker() +{ + invokers[std::tuple(c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; +} +struct InvokerInitializer { + InvokerInitializer() + { + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + } +} _invoker_initializer; + +void invoke(std::shared_ptr opt, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + size_t param_size) +{ + c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); + c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype()); + + auto it = invokers.find(std::tuple(params_type, state_type)); + if (it == invokers.end()) { + throw std::runtime_error("Lion optimizer with param type "s + c10::toString(params_type) + + " and state type "s + c10::toString(state_type) + + " is not supported on current hardware"s); + } + + it->second(opt, params.data_ptr(), grads.data_ptr(), exp_avg.data_ptr(), param_size); +} + +int ds_lion_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float weight_decay, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg) +{ + auto params_c = params.contiguous(); + auto grads_c = grads.contiguous(); + auto exp_avg_c = exp_avg.contiguous(); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + opt->IncrementStep(step, beta1, beta2); + opt->update_state(lr, weight_decay); + + invoke(opt, params_c, grads_c, exp_avg_c, params_c.numel()); + + return 0; +} + +int destroy_lion_optimizer(int optimizer_id) +{ + s_optimizers.erase(optimizer_id); + + return 0; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/fused_lion_frontend.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/fused_lion_frontend.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e523f97ca3098444a07b21394cf6a773ddcc700d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/fused_lion_frontend.cpp @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +void multi_tensor_lion_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const int step, + const float weight_decay); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_lion", + &multi_tensor_lion_cuda, + "Compute and apply gradient update to parameters for Lion optimizer"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_apply.cuh b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_apply.cuh new file mode 100644 index 0000000000000000000000000000000000000000..12f41cb49c6bf505db48f1d21e312578f19da836 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_apply.cuh @@ -0,0 +1,132 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +#include +#include "compat.h" + +#include + +// #include + +// This header is the one-stop shop for all your multi-tensor apply needs. + +// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) +constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; +constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; + +template +struct TensorListMetadata { + void* addresses[n][depth_to_max_tensors[n - 1]]; + int sizes[depth_to_max_tensors[n - 1]]; + unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. + int start_tensor_this_launch; +}; + +template +__global__ void multi_tensor_apply_kernel(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) +{ + // Hand the chunk information to the user-supplied functor to process however it likes. + callable(chunk_size, noop_flag, tl, args...); +} + +template +void multi_tensor_apply(int block_size, + int chunk_size, + const at::Tensor& noop_flag, + const std::vector>& tensor_lists, + T callable, + ArgTypes... args) +{ + TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); + int len0 = tensor_lists[0].size(); + TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); + auto ref_device = tensor_lists[0][0].device(); + TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda"); + for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices + { + TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); + for (int t = 0; t < tensor_lists[l].size(); t++) { + // TODO: Print which tensor fails. + bool contiguous_memory = tensor_lists[l][t].is_contiguous(); +#ifdef VERSION_GE_1_5 + contiguous_memory = (contiguous_memory || + tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); +#endif + TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); + TORCH_CHECK(tensor_lists[l][t].device() == ref_device, + "A tensor was not on the same device as the first tensor"); + TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); + } + } + + int ntensors = tensor_lists[0].size(); + + TensorListMetadata tl; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); + auto stream = at::cuda::getCurrentCUDAStream(); + + tl.start_tensor_this_launch = 0; + int loc_block_info = 0; + int loc_tensor_info = 0; + for (int t = 0; t < ntensors; t++) { + tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); + for (int d = 0; d < depth; d++) + tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); + loc_tensor_info++; + + int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; + + for (int chunk = 0; chunk < chunks_this_tensor; chunk++) { + // std::cout << chunks_this_tensor << std::endl; + tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; + tl.block_to_chunk[loc_block_info] = chunk; + loc_block_info++; + + bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && + chunk == chunks_this_tensor - 1); + bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); + bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); + if (tensors_full || blocks_full || last_chunk) { + // using accscalar_t = acc_type; + multi_tensor_apply_kernel<<>>( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); + + AT_CUDA_CHECK(cudaGetLastError()); + + // Reset. The control flow possibilities here make my brain hurt. + loc_block_info = 0; + if (chunk == chunks_this_tensor - 1) { + // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + loc_tensor_info = 0; + tl.start_tensor_this_launch = t + 1; + } else { + // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; + for (int d = 0; d < depth; d++) + tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; + loc_tensor_info = 1; + tl.start_tensor_this_launch = t; + } + } + } + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_lion.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_lion.cu new file mode 100644 index 0000000000000000000000000000000000000000..f5fe6dfdd7d0c3fa9200da8592d0a3f56b138428 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_lion.cu @@ -0,0 +1,126 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +// Another possibility: +// #include + +#include + +#include "multi_tensor_apply.cuh" +#include "type_shim.h" + +#define BLOCK_SIZE 512 +#define ILP 4 + +using MATH_T = float; + +template +struct LionFunctor { + __device__ __forceinline__ void operator()(int chunk_size, + volatile int* noop_gmem, + TensorListMetadata<3>& tl, + const float beta1, + const float beta2, + const float lr, + const float decay) + { + // I'd like this kernel to propagate infs/nans. + // if(*noop_gmem == 1) + // return; + + int tensor_loc = tl.block_to_tensor[blockIdx.x]; + + // potentially use to pass in list of scalar + // int tensor_num = tl.start_tensor_this_launch + tensor_loc; + + int chunk_idx = tl.block_to_chunk[blockIdx.x]; + int n = tl.sizes[tensor_loc]; + + T* g = (T*)tl.addresses[0][tensor_loc]; + g += chunk_idx * chunk_size; + + T* p = (T*)tl.addresses[1][tensor_loc]; + p += chunk_idx * chunk_size; + + T* m = (T*)tl.addresses[2][tensor_loc]; + m += chunk_idx * chunk_size; + + n -= chunk_idx * chunk_size; + + MATH_T after_decay = 1.0f - lr * decay; + + // see note in multi_tensor_scale_kernel.cu + for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) { + MATH_T r_g[ILP]; + MATH_T r_p[ILP]; + MATH_T r_m[ILP]; +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) { + r_g[ii] = g[i]; + r_p[ii] = p[i]; + r_m[ii] = m[i]; + } else { + r_g[ii] = MATH_T(0); + r_p[ii] = MATH_T(0); + r_m[ii] = MATH_T(0); + } + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + MATH_T c = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + MATH_T update = c > 0 ? (-lr) : lr; + r_p[ii] = r_p[ii] * after_decay + update; + r_m[ii] = beta2 * r_m[ii] + (1 - beta2) * r_g[ii]; + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) { + p[i] = r_p[ii]; + m[i] = r_m[ii]; + } + } + } + } +}; + +void multi_tensor_lion_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const int step, + const float weight_decay) +{ + using namespace at; + + // Assume single type across p,g,m1,m2 now + DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), + 0, + "lion", + multi_tensor_apply<3>(BLOCK_SIZE, + chunk_size, + noop_flag, + tensor_lists, + LionFunctor(), + beta1, + beta2, + lr, + weight_decay);) + + AT_CUDA_CHECK(cudaGetLastError()); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..783ffb7a41e396d73e9c0254826012a1ddc6655c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "dequantization_utils.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +template +__global__ void dequantize_kernel(T* __restrict__ dequant_data, + const int8_t* __restrict__ q_data, + const float* __restrict__ q_params, + int elems_per_group, + int total_elems) +{ + dequantize::to_global( + dequant_data, q_data, q_params, elems_per_group, total_elems); +} + +#define LAUNCH_DEQUANT_KERNEL(num_bits, q_type) \ + dequantize_kernel<<>>( \ + dequant_data, q_data, q_params, elems_per_group, total_elems); + +template +void launch_dequantize_kernel(T* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream) +{ + constexpr int unroll = 8; + constexpr int threads = 512; + constexpr int elems_per_block = unroll * threads * dequantize::granularity / (sizeof(T)); + + const dim3 block(threads); + const dim3 grid((total_elems + elems_per_block - 1) / elems_per_block); + + // TODO(cmikeh2): It may make sense to tune unroll, there is perf benefit for large + // problem sizes with this large unroll value. + if (num_bits == 8 && q_type == quantize::Type::Symmetric) { + LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Symmetric); + } else if (num_bits == 8 && q_type == quantize::Type::Asymmetric) { + LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Asymmetric); + } else if (num_bits == 4 && q_type == quantize::Type::Symmetric) { + LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Symmetric); + } else if (num_bits == 4 && q_type == quantize::Type::Asymmetric) { + LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Asymmetric); + } +} + +template void launch_dequantize_kernel(__half* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream); + +template void launch_dequantize_kernel(float* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu new file mode 100644 index 0000000000000000000000000000000000000000..4c08cd4cc3d28d9095813e9190a5034a59c5660e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu @@ -0,0 +1,1028 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "custom_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +__global__ void fake_quantize_kernel(__half* vals, int group_size, int num_bits) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); // tb + cg::thread_block_tile<32> g = + cg::tiled_partition<32>(b); // warp, 32 not optimal for AMD which should be 64. + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(__half); + + __half data[vals_per_access]; + + int group_id = blockIdx.x; + + int thread_index = id * vals_per_access; + int reg_count = 0; + int offset = group_id * group_size; + float max = -10000.0; + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); + +#pragma unroll + for (int i = 0; i < vals_per_access; i++) { + if (abs((float)data[i]) > max) max = abs((float)data[i]); + } + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale = (float)(1 << num_bits) / (2 * max + 1e-5); + float q_scale_inv = 1 / q_scale; + int q_range_max = (1 << (num_bits - 1)) - 1; + int q_range_min = -(1 << (num_bits - 1)); + + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); +#pragma unroll + for (int j = 0; j < vals_per_access; j++) { + float q_data; + q_data = __half2float(data[j]); + q_data = __float2int_rn(q_data * q_scale); + q_data = q_data > (q_range_max) ? (q_range_max) + : (q_data < (q_range_min) ? (q_range_min) : q_data); + data[j] = __float2half_rn(q_data * q_scale_inv); + } + mem_access::store_global(vals + offset + thread_index, data); + } + +#endif +} + +__global__ void fake_quantize_kernel(float* vals, int group_size, int num_bits) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(float); + + float data[vals_per_access]; + + int bid = blockIdx.x; + + int thread_index = id * vals_per_access; + + int reg_count = 0; + + int offset = bid * group_size; + + float max = -10000.0; + + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); + +#pragma unroll + for (int i = 0; i < vals_per_access; i++) { + if (abs(data[i]) > max) max = abs(data[i]); + } + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + + b.sync(); + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale = (1 << num_bits) / (2 * max + 1e-5); + float q_scale_inv = 1 / q_scale; + + int q_range_max = (1 << (num_bits - 1)) - 1; + int q_range_min = -(1 << (num_bits - 1)); + + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); +#pragma unroll + for (int j = 0; j < vals_per_access; j++) { + float q_data; + q_data = __float2int_rn(data[j] * q_scale); + q_data = q_data > (q_range_max) ? (q_range_max) + : (q_data < (q_range_min) ? (q_range_min) : q_data); + data[j] = roundf(q_data * q_scale_inv); + } + mem_access::store_global(vals + offset + thread_index, data); + } +} + +template +void launch_fake_quantize_kernel(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 grid_dim(group_num); + dim3 block_dim(1024); + + fake_quantize_kernel<<>>( + vals, total_count / group_num, num_bits); +} + +template void launch_fake_quantize_kernel(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_fake_quantize_kernel(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +__global__ void sr_fake_quantize_kernel(__half* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + float2* vals_cast = reinterpret_cast(vals); + + __half2 data_low[128]; + __half2 data_high[128]; + + int bid = blockIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + unsigned int tid = threadIdx.x; + int reg_count = 0; + int offset = bid * token_size; + int group_index = bid * token_size + tid; + + int total_count = token_size * token_num; + if (group_index < total_count) { + // float min = 10000.0; + float max = -10000.0; + while (tid < token_size) { + float2 data = vals_cast[offset + tid]; + __half2* data_h = reinterpret_cast<__half2*>(&data); + data_low[reg_count] = data_h[0]; + data_high[reg_count] = data_h[1]; + + float2 data_f[2]; + data_f[0] = __half22float2(data_h[0]); + data_f[1] = __half22float2(data_h[1]); + + if (abs((float)data_f[0].x) > max) max = abs((float)data_f[0].x); + if (abs((float)data_f[0].y) > max) max = abs((float)data_f[0].y); + if (abs((float)data_f[1].x) > max) max = abs((float)data_f[1].x); + if (abs((float)data_f[1].y) > max) max = abs((float)data_f[1].y); + + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale_val = (float)(1 << num_bits) / (max * 2 + 1e-5); + float high_q = (float)((1 << (num_bits - 1)) - 1); + float low_q = (float)(-((1 << (num_bits - 1)))); + + for (int i = 0; i < reg_count; i++) { + int token_index = i * blockDim.x + threadIdx.x; + if (token_index < token_size) { + float2 data_f[2]; + data_f[0] = __half22float2(data_low[i]); + data_f[1] = __half22float2(data_high[i]); + + float2 q_data_int[2]; + q_data_int[0].x = (float)((int)(data_f[0].x * q_scale_val)); + q_data_int[0].y = (float)((int)(data_f[0].y * q_scale_val)); + q_data_int[1].x = (float)((int)(data_f[1].x * q_scale_val)); + q_data_int[1].y = (float)((int)(data_f[1].y * q_scale_val)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = abs(data_f[0].x - (q_data_int[0].x / q_scale_val)) * q_scale_val; + q_error[1] = abs(data_f[0].y - (q_data_int[0].y / q_scale_val)) * q_scale_val; + q_error[2] = abs(data_f[1].x - (q_data_int[1].x / q_scale_val)) * q_scale_val; + q_error[3] = abs(data_f[1].y - (q_data_int[1].y / q_scale_val)) * q_scale_val; + + q_data_int[0].x = + (rand.x < q_error[0] && q_data_int[0].x > low_q && q_data_int[0].x < high_q) + ? (q_data_int[0].x + (data_f[0].x > 0 ? 1 : -1)) + : q_data_int[0].x; + q_data_int[0].y = + (rand.y < q_error[1] && q_data_int[0].y > low_q && q_data_int[0].y < high_q) + ? (q_data_int[0].y + (data_f[0].y > 0 ? 1 : -1)) + : q_data_int[0].y; + q_data_int[1].x = + (rand.w < q_error[2] && q_data_int[1].x > low_q && q_data_int[1].x < high_q) + ? (q_data_int[1].x + (data_f[1].x > 0 ? 1 : -1)) + : q_data_int[1].x; + q_data_int[1].y = + (rand.z < q_error[3] && q_data_int[1].y > low_q && q_data_int[1].y < high_q) + ? (q_data_int[1].y + (data_f[1].y > 0 ? 1 : -1)) + : q_data_int[1].y; + + data_f[0].x = q_data_int[0].x / q_scale_val; + data_f[0].y = q_data_int[0].y / q_scale_val; + data_f[1].x = q_data_int[1].x / q_scale_val; + data_f[1].y = q_data_int[1].y / q_scale_val; + + float2 result; + __half2* result_h = reinterpret_cast<__half2*>(&result); + result_h[0] = __float22half2_rn(data_f[0]); + result_h[1] = __float22half2_rn(data_f[1]); + + vals_cast[offset + token_index] = result; + } + } + } +#endif +} + +__global__ void sr_fake_quantize_kernel(float* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + int idx = blockIdx.x * blockDim.x + id; + + float4* vals_cast = reinterpret_cast(vals); + + float4 data[128]; + + int bid = blockIdx.x; + int tid = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + int group_index = bid * token_size + threadIdx.x; + int reg_count = 0; + int total_count = token_size * token_num; + if (group_index < total_count) { + // float min = 10000.0; + float max = -10000.0; + + while (tid < token_size) { + data[reg_count] = vals_cast[group_index]; + + if (abs(data[reg_count].x) > max) max = abs(data[reg_count].x); + if (abs(data[reg_count].y) > max) max = abs(data[reg_count].y); + if (abs(data[reg_count].z) > max) max = abs(data[reg_count].z); + if (abs(data[reg_count].w) > max) max = abs(data[reg_count].w); + + group_index += blockDim.x; + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale_val = (float)(1 << num_bits) / (max * 2 + 1e-5); + float high_q = (float)((1 << (num_bits - 1)) - 1); + float low_q = (float)(-((1 << (num_bits - 1)))); + + int offset = (bid)*token_size; + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + threadIdx.x; + if (group_index < token_size) { + float4 q_data = data[i]; + + float4 q_data_int; + q_data_int.x = (float)((int)(q_data.x * q_scale_val)); + q_data_int.y = (float)((int)(q_data.y * q_scale_val)); + q_data_int.w = (float)((int)(q_data.w * q_scale_val)); + q_data_int.z = (float)((int)(q_data.z * q_scale_val)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = abs(q_data.x - (q_data_int.x / q_scale_val)) * q_scale_val; + q_error[1] = abs(q_data.y - (q_data_int.y / q_scale_val)) * q_scale_val; + q_error[2] = abs(q_data.w - (q_data_int.w / q_scale_val)) * q_scale_val; + q_error[3] = abs(q_data.z - (q_data_int.z / q_scale_val)) * q_scale_val; + + q_data_int.x = + (rand.x < q_error[0] && q_data_int.x > low_q && q_data_int.x < high_q) + ? (q_data_int.x + (q_data.x > 0 ? 1 : -1)) + : q_data_int.x; + q_data_int.y = + (rand.y < q_error[1] && q_data_int.y > low_q && q_data_int.y < high_q) + ? (q_data_int.y + (q_data.y > 0 ? 1 : -1)) + : q_data_int.y; + q_data_int.w = + (rand.w < q_error[2] && q_data_int.w > low_q && q_data_int.w < high_q) + ? (q_data_int.w + (q_data.w > 0 ? 1 : -1)) + : q_data_int.w; + q_data_int.z = + (rand.z < q_error[3] && q_data_int.z > low_q && q_data_int.z < high_q) + ? (q_data_int.z + (q_data.z > 0 ? 1 : -1)) + : q_data_int.z; + + q_data_int.x /= q_scale_val; + q_data_int.y /= q_scale_val; + q_data_int.w /= q_scale_val; + q_data_int.z /= q_scale_val; + + vals_cast[group_index + offset] = q_data_int; + } + } + } +} + +template +void launch_sr_fake_quantize_kernel(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 block_dim(1024); + dim3 grid_dim(group_num); + + uint64_t inc = total_count / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + + sr_fake_quantize_kernel<<>>( + vals, (total_count / group_num) / 4, group_num, num_bits, seed); +} +template void launch_sr_fake_quantize_kernel(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_sr_fake_quantize_kernel(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +__global__ void fake_quantize_kernel_asym(__half* vals, int group_size, int num_bits) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + float2* vals_cast = reinterpret_cast(vals); + + float2 data[MAX_REG]; + + int group_id = blockIdx.x; + + { + int group_index = id; + int reg_count = 0; + int offset = group_id * group_size; + float max = -10000.0; + float min = 10000.0; + + while (group_index < group_size && reg_count < MAX_REG) { + data[reg_count] = vals_cast[offset + group_index]; + __half* data_h = reinterpret_cast<__half*>(&data[reg_count]); + + if (((float)data_h[0]) > max) max = (float)data_h[0]; + if (((float)data_h[1]) > max) max = (float)data_h[1]; + if (((float)data_h[2]) > max) max = (float)data_h[2]; + if (((float)data_h[3]) > max) max = (float)data_h[3]; + + if (((float)data_h[0]) < min) min = (float)data_h[0]; + if (((float)data_h[1]) < min) min = (float)data_h[1]; + if (((float)data_h[2]) < min) min = (float)data_h[2]; + if (((float)data_h[3]) < min) min = (float)data_h[3]; + + group_index += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale = ((max - min) + 1e-5) / (float)(1 << num_bits); + float q_scale_inv = 1 / q_scale; + + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + id; + if (group_index < group_size) { + __half2* data_h = reinterpret_cast<__half2*>(&data[i]); + float2 q_data[2]; + q_data[0] = __half22float2(data_h[0]); + q_data[1] = __half22float2(data_h[1]); + + float2 q_data_int[2]; + + q_data_int[0].x = roundf((q_data[0].x - min) * q_scale_inv); + q_data_int[0].y = roundf((q_data[0].y - min) * q_scale_inv); + q_data_int[1].x = roundf((q_data[1].x - min) * q_scale_inv); + q_data_int[1].y = roundf((q_data[1].y - min) * q_scale_inv); + + q_data_int[0].x = q_data_int[0].x * q_scale + min; + q_data_int[0].y = q_data_int[0].y * q_scale + min; + q_data_int[1].x = q_data_int[1].x * q_scale + min; + q_data_int[1].y = q_data_int[1].y * q_scale + min; + + data_h[0] = __float22half2_rn(q_data_int[0]); + data_h[1] = __float22half2_rn(q_data_int[1]); + + vals_cast[offset + group_index] = data[i]; + } + } + } +#endif +} + +__global__ void fake_quantize_kernel_asym(float* vals, int group_size, int num_bits) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + float4* vals_cast = reinterpret_cast(vals); + + float4 data[MAX_REG]; + + int bid = blockIdx.x; + + int group_index = bid * group_size + id; + int reg_count = 0; + + float max = -10000.0; + float min = 10000.0; + + while (id < group_size && reg_count < MAX_REG) { + float4 data_reg = vals_cast[group_index]; + data[reg_count] = data_reg; + + if (data_reg.x > max) max = data_reg.x; + if (data_reg.y > max) max = data_reg.y; + if (data_reg.w > max) max = data_reg.w; + if (data_reg.z > max) max = data_reg.z; + + if (data_reg.x < min) min = data_reg.x; + if (data_reg.y < min) min = data_reg.y; + if (data_reg.w < min) min = data_reg.w; + if (data_reg.z < min) min = data_reg.z; + + group_index += blockDim.x; + id += blockDim.x; + reg_count++; + } + id = threadIdx.x; + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale = ((max - min) + 1e-5) / (float)(1 << num_bits); + float q_scale_inv = 1 / q_scale; + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + id; + if (group_index < group_size) { + float4 q_data; + q_data = data[i]; + + float4 q_data_int; + q_data_int.x = roundf((q_data.x - min) * q_scale_inv); + q_data_int.y = roundf((q_data.y - min) * q_scale_inv); + q_data_int.w = roundf((q_data.w - min) * q_scale_inv); + q_data_int.z = roundf((q_data.z - min) * q_scale_inv); + + q_data.x = q_data_int.x * q_scale + min; + q_data.y = q_data_int.y * q_scale + min; + q_data.w = q_data_int.w * q_scale + min; + q_data.z = q_data_int.z * q_scale + min; + + vals_cast[group_index + bid * group_size] = q_data; + } + } +} + +template +void launch_fake_quantize_kernel_asym(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 grid_dim(group_num); + dim3 block_dim(1024); + + fake_quantize_kernel_asym<<>>( + vals, (total_count / group_num) / 4, num_bits); +} + +template void launch_fake_quantize_kernel_asym(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_fake_quantize_kernel_asym(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +__global__ void sr_fake_quantize_kernel_asym(__half* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + float2* vals_cast = reinterpret_cast(vals); + + __half2 data_low[128]; + __half2 data_high[128]; + + int bid = blockIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + unsigned int tid = threadIdx.x; + int reg_count = 0; + int offset = bid * token_size; + int group_index = bid * token_size + tid; + + int total_count = token_size * token_num; + if (group_index < total_count) { + float min = 10000.0; + float max = -10000.0; + while (tid < token_size) { + float2 data = vals_cast[offset + tid]; + __half2* data_h = reinterpret_cast<__half2*>(&data); + data_low[reg_count] = data_h[0]; + data_high[reg_count] = data_h[1]; + + float2 data_f[2]; + data_f[0] = __half22float2(data_h[0]); + data_f[1] = __half22float2(data_h[1]); + + if (((float)data_f[0].x) > max) max = (float)data_f[0].x; + if (((float)data_f[0].y) > max) max = (float)data_f[0].y; + if (((float)data_f[1].x) > max) max = (float)data_f[1].x; + if (((float)data_f[1].y) > max) max = (float)data_f[1].y; + + if (((float)data_f[0].x) < min) min = (float)data_f[0].x; + if (((float)data_f[0].y) < min) min = (float)data_f[0].y; + if (((float)data_f[1].x) < min) min = (float)data_f[1].x; + if (((float)data_f[1].y) < min) min = (float)data_f[1].y; + + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale_val = ((max - min) + 1e-5) / (float)(1 << num_bits); + float q_scale_val_inv = 1 / q_scale_val; + float high_q = (float)((1 << num_bits) - 1); + + for (int i = 0; i < reg_count; i++) { + int token_index = i * blockDim.x + threadIdx.x; + if (token_index < token_size) { + float2 data_f[2]; + data_f[0] = __half22float2(data_low[i]); + data_f[1] = __half22float2(data_high[i]); + + float2 q_data_int[2]; + q_data_int[0].x = (float)((unsigned int)((data_f[0].x - min) * q_scale_val_inv)); + q_data_int[0].y = (float)((unsigned int)((data_f[0].y - min) * q_scale_val_inv)); + q_data_int[1].x = (float)((unsigned int)((data_f[1].x - min) * q_scale_val_inv)); + q_data_int[1].y = (float)((unsigned int)((data_f[1].y - min) * q_scale_val_inv)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = + abs(data_f[0].x - ((q_data_int[0].x * q_scale_val) + min)) * q_scale_val_inv; + q_error[1] = + abs(data_f[0].y - ((q_data_int[0].y * q_scale_val) + min)) * q_scale_val_inv; + q_error[2] = + abs(data_f[1].x - ((q_data_int[1].x * q_scale_val) + min)) * q_scale_val_inv; + q_error[3] = + abs(data_f[1].y - ((q_data_int[1].y * q_scale_val) + min)) * q_scale_val_inv; + + q_data_int[0].x = (rand.x < q_error[0] && q_data_int[0].x < high_q) + ? (q_data_int[0].x + 1) + : q_data_int[0].x; + q_data_int[0].y = (rand.y < q_error[1] && q_data_int[0].y < high_q) + ? (q_data_int[0].y + 1) + : q_data_int[0].y; + q_data_int[1].x = (rand.w < q_error[2] && q_data_int[1].x < high_q) + ? (q_data_int[1].x + 1) + : q_data_int[1].x; + q_data_int[1].y = (rand.z < q_error[3] && q_data_int[1].y < high_q) + ? (q_data_int[1].y + 1) + : q_data_int[1].y; + + data_f[0].x = q_data_int[0].x * q_scale_val + min; + data_f[0].y = q_data_int[0].y * q_scale_val + min; + data_f[1].x = q_data_int[1].x * q_scale_val + min; + data_f[1].y = q_data_int[1].y * q_scale_val + min; + + float2 result; + __half2* result_h = reinterpret_cast<__half2*>(&result); + result_h[0] = __float22half2_rn(data_f[0]); + result_h[1] = __float22half2_rn(data_f[1]); + + vals_cast[offset + token_index] = result; + } + } + } +#endif +} + +__global__ void sr_fake_quantize_kernel_asym(float* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + int idx = blockIdx.x * blockDim.x + id; + + float4* vals_cast = reinterpret_cast(vals); + + float4 data[128]; + + int bid = blockIdx.x; + int tid = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + int group_index = bid * token_size + threadIdx.x; + int reg_count = 0; + int total_count = token_size * token_num; + if (group_index < total_count) { + float min = 10000.0; + float max = -10000.0; + + while (tid < token_size) { + float4 data_reg = vals_cast[group_index]; + data[reg_count] = data_reg; + if (data_reg.x > max) max = data_reg.x; + if (data_reg.y > max) max = data_reg.y; + if (data_reg.w > max) max = data_reg.w; + if (data_reg.z > max) max = data_reg.z; + + if (data_reg.x < min) min = data_reg.x; + if (data_reg.y < min) min = data_reg.y; + if (data_reg.w < min) min = data_reg.w; + if (data_reg.z < min) min = data_reg.z; + + group_index += blockDim.x; + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale_val = ((max - min) + 1e-5) / (float)(1 << num_bits); + float high_q = (float)((1 << num_bits) - 1); + + int offset = (bid)*token_size; + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + threadIdx.x; + if (group_index < token_size) { + float4 q_data = data[i]; + + float4 q_data_int; + q_data_int.x = (float)((int)((q_data.x - min) / q_scale_val)); + q_data_int.y = (float)((int)((q_data.y - min) / q_scale_val)); + q_data_int.w = (float)((int)((q_data.w - min) / q_scale_val)); + q_data_int.z = (float)((int)((q_data.z - min) / q_scale_val)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = abs(q_data.x - ((q_data_int.x * q_scale_val) + min)) / q_scale_val; + q_error[1] = abs(q_data.y - ((q_data_int.y * q_scale_val) + min)) / q_scale_val; + q_error[2] = abs(q_data.w - ((q_data_int.w * q_scale_val) + min)) / q_scale_val; + q_error[3] = abs(q_data.z - ((q_data_int.z * q_scale_val) + min)) / q_scale_val; + + q_data_int.x = (rand.x < q_error[0] && q_data_int.x < high_q) ? (q_data_int.x + 1) + : q_data_int.x; + q_data_int.y = (rand.y < q_error[1] && q_data_int.y < high_q) ? (q_data_int.y + 1) + : q_data_int.y; + q_data_int.w = (rand.w < q_error[2] && q_data_int.w < high_q) ? (q_data_int.w + 1) + : q_data_int.w; + q_data_int.z = (rand.z < q_error[3] && q_data_int.z < high_q) ? (q_data_int.z + 1) + : q_data_int.z; + + q_data_int.x = q_data_int.x * q_scale_val + min; + q_data_int.y = q_data_int.y * q_scale_val + min; + q_data_int.w = q_data_int.w * q_scale_val + min; + q_data_int.z = q_data_int.z * q_scale_val + min; + + vals_cast[group_index + offset] = q_data_int; + } + } + } +} +template +void launch_sr_fake_quantize_kernel_asym(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 block_dim(1024); + dim3 grid_dim(group_num); + + uint64_t inc = total_count / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + + sr_fake_quantize_kernel<<>>( + vals, (total_count / group_num) / 4, group_num, num_bits, seed); +} +template void launch_sr_fake_quantize_kernel_asym(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_sr_fake_quantize_kernel_asym(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..b48eaacd0881c68ce06052e894f9f32797233cc8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp @@ -0,0 +1,404 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "quantization.h" + +template +at::Tensor ds_quantize(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if ((((size / groups) - 1) / 4096 + 1) <= 256) { + launch_fake_quantize_kernel( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +template +at::Tensor ds_sr_quantize(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if (((size / groups) / 4 / 1024) <= 256) { + launch_sr_fake_quantize_kernel( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +template +at::Tensor ds_quantize_asym(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if ((((size / groups) - 1) / 4096 + 1) <= 256) { + launch_fake_quantize_kernel_asym( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +template +at::Tensor ds_sr_quantize_asym(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if (((size / groups) / 4 / 1024) <= 256) { + launch_sr_fake_quantize_kernel_asym( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +std::vector quantize_kernel(at::Tensor& input_vals, + int groups, + int numBits, + quantize::Type quantType) +{ + auto dtype = at::kFloat; + auto params_options = at::TensorOptions() + .dtype(dtype) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int param_elems = (quantize::requires_offset(quantType)) ? 2 : 1; + auto params = torch::empty({groups, param_elems}, params_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output_sizes = input_vals.sizes().vec(); + output_sizes[output_sizes.size() - 1] /= numBits == 8 ? 1 : 2; + auto output = torch::empty(output_sizes, output_options); + + const int elems_per_group = at::numel(input_vals) / groups; + + launch_quant((int8_t*)output.data_ptr(), + (float*)params.data_ptr(), + (__half*)input_vals.data_ptr(), + groups, + elems_per_group, + numBits, + quantType, + at::cuda::getCurrentCUDAStream()); + + return {output, params}; +} + +template +at::Tensor dequantize(at::Tensor& quantized_data, + at::Tensor& params, + int groups, + int num_bits, + quantize::Type quant_type) +{ + auto dtype = (std::is_same::value) ? torch::kFloat32 : torch::kFloat16; + auto output_options = at::TensorOptions() + .dtype(dtype) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output_sizes = quantized_data.sizes().vec(); + output_sizes[output_sizes.size() - 1] *= num_bits == 8 ? 1 : 2; + auto output = torch::empty(output_sizes, output_options); + + const int total_elems = at::numel(output); + const int elems_per_group = total_elems / groups; + + launch_dequantize_kernel((T*)output.data_ptr(), + (const int8_t*)quantized_data.data_ptr(), + (const float*)params.data_ptr(), + quant_type, + num_bits, + elems_per_group, + total_elems, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor dequantize_int4_to_half_experimental(at::Tensor& data_in, + at::Tensor& scale_buffer, + at::Tensor& min_val_buffer, + int num_group, + int group_size) +{ + auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA); + auto output = torch::empty({num_group, group_size}, output_options); + + launch_dequantize_int4_to_half_experimental((uint8_t*)data_in.data_ptr(), + (half*)output.data_ptr(), + (half*)scale_buffer.data_ptr(), + (half*)min_val_buffer.data_ptr(), + num_group, + group_size, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor dequantize_int8_to_half_experimental(at::Tensor& data_in, + at::Tensor& scale_buffer, + at::Tensor& min_val_buffer, + int num_group, + int group_size) +{ + auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA); + auto output = torch::empty({num_group, group_size}, output_options); + + launch_dequantize_int8_to_half_experimental((uint8_t*)data_in.data_ptr(), + (half*)output.data_ptr(), + (half*)scale_buffer.data_ptr(), + (half*)min_val_buffer.data_ptr(), + num_group, + group_size, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +std::vector ds_loco_swizzle_quant(at::Tensor& input_vals, + at::Tensor& error_feedback, + float err_beta, + int groups, + int num_bits, + quantize::Type quant_type, + int pipeline_size, + int nodes, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + auto scales = torch::empty({groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + const int quantization_scalar = 8 / num_bits; + const int compressed_vals = at::numel(input_vals) / quantization_scalar; + + auto output = torch::empty({compressed_vals}, output_options); + const int elems_per_group = at::numel(input_vals) / groups; + + launch_loco_swizzled_quant(reinterpret_cast(output.data_ptr()), + reinterpret_cast(scales.data_ptr()), + reinterpret_cast(input_vals.data_ptr()), + reinterpret_cast<__half*>(error_feedback.data_ptr()), + err_beta, + num_bits, + quant_type, + groups, + elems_per_group, + pipeline_size, + nodes, + devices_per_node, + at::cuda::getCurrentCUDAStream()); + + return {output, scales}; +} + +std::vector ds_swizzle_quant(at::Tensor& input_vals, + int groups, + int num_bits, + quantize::Type quant_type, + int pipeline_size, + int nodes, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + auto scales = torch::empty({groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + const int quantization_scalar = 8 / num_bits; + const int compressed_vals = at::numel(input_vals) / quantization_scalar; + + auto output = torch::empty({compressed_vals}, output_options); + const int elems_per_group = at::numel(input_vals) / groups; + + launch_swizzled_quant((int8_t*)output.data_ptr(), + (float*)scales.data_ptr(), + (__half*)input_vals.data_ptr(), + num_bits, + quant_type, + groups, + elems_per_group, + pipeline_size, + nodes, + devices_per_node, + at::cuda::getCurrentCUDAStream()); + + return {output, scales}; +} + +std::vector quantized_reduction(at::Tensor& input_vals, + at::Tensor& input_scales, + int in_groups, + int out_groups, + int num_bits, + quantize::Type quant_type, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + auto scales = torch::empty({out_groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + std::vector sz(input_vals.sizes().begin(), input_vals.sizes().end()); + sz[sz.size() - 1] = sz.back() / devices_per_node; // num of GPU per nodes + const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node; + auto output = torch::empty(sz, output_options); + + const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node); + const int elems_per_out_group = elems_per_in_tensor / out_groups; + + launch_dequant_reduce((int8_t*)output.data_ptr(), + (float*)scales.data_ptr(), + (const int8_t*)input_vals.data_ptr(), + (const float*)input_scales.data_ptr(), + devices_per_node, + num_bits, + quant_type, + out_groups, + elems_per_out_group, + elems_per_in_tensor, + in_groups / devices_per_node, + elems_per_in_group, + at::cuda::getCurrentCUDAStream()); + return {output, scales}; +} + +std::vector loco_quantized_reduction(at::Tensor& input_vals, + at::Tensor& input_scales, + at::Tensor& error_feedback, + float err_beta, + int in_groups, + int out_groups, + int num_bits, + quantize::Type quant_type, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + + auto scales = torch::empty({out_groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + std::vector sz(input_vals.sizes().begin(), input_vals.sizes().end()); + sz[sz.size() - 1] = sz.back() / devices_per_node; + + const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node; + + auto output = torch::empty(sz, output_options); + + const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node); + const int elems_per_out_group = elems_per_in_tensor / out_groups; + + launch_loco_dequant_reduce((int8_t*)output.data_ptr(), + (float*)scales.data_ptr(), + (const int8_t*)input_vals.data_ptr(), + (const float*)input_scales.data_ptr(), + devices_per_node, + num_bits, + quant_type, + out_groups, + elems_per_out_group, + elems_per_in_tensor, + in_groups / devices_per_node, + elems_per_in_group, + (__half2*)error_feedback.data_ptr(), + err_beta, + at::cuda::getCurrentCUDAStream()); + + return {output, scales}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("ds_quantize_fp32", &ds_quantize, "DeepSpeed Quantize with fp32 (CUDA)"); + m.def("ds_quantize_fp16", &ds_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); + m.def("ds_sr_quantize_fp32", &ds_sr_quantize, "DeepSpeed Quantize with fp32 (CUDA)"); + m.def("ds_sr_quantize_fp16", &ds_sr_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); + m.def("ds_quantize_asym_fp32", &ds_quantize_asym, "DeepSpeed Quantize with fp32 (CUDA)"); + m.def( + "ds_quantize_asym_fp16", &ds_quantize_asym<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); + m.def("ds_sr_quantize_asym_fp32", + &ds_sr_quantize_asym, + "DeepSpeed Quantize with fp32 (CUDA)"); + m.def("ds_sr_quantize_asym_fp16", + &ds_sr_quantize_asym<__half>, + "DeepSpeed Quantize with fp16 (CUDA)"); + pybind11::enum_(m, "QuantizationType") + .value("Symmetric", quantize::Type::Symmetric) + .value("Asymmetric", quantize::Type::Asymmetric) + .export_values(); + m.def("quantize", &quantize_kernel); + m.def("dequantize", &dequantize<__half>); + m.def("dequantize_fp32", &dequantize); + m.def("dequantize_int4_to_half_experimental", + &dequantize_int4_to_half_experimental, + "Dequantize int4 to half (experimental)"); + m.def("dequantize_int8_to_half_experimental", + &dequantize_int8_to_half_experimental, + "Dequantize int8 to half (experimental)"); + m.def("swizzle_quant", &ds_swizzle_quant); + m.def("quantized_reduction", &quantized_reduction); + m.def("loco_swizzle_quant", &ds_loco_swizzle_quant, "LoCo Swizzled Quantization Kernel"); + m.def("loco_quantized_reduction", + &loco_quantized_reduction, + "LoCo Quantization and Reduction Kernel"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu new file mode 100644 index 0000000000000000000000000000000000000000..4100c5174b809e0db4dcf138dfcfde4f42e75843 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu @@ -0,0 +1,557 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "dequantization_utils.h" +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" +#include "quantization_utils.h" +#include "reduction_utils.h" + +using rop = reduce::ROpType; + +/* +TODO(cmikeh2): Add implementation that better handles larger nodes. It would like make sense +to leverage some parallel reductions here to improve performance. +*/ + +template +__global__ void __launch_bounds__(1024) dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // NOTE(cmikeh2): This probably could be hardcoded to a larger number, + // but that means even stronger restrictions on the number of elements per group + // A performance analysis here might be beneficial + constexpr int mem_granularity = (numBits == 8) ? 8 : 4; + constexpr int elems_per_load = mem_granularity / sizeof(int8_t); // div by 1 + constexpr int storage_values = 16 / sizeof(__half2); + + const int block_offset = tb.group_index().x * elems_per_out_group; + const int elem_offset = tb.thread_index().x * elems_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.group_dim().x * elems_per_load; + + __half2 local_buffer[totalChunks * storage_values]; + + quantize::GroupStats stats; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * storage_values; + +#pragma unroll + for (int j = 0; j < storage_values; j++) { + iteration_buffer[j] = reduce::init(); + } + + const int iter_offset = i * stride + base_offset; + const int iter_scale_idx = iter_offset / elems_per_in_group; + bool do_loads = i * stride + elem_offset < elems_per_out_group; + + if (numTensors > 0) { +#pragma unroll + for (int j = 0; j < numTensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } else { +#pragma unroll 4 + for (int j = 0; j < num_tensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } + +#pragma unroll + for (int j = 0; j < storage_values; j++) { stats.update(iteration_buffer[j]); } + } + + auto params = stats.template get_params(tb, warp); + + if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); } + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + const int iter_offset = i * stride + base_offset; + if (i * stride + elem_offset < elems_per_out_group) { + int8_t local_output[elems_per_load]; + quantize::_chunk( + local_output, local_buffer + i * storage_values, params); + mem_access::store_global(reduced_data + iter_offset, local_output); + } + } +} + +template +int32_t pow2_round(int32_t raw_value) +{ + return (((raw_value - 1) >> Power) + 1) << Power; +} + +#define LAUNCH_DEQUANT_REDUCE(num_chunks) \ + dequant_reduce \ + <<>>(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_tensors); + +template +void launch_dequant_reduce_impl(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors, + cudaStream_t stream) +{ + // This is a coincidence. This is derived by 8 halves per 16 bytes with 2-way packing for int4 + constexpr int elems_per_thread = numBits; + const int one_step_threads = + next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread)); + // TODO(cmikeh2): Tune this + const int threads = (one_step_threads < 1024) ? one_step_threads : 1024; + + dim3 block(threads); + dim3 grid(out_groups); + + const int elems_per_step = threads * elems_per_thread; + const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step; + + const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw; + + if (unroll == 1) { + // 0-4096 elems + LAUNCH_DEQUANT_REDUCE(1); + } else if (unroll == 2) { + // 4097-8192 etc... + LAUNCH_DEQUANT_REDUCE(2); + } else if (unroll == 3) { + LAUNCH_DEQUANT_REDUCE(3); + } else if (unroll == 4) { + LAUNCH_DEQUANT_REDUCE(4); + } else if (unroll == 6) { + LAUNCH_DEQUANT_REDUCE(6); + } else if (unroll == 8) { + LAUNCH_DEQUANT_REDUCE(8); + } else if (unroll == 10) { + LAUNCH_DEQUANT_REDUCE(10); + } else if (unroll == 12) { + // 48k limit + LAUNCH_DEQUANT_REDUCE(12); + } else { + assert(false); + } +} + +#define LAUNCH_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \ + launch_dequant_reduce_impl(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + out_groups, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_gpus, \ + stream); + +void launch_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int num_gpus, + int num_bits, + quantize::Type quant_type, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + cudaStream_t stream) +{ + if (quant_type == quantize::Type::Symmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric); + } + } + } else if (quant_type == quantize::Type::Asymmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric); + } + } + } +} + +/* +Modified loco_dequant_reduce function that performs dequantization and reduction, +and incorporates error-feedback by updating the error_feedback tensor in-place. +*/ + +template +__global__ void __launch_bounds__(1024) loco_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors, + __half2* error_feedback, + const float err_beta) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + constexpr int mem_granularity = (numBits == 8) ? 8 : 4; + constexpr int elems_per_load = mem_granularity / sizeof(int8_t); + constexpr int storage_values = 16 / sizeof(__half2); + + const int block_offset = tb.group_index().x * elems_per_out_group; + const int elem_offset = tb.thread_index().x * elems_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.group_dim().x * elems_per_load; + + constexpr int scaling_factor = elems_per_load / storage_values; + const int block_offset_err = block_offset / scaling_factor; + const int elem_offset_err = tb.thread_index().x * storage_values; + const int base_offset_err = block_offset_err + elem_offset_err; + const int stride_err = tb.group_dim().x * storage_values; + + __half2 local_buffer[totalChunks * storage_values]; + __half2 err_buffer[totalChunks * storage_values]; + + quantize::GroupStats stats; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * storage_values; + __half2* iter_err_buffer = err_buffer + i * storage_values; + +#pragma unroll + for (int j = 0; j < storage_values; j++) { + iteration_buffer[j] = reduce::init(); + } + + const int iter_offset = i * stride + base_offset; + const int iter_offset_err = i * stride_err + base_offset_err; + const int iter_scale_idx = iter_offset / elems_per_in_group; + bool do_loads = i * stride + elem_offset < elems_per_out_group; + + if (numTensors > 0) { +#pragma unroll + for (int j = 0; j < numTensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } else { +#pragma unroll 4 + for (int j = 0; j < num_tensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } + mem_access::load_global( + iter_err_buffer, error_feedback + iter_offset_err, do_loads); +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = __hadd2(iteration_buffer[k], iter_err_buffer[k]); + stats.update(iteration_buffer[k]); + } + } + + auto params = stats.template get_params(tb, warp); + + // Initialize dequantization parameters based on params + auto de_params = params; + de_params.scale = 1.0f / params.scale; + if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; } + + if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); } + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + const int iter_offset = i * stride + base_offset; + const int iter_offset_err = i * stride_err + base_offset_err; + __half2* iteration_buffer = local_buffer + i * storage_values; + __half2* iter_err_buffer = err_buffer + i * storage_values; + + if (i * stride + elem_offset < elems_per_out_group) { + // ----------- Begin Error-Feedback Modification ----------- + int8_t local_output[elems_per_load]; + quantize::_chunk(local_output, iteration_buffer, params); + mem_access::store_global(reduced_data + iter_offset, local_output); + + // Dequantize the quantized output to compute the dequantized value + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, local_output, de_params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + // __half2 to float2 + float2 iter_buf_f = __half22float2(iteration_buffer[k]); + float2 dequant_buf_f = __half22float2(dequant_buffer[k]); + + // Update within float precision + float2 new_error_f; + new_error_f.x = iter_buf_f.x - dequant_buf_f.x; + new_error_f.y = iter_buf_f.y - dequant_buf_f.y; + + float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]); + + iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x; + iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y; + + // float2 back to __half2 + iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f); + } + mem_access::store_global(error_feedback + iter_offset_err, + iter_err_buffer); + } + } +} + +#define LAUNCH_LOCO_DEQUANT_REDUCE(num_chunks) \ + loco_dequant_reduce \ + <<>>(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_tensors, \ + error_feedback, \ + err_beta); + +template +void launch_loco_dequant_reduce_impl(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors, + __half2* error_feedback, + const float err_beta, + cudaStream_t stream) +{ + constexpr int elems_per_thread = numBits; + const int one_step_threads = + next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread)); + const int threads = (one_step_threads < 1024) ? one_step_threads : 1024; + + dim3 block(threads); + dim3 grid(out_groups); + + const int elems_per_step = threads * elems_per_thread; + const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step; + + const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw; + + if (unroll == 1) { + LAUNCH_LOCO_DEQUANT_REDUCE(1); + } else if (unroll == 2) { + LAUNCH_LOCO_DEQUANT_REDUCE(2); + } else if (unroll == 3) { + LAUNCH_LOCO_DEQUANT_REDUCE(3); + } else if (unroll == 4) { + LAUNCH_LOCO_DEQUANT_REDUCE(4); + } else if (unroll == 6) { + LAUNCH_LOCO_DEQUANT_REDUCE(6); + } else if (unroll == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE(8); + } else if (unroll == 10) { + LAUNCH_LOCO_DEQUANT_REDUCE(10); + } else if (unroll == 12) { + LAUNCH_LOCO_DEQUANT_REDUCE(12); + } else { + assert(false); + } +} + +#define LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \ + launch_loco_dequant_reduce_impl(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + out_groups, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_gpus, \ + error_feedback, \ + err_beta, \ + stream); + +void launch_loco_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int num_gpus, + int num_bits, + quantize::Type quant_type, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + __half2* error_feedback, + const float err_beta, + cudaStream_t stream) +{ + if (quant_type == quantize::Type::Symmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric); + } + } + } else if (quant_type == quantize::Type::Asymmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric); + } + } + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..b04d0e968ba58a203f2b5e24790053cbd43bc74a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize.cu @@ -0,0 +1,151 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" +#include "quantization.h" +#include "quantization_utils.h" +#include "reduction_utils.h" + +namespace cg = cooperative_groups; + +/* +Pure quantization kernel with no fusion. +*/ +template +__global__ void cached_quantization(int8_t* __restrict__ output_data, + float* __restrict__ params, + const __half* __restrict__ input_data, + int groups, + int elems_per_group) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Indexing offsets + const int block_offset = + (tb.group_index().x * (max_threads / threads_per_group) * elems_per_group) + + (tb.thread_index().y * elems_per_group); + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.size() * quantize::h_per_load; + + const __half* input_base = input_data + base_offset; //.. + + __half2 local_buffer[UNROLL * internal_unroll * quantize::h2_per_load]; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + // Convenience helper, should resolve to register indices and not realize. + __half2* iteration_buffer = local_buffer + i * internal_unroll * quantize::h2_per_load; +#pragma unroll + for (int j = 0; j < internal_unroll; j++) { + const int iteration = i * internal_unroll + j; + mem_access::load_global( + iteration_buffer + j * quantize::h2_per_load, + input_base + iteration * stride, + elem_offset + iteration * stride < elems_per_group); + } + } + + quantize:: + local_array( + local_buffer, params, output_data, elems_per_group, groups); +} + +/********* Launcher methods ***********/ +#define LAUNCH_CACHED_QUANT_CALL(q_bits, quant_type) \ + cached_quantization \ + <<>>(output_data, params, input_data, groups, elems_per_group); + +#define LAUNCH_CACHED_QUANT( \ + q_bits, quant_type, unroll_factor_in, internal_unroll_in, threads_per_group_in) \ + const int unroll_factor = unroll_factor_in; \ + const int internal_unroll_l = internal_unroll_in; \ + const int threads_per_group = threads_per_group_in; \ + if (q_bits == 4) { \ + if (quant_type == quantize::Type::Asymmetric) { \ + LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Asymmetric) \ + } else { \ + LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Symmetric) \ + } \ + } else { \ + if (quant_type == quantize::Type::Asymmetric) { \ + LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Asymmetric) \ + } else { \ + LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Symmetric) \ + } \ + } + +void launch_quant(int8_t* output_data, + float* params, + const __half* input_data, + const int groups, + const int elems_per_group, + const int num_bits, + const quantize::Type quant_type, + cudaStream_t stream) +{ + constexpr int max_threads = 256; + + constexpr int internal_unroll = 2; + + const bool is_subblock_schedule = (elems_per_group <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? quantize::h_per_load + : quantize::h_per_load * internal_unroll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_group + h_per_step - 1) / h_per_step); + const int threads_per_group = (one_step_threads < max_threads) ? one_step_threads : max_threads; + + const int groups_per_block = + is_subblock_schedule ? (max_threads + threads_per_group - 1) / threads_per_group : 1; + const int groups_launch = (groups_per_block + groups - 1) / groups_per_block; + + dim3 block(threads_per_group, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threads_per_group * h_per_step; + const int external_unroll = (elems_per_group + elems_per_step - 1) / elems_per_step; + + if (is_subblock_schedule) { + // <=128 + if (threads_per_group == 1) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 1); + } else if (threads_per_group == 2) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 2); + } else if (threads_per_group == 4) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 4); + } else if (threads_per_group == 8) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 8); + } else if (threads_per_group == 16) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 16); + } + } else if (external_unroll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, internal_unroll, max_threads); + } else if (external_unroll == 2) { + // 4097 - 8192 elems + LAUNCH_CACHED_QUANT(num_bits, quant_type, 2, internal_unroll, max_threads); + } else if (external_unroll == 3) { + // 8193 - 12288 elems + LAUNCH_CACHED_QUANT(num_bits, quant_type, 3, internal_unroll, max_threads); + } else if (external_unroll == 4) { + // 12289 - 16384 elems + LAUNCH_CACHED_QUANT(num_bits, quant_type, 4, internal_unroll, max_threads); + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu new file mode 100644 index 0000000000000000000000000000000000000000..b26151ab5c8c327d7b5b54239b73ee3ebd6bea8e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu @@ -0,0 +1,281 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include "memory_access_utils.h" + +template +struct alignas(sizeof(T) * N) AlignedArray { + using Element = T; + static const int kElements = N; + + __device__ __host__ AlignedArray() {} + + __device__ __host__ AlignedArray(const T& rhs) + { +#pragma unroll + for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = rhs; } + } + + __device__ __host__ T& operator[](int offset) + { + return reinterpret_cast(this->buffer[offset]); + } + + __device__ __host__ const T& operator[](int offset) const + { + return reinterpret_cast(this->buffer[offset]); + } + + __device__ __host__ T& at(int offset) { return reinterpret_cast(this->buffer[offset]); } + + __device__ __host__ const T& at(int offset) const + { + return reinterpret_cast(this->buffer[offset]); + } + + __device__ __host__ AlignedArray operator+(const AlignedArray& rhs) const + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < kElements; ++idx) { ret[idx] = this->at(idx) + rhs.at(idx); } + + return ret; + } + + __device__ __forceinline__ void clear() + { +#pragma unroll + for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = Element(0); } + } + + Element buffer[N]; +}; + +template +struct reduce_max { + __device__ __forceinline__ T operator()(const T& lhs, const T& rhs) + { + return lhs > rhs ? lhs : rhs; + } +}; + +template +struct reduce_min { + __device__ __forceinline__ T operator()(const T& lhs, const T& rhs) + { + return lhs < rhs ? lhs : rhs; + } +}; + +template +struct subtract { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] - rhs; } + + return ret; + } +}; + +template +struct plus { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] + rhs; } + + return ret; + } +}; + +template +struct multiply { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] * rhs; } + + return ret; + } +}; + +template +struct clamp { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& min_val, + const T& max_val) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { + ret[idx] = reduce_max()(reduce_min()(lhs[idx], max_val), min_val); + } + + return ret; + } +}; + +template +struct round_int; + +template +struct round_int { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = hrint(lhs[idx]); } + + return ret; + } +}; + +template +struct divide { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] / rhs; } + + return ret; + } +}; + +template +__device__ __forceinline__ T to_scalar(const AlignedArray& data) +{ + Reducer re; + T res = data[0]; + +#pragma unroll + for (int idx = 1; idx < N; ++idx) { res = re(res, data[idx]); } + + return res; +} + +template +__device__ __forceinline__ AlignedArray int4_to_half( + const AlignedArray& data) +{ + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N * 2; idx += 2) { + ret[idx] = half(int(data[idx / 2] >> 4)); + ret[idx + 1] = half(int(data[idx / 2] & 0xf)); + } + + return ret; +} + +__global__ void dequantize_int4_to_half(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size) +{ + using AccessType = AlignedArray; + using AccessTypeOut = AlignedArray; + + for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8; + idx += blockDim.x * gridDim.x) { + int id_group = idx / (group_size / 8); + AccessType value = reinterpret_cast(data_in)[idx]; + half scale = scale_buffer[id_group]; + half min_value = min_val_buffer[id_group]; + + AccessTypeOut output = int4_to_half(value); + output = divide()(output, scale); + output = plus()(output, min_value); + + reinterpret_cast(data_out)[idx] = output; + } +} + +void launch_dequantize_int4_to_half_experimental(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size, + cudaStream_t stream) +{ + int num_warp = num_group / 4; + int num_block = num_warp / 8; // 256 trd / block + + dequantize_int4_to_half<<>>( + data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size); +} + +template +__device__ __forceinline__ AlignedArray int8_to_half(const AlignedArray& data) +{ + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; idx += 1) { ret[idx] = half(int(data[idx])); } + + return ret; +} + +__global__ void dequantize_int8_to_half(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size) +{ + using AccessType = AlignedArray; + using AccessTypeOut = AlignedArray; + + for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8; + idx += blockDim.x * gridDim.x) { + int id_group = idx / (group_size / 8); + AccessType value = reinterpret_cast(data_in)[idx]; + half scale = scale_buffer[id_group]; + half min_value = min_val_buffer[id_group]; + + AccessTypeOut output = int8_to_half(value); + output = divide()(output, scale); + output = plus()(output, min_value); + + reinterpret_cast(data_out)[idx] = output; + } +} + +void launch_dequantize_int8_to_half_experimental(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size, + cudaStream_t stream) +{ + int num_warp = num_group / 4; + int num_block = num_warp / 8; // 256 trd / block + + dequantize_int8_to_half<<>>( + data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..a4b6096c81af102cab9f71a946024f7117902e2e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu @@ -0,0 +1,427 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "dequantization_utils.h" +#include "memory_access_utils.h" +#include "quantization_utils.h" +#include "reduction_utils.h" + +using rop = reduce::ROpType; + +namespace swiz_quant { +constexpr int max_threads = 512; +constexpr int min_threads = 32; + +constexpr int step_granularity = 2; +constexpr int h_per_step = step_granularity * quantize::h_per_load; +} // namespace swiz_quant + +template +__global__ void swizzled_quant_kernel(int8_t* quantized_data, + float* quantized_scales, + const __half* uncompressed_data, + int elems_per_group, + int nodes, + int devices_per_node) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Indexing offsets, same as normal quantization for in-case + const int block_rank = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; + const int block_offset = block_rank * elems_per_group; + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.size() * quantize::h_per_load; + const __half* input_base = uncompressed_data + base_offset; + + // Local buffer + __half2 local_buffer[totalChunks * quantize::h2_per_load]; + + quantize::GroupStats stats; +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; + + mem_access::load_global( + iteration_buffer, input_base + i * stride, elem_offset + i * stride < elems_per_group); + +#pragma unroll + for (int j = 0; j < quantize::h2_per_load; j++) { stats.update(iteration_buffer[j]); } + } + + auto params = stats.template get_params(tb, warp); + + const int partition_id = blockIdx.z; + const int partition_offset = partition_id / devices_per_node; + const int partition_base = (partition_id % devices_per_node) * nodes; + const int pipelining_offset = blockIdx.y * (devices_per_node * nodes); + const int output_partition = (pipelining_offset + partition_base + partition_offset); + + constexpr int out_scalar_effect = 8 / numBits; + const int out_block_rank = output_partition * gridDim.x + blockIdx.x; + const int out_block_offset = out_block_rank * elems_per_group / out_scalar_effect; + const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect; + int8_t* out_base = quantized_data + out_base_offset; + + const int out_stride = stride / out_scalar_effect; + constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect; + + if (tb.thread_index().x == 0) { params.store(quantized_scales, out_block_rank); } + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + if (i * stride + elem_offset < elems_per_group) { + int8_t local_output[quantize::h_per_load / out_scalar_effect]; + quantize::_chunk( + local_output, local_buffer + i * quantize::h2_per_load, params); + mem_access::store_global(out_base + i * out_stride, local_output); + } + } +} + +#define LAUNCH_SWIZZLE_QUANT(total_chunks, threads) \ + swizzled_quant_kernel<<>>( \ + q_data, q_scales, input_data, elems_per_group, nodes, devices_per_node); + +/* +Swizzled quantization reorganizes the quantized groups in order to better facilitate +communication. As an example of the partitioning scheme we have the following example +of 2 node, 4 device swizzling: + + --- --- --- --- --- --- --- --- +| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | + --- --- --- --- --- --- --- --- +becomes + --- --- --- --- --- --- --- --- +| 0 | 4 | 1 | 5 | 2 | 6 | 3 | 7 | + --- --- --- --- --- --- --- --- + +Multiple quantization groups may be mapped into a single partition. In order to better support +later pipelining, we may also perform an additional slicing. In two-way slicing, for instance, +the first halves of each partition are concatenated. +*/ + +template +void launch_swizzled_quant_impl(int8_t* q_data, + float* q_scales, + const __half* input_data, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + const int one_step_threads = + next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step); + const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads + : swiz_quant::max_threads; + const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads + : max_threads; + + dim3 block(threads); + const int groups_per_partition = groups / (nodes * devices_per_node); + assert(groups_per_partition % pipelining == 0); + const int contiguous_groups = groups_per_partition / pipelining; + const int partitions = nodes * devices_per_node; + dim3 grid(contiguous_groups, pipelining, partitions); + + const int elems_per_step = threads * swiz_quant::h_per_step; + const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step); + const int total_unroll = external_unroll * swiz_quant::step_granularity; + + assert(total_unroll % 2 == 0); + + if (threads == 32) { + LAUNCH_SWIZZLE_QUANT(2, 32); + } else if (threads == 64) { + LAUNCH_SWIZZLE_QUANT(2, 64); + } else if (threads == 128) { + LAUNCH_SWIZZLE_QUANT(2, 128); + } else if (threads == 256) { + LAUNCH_SWIZZLE_QUANT(2, 256); + } else if (threads == 512) { + if (total_unroll == 2) { + LAUNCH_SWIZZLE_QUANT(2, 512); + } else if (total_unroll == 4) { + LAUNCH_SWIZZLE_QUANT(4, 512); + } else if (total_unroll == 6) { + LAUNCH_SWIZZLE_QUANT(6, 512); + } else if (total_unroll == 8) { + LAUNCH_SWIZZLE_QUANT(8, 512); + } else if (total_unroll == 10) { + LAUNCH_SWIZZLE_QUANT(10, 512); + } + } +} + +#define DISPATCH_SWIZZLE_QUANT(num_bits, qtype) \ + launch_swizzled_quant_impl(q_data, \ + q_scales, \ + input_data, \ + groups, \ + elems_per_group, \ + pipelining, \ + nodes, \ + devices_per_node, \ + stream); + +void launch_swizzled_quant(int8_t* q_data, + float* q_scales, + const __half* input_data, + int num_bits, + quantize::Type q_type, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + if (num_bits == 4) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Symmetric); + } + } +} + +template +__global__ void loco_swizzled_quant_kernel(int8_t* quantized_data, + float* quantized_scales, + const __half* uncompressed_data, + __half* error_feedback, + const float err_beta, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Indexing offsets, same as normal quantization for in-case + const int block_rank_data = + blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; + const int block_offset_data = block_rank_data * elems_per_group; + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset_data = block_offset_data + elem_offset; + const int stride = tb.size() * quantize::h_per_load; + const __half* uncompressed_data_base = uncompressed_data + base_offset_data; + + const int partition_id = blockIdx.z; + const int partition_offset = partition_id / devices_per_node; + const int partition_base = (partition_id % devices_per_node) * nodes; + const int pipelining_offset = blockIdx.y * (devices_per_node * nodes); + const int output_partition = (pipelining_offset + partition_base + partition_offset); + const int block_rank_err = output_partition * gridDim.x + blockIdx.x; + + const int block_offset_err = block_rank_err * elems_per_group; + const int base_offset_err = block_offset_err + elem_offset; + __half* error_feedback_base = error_feedback + base_offset_err; + + __half2 local_buffer[totalChunks * quantize::h2_per_load]; + __half2 err_buffer[totalChunks * quantize::h2_per_load]; + + quantize::GroupStats stats; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; + __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load; + const int i_stride = i * stride; + bool do_loads = (elem_offset + i_stride) < elems_per_group; + + mem_access::load_global( + iteration_buffer, uncompressed_data_base + i_stride, do_loads); + + mem_access::load_global( + iter_err_buffer, error_feedback_base + i_stride, do_loads); + +#pragma unroll + for (int j = 0; j < quantize::h2_per_load; j++) { + iteration_buffer[j] = __hadd2(iteration_buffer[j], iter_err_buffer[j]); + stats.update(iteration_buffer[j]); + } + } + + auto params = stats.template get_params(tb, warp); + + // Initialize dequantization parameters based on params + auto de_params = params; + de_params.scale = 1.0f / params.scale; + if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; } + + if (threadIdx.x == 0) { params.store(quantized_scales, block_rank_err); } + + constexpr int out_scalar_effect = 8 / numBits; + const int out_block_offset = block_rank_err * elems_per_group / out_scalar_effect; + const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect; + int8_t* out_base = quantized_data + out_base_offset; + + const int out_stride = stride / out_scalar_effect; + constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + const int i_stride = i * stride; + __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; + __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load; + + if (i_stride + elem_offset < elems_per_group) { + int8_t local_output[quantize::h_per_load / out_scalar_effect]; + quantize::_chunk(local_output, iteration_buffer, params); + mem_access::store_global(out_base + i * out_stride, local_output); + + // Dequantize the quantized output to compute the dequantized value + __half2 dequant_buffer[quantize::h2_per_load]; + dequantize::chunk(dequant_buffer, local_output, de_params); + +// Compute new error: sum - dequant_buffer +#pragma unroll + for (int k = 0; k < quantize::h2_per_load; k++) { + // __half2 to float2 + float2 iter_buf_f = __half22float2(iteration_buffer[k]); + float2 dequant_buf_f = __half22float2(dequant_buffer[k]); + + // Update within float precision + float2 new_error_f; + new_error_f.x = iter_buf_f.x - dequant_buf_f.x; + new_error_f.y = iter_buf_f.y - dequant_buf_f.y; + + float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]); + + iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x; + iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y; + + // float2 back to __half2 + iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f); + } + __half2* error_feedback_base_h2 = reinterpret_cast<__half2*>(error_feedback_base); + mem_access::store_global(error_feedback_base_h2 + i_stride / 2, + iter_err_buffer); + } + } +} + +#define LAUNCH_LOCO_SWIZZLE_QUANT(total_chunks, threads) \ + loco_swizzled_quant_kernel \ + <<>>(output_data, \ + params, \ + input_data, \ + error_feedback, \ + err_beta, \ + groups, \ + elems_per_group, \ + pipelining, \ + nodes, \ + devices_per_node); + +template +void launch_loco_swizzled_quant_impl(int8_t* output_data, + float* params, + const __half* input_data, + __half* error_feedback, + const float err_beta, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + const int one_step_threads = + next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step); + const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads + : swiz_quant::max_threads; + const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads + : max_threads; + + dim3 block(threads); + const int groups_per_partition = groups / (nodes * devices_per_node); + assert(groups_per_partition % pipelining == 0); + const int contiguous_groups = groups_per_partition / pipelining; + const int partitions = nodes * devices_per_node; + dim3 grid(contiguous_groups, pipelining, partitions); + + const int elems_per_step = threads * swiz_quant::h_per_step; + const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step); + const int total_unroll = external_unroll * swiz_quant::step_granularity; + + assert(total_unroll % 2 == 0); + + if (threads == 32) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 32); + } else if (threads == 64) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 64); + } else if (threads == 128) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 128); + } else if (threads == 256) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 256); + } else if (threads == 512) { + if (total_unroll == 2) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 512); + } else if (total_unroll == 4) { + LAUNCH_LOCO_SWIZZLE_QUANT(4, 512); + } else if (total_unroll == 6) { + LAUNCH_LOCO_SWIZZLE_QUANT(6, 512); + } else if (total_unroll == 8) { + LAUNCH_LOCO_SWIZZLE_QUANT(8, 512); + } else if (total_unroll == 10) { + LAUNCH_LOCO_SWIZZLE_QUANT(10, 512); + } + } +} + +#define DISPATCH_LOCO_SWIZZLE_QUANT(num_bits, qtype) \ + launch_loco_swizzled_quant_impl(output_data, \ + params, \ + input_data, \ + error_feedback, \ + err_beta, \ + groups, \ + elems_per_group, \ + pipelining, \ + nodes, \ + devices_per_node, \ + stream); + +void launch_loco_swizzled_quant(int8_t* output_data, + float* params, + const __half* input_data, + __half* error_feedback, + const float err_beta, + int num_bits, + quantize::Type q_type, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + if (num_bits == 4) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Symmetric); + } + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/gather_scatter.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/gather_scatter.cu new file mode 100644 index 0000000000000000000000000000000000000000..291c8eb063b0e0a56843fb43d83caa5fdb00f6cb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/gather_scatter.cu @@ -0,0 +1,186 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "custom_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +namespace td_data { +constexpr int granularity = 16; +} + +template +__global__ void gather_tokens_impl(T* retained_tokens, + const T* activations, + int32_t* gather_indices, + int32_t sampled_tokens, + int32_t channels, + int32_t read_batch_stride, + int32_t read_seq_stride, + int32_t write_batch_stride, + int32_t write_seq_stride) +{ + constexpr int mem_vals_t = td_data::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + + const int gather_idx = gather_indices[tb.group_index().x * sampled_tokens + tb.group_index().y]; + + const int read_offset = read_batch_stride * tb.group_index().x + read_seq_stride * gather_idx; + const int write_offset = + write_batch_stride * tb.group_index().x + write_seq_stride * tb.group_index().y; + + for (int i = tb.thread_index().x * mem_vals_t; i < channels; i += blockDim.x * mem_vals_t) { + T local_data[mem_vals_t]; + mem_access::load_global(local_data, activations + read_offset + i); + mem_access::store_global(retained_tokens + write_offset + i, + local_data); + } +} + +template +void launch_gather_tokens(T* retained_tokens, + T* activations, + int32_t* gather_indices, + int32_t batch_size, + int32_t sampled_tokens, + int32_t channels, + int32_t read_batch_stride, + int32_t read_seq_stride, + int32_t write_batch_stride, + int32_t write_seq_stride, + cudaStream_t stream) +{ + constexpr int mem_vals_t = td_data::granularity / sizeof(T); + + const int load_steps = (channels + mem_vals_t - 1) / mem_vals_t; + const int threads = (load_steps >= 1024) ? 1024 : load_steps; + + dim3 block(threads); + dim3 grid(batch_size, sampled_tokens); + + gather_tokens_impl<<>>(retained_tokens, + activations, + gather_indices, + sampled_tokens, + channels, + read_batch_stride, + read_seq_stride, + write_batch_stride, + write_seq_stride); +} + +template void launch_gather_tokens(float*, + float*, + int32_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + cudaStream_t); + +template void launch_gather_tokens<__half>(__half*, + __half*, + int32_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + cudaStream_t); + +template +__global__ void scatter_tokens_impl(T* all_activations, + const T* layer_activations, + int32_t* gather_indices, + int32_t retained_tokens, + int32_t channels, + int32_t read_batch_stride, + int32_t read_seq_stride, + int32_t write_batch_stride, + int32_t write_seq_stride) +{ + constexpr int mem_vals_t = td_data::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + + const int gather_idx = + gather_indices[tb.group_index().x * retained_tokens + tb.group_index().y]; + + const int read_offset = + read_batch_stride * tb.group_index().x + read_seq_stride * tb.group_index().y; + const int write_offset = + write_batch_stride * tb.group_index().x + write_seq_stride * gather_idx; + + for (int i = tb.thread_index().x * mem_vals_t; i < channels; i += mem_vals_t * blockDim.x) { + T local_data[mem_vals_t]; + mem_access::load_global(local_data, + layer_activations + read_offset + i); + mem_access::store_global(all_activations + write_offset + i, + local_data); + } +} + +template +void launch_scatter_tokens(T* all_activations, + T* layer_activations, + int32_t* gather_indices, + int32_t batch_size, + int32_t sampled_tokens, + int32_t channels, + int32_t read_batch_stride, + int32_t read_seq_stride, + int32_t write_batch_stride, + int32_t write_seq_stride, + cudaStream_t stream) +{ + constexpr int mem_vals_t = td_data::granularity / sizeof(T); + + const int load_steps = (channels + mem_vals_t - 1) / mem_vals_t; + const int threads = (load_steps >= 1024) ? 1024 : load_steps; + + dim3 block(threads); + dim3 grid(batch_size, sampled_tokens); + + scatter_tokens_impl<<>>(all_activations, + layer_activations, + gather_indices, + sampled_tokens, + channels, + read_batch_stride, + read_seq_stride, + write_batch_stride, + write_seq_stride); +} + +template void launch_scatter_tokens(float*, + float*, + int32_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + cudaStream_t); + +template void launch_scatter_tokens<__half>(__half*, + __half*, + int32_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + int32_t, + cudaStream_t); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/pt_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..b50476719872e80fe785cac2c17501b24e1a01e4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/pt_binding.cpp @@ -0,0 +1,216 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include "custom_cuda_layers.h" + +torch::Tensor token_sort_(torch::Tensor& unsorted_token_ids, int64_t original_tokens) +{ + const int layers = unsorted_token_ids.size(0); + const int batch_size = unsorted_token_ids.size(1); + const int reserved_tokens = unsorted_token_ids.size(2); + + launch_token_sort(unsorted_token_ids.data_ptr(), + layers, + batch_size, + reserved_tokens, + original_tokens, + c10::cuda::getCurrentCUDAStream()); + + return unsorted_token_ids; +} + +torch::Tensor token_gather(torch::Tensor& activations, + torch::Tensor& sorted_indices, + bool batch_first) +{ + // Activations may be in either [N, S, C] or [S, N, C] while sorted_indices is + // always in [N, retained] + /* + TORCH_CHECK(sorted_indices.size(0) == activations.size(0) || + sorted_indices.size(0) == activations.size(1), + "Unable to match the batch size of the sorted indices to the activation + shape."); TORCH_CHECK(activations.size(2) % 8 == 0, "Channels must be divisible by 8 to align + with vectorized loads."); + */ + // bool batch_first = sorted_indices.size(0) == activations.size(0); + + const int64_t dim_0 = (batch_first) ? sorted_indices.size(0) : sorted_indices.size(1); + const int64_t dim_1 = (batch_first) ? sorted_indices.size(1) : sorted_indices.size(0); + const int64_t dim_2 = activations.size(2); + + auto output = torch::empty({dim_0, dim_1, dim_2}, activations.options()); + + const int batch_size = sorted_indices.size(0); + const int channels = dim_2; + const int retained_tokens = sorted_indices.size(1); + const int read_batch_stride = (batch_first) ? activations.stride(0) : activations.stride(1); + const int read_seq_stride = (batch_first) ? activations.stride(1) : activations.stride(0); + const int write_batch_stride = (batch_first) ? output.stride(0) : output.stride(1); + const int write_seq_stride = (batch_first) ? output.stride(1) : output.stride(0); + + if (activations.options().dtype() == torch::kFloat) { + launch_gather_tokens((float*)output.data_ptr(), + (float*)activations.data_ptr(), + (int32_t*)sorted_indices.data_ptr(), + batch_size, + retained_tokens, + channels, + read_batch_stride, + read_seq_stride, + write_batch_stride, + write_seq_stride, + c10::cuda::getCurrentCUDAStream()); + } else { + launch_gather_tokens((__half*)output.data_ptr(), + (__half*)activations.data_ptr(), + (int32_t*)sorted_indices.data_ptr(), + batch_size, + retained_tokens, + channels, + read_batch_stride, + read_seq_stride, + write_batch_stride, + write_seq_stride, + c10::cuda::getCurrentCUDAStream()); + } + + return output; +} + +torch::Tensor token_scatter_(torch::Tensor& all_activations, + torch::Tensor& layer_activations, + torch::Tensor& sorted_indices, + bool batch_first) +{ + // Activations may be in either [N, S, C] or [S, N, C] while sorted_indices is + // always in [N, retained] + /* + TORCH_CHECK(sorted_indices.size(0) == all_activations.size(0) || + sorted_indices.size(0) == all_activations.size(1), + "Unable to match the batch size of the sorted indices to the activation + shape."); TORCH_CHECK(all_activations.size(2) % 8 != 0, "Channels must be divisible by 8 to + align with vectorized loads."); + */ + // bool batch_first = sorted_indices.size(0) == all_activations.size(0); + + const int batch_size = sorted_indices.size(0); + const int channels = all_activations.size(2); + const int retained_tokens = sorted_indices.size(1); + const int read_batch_stride = (batch_first) ? layer_activations.stride(0) + : layer_activations.stride(1); + const int read_seq_stride = (batch_first) ? layer_activations.stride(1) + : layer_activations.stride(0); + const int write_batch_stride = (batch_first) ? all_activations.stride(0) + : all_activations.stride(1); + const int write_seq_stride = (batch_first) ? all_activations.stride(1) + : all_activations.stride(0); + + if (all_activations.options().dtype() == torch::kFloat) { + launch_scatter_tokens((float*)all_activations.data_ptr(), + (float*)layer_activations.data_ptr(), + (int32_t*)sorted_indices.data_ptr(), + batch_size, + retained_tokens, + channels, + read_batch_stride, + read_seq_stride, + write_batch_stride, + write_seq_stride, + c10::cuda::getCurrentCUDAStream()); + } else { + launch_scatter_tokens((__half*)all_activations.data_ptr(), + (__half*)layer_activations.data_ptr(), + (int32_t*)sorted_indices.data_ptr(), + batch_size, + retained_tokens, + channels, + read_batch_stride, + read_seq_stride, + write_batch_stride, + write_seq_stride, + c10::cuda::getCurrentCUDAStream()); + } + + return all_activations; +} + +torch::Tensor mask_gather_bert(torch::Tensor& dense_mask, torch::Tensor& sorted_indices) +{ + // TORCH_CHECK(dense_mask.dim() == 4) + + const int batch_size = dense_mask.size(0); + const int layers = sorted_indices.size(0); + /* + TORCH_CHECK(layers * batch_size == sorted_indices.size(0), + "Mismatch between the indices and the mask"); + */ + const int orig_seq_len = dense_mask.size(3); + const int truncated_seq_len = sorted_indices.size(2); + + auto output = torch::empty({layers, batch_size, 1, truncated_seq_len, truncated_seq_len}, + dense_mask.options()); + + if (dense_mask.options().dtype() == torch::kFloat) { + launch_slice_bert_mask((float*)output.data_ptr(), + (const float*)dense_mask.data_ptr(), + (const int32_t*)sorted_indices.data_ptr(), + layers, + batch_size, + truncated_seq_len, + orig_seq_len, + c10::cuda::getCurrentCUDAStream()); + } else { + launch_slice_bert_mask((__half*)output.data_ptr(), + (const __half*)dense_mask.data_ptr(), + (const int32_t*)sorted_indices.data_ptr(), + layers, + batch_size, + truncated_seq_len, + orig_seq_len, + c10::cuda::getCurrentCUDAStream()); + } + + return output; +} + +torch::Tensor mask_gather_gpt(torch::Tensor dense_mask, int truncated_seq_len) +{ + // TORCH_CHECK(dense_mask.dim() == 4) + + const int batch_size = dense_mask.size(0); + const int orig_seq_len = dense_mask.size(3); + + auto output = + torch::empty({batch_size, 1, truncated_seq_len, truncated_seq_len}, dense_mask.options()); + + if (dense_mask.options().dtype() == torch::kFloat) { + launch_slice_gpt_mask((float*)output.data_ptr(), + (const float*)dense_mask.data_ptr(), + batch_size, + truncated_seq_len, + orig_seq_len, + c10::cuda::getCurrentCUDAStream()); + } else { + launch_slice_gpt_mask((__half*)output.data_ptr(), + (const __half*)dense_mask.data_ptr(), + batch_size, + truncated_seq_len, + orig_seq_len, + c10::cuda::getCurrentCUDAStream()); + } + + return output; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("token_sort_", &token_sort_, "Comparison free sorting algorithm (CUDA)"); + m.def("token_gather", &token_gather, "Parallel gather of tokens (CUDA)"); + m.def("token_scatter_", &token_scatter_, "Parallel scatter of tokens (CUDA)"); + m.def("mask_gather_bert", &mask_gather_bert, "Token-based mask gather for BERT masking (CUDA)"); + m.def("mask_gather_gpt", &mask_gather_gpt, "Token-based mask gather for GPT masking (CUDA)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu new file mode 100644 index 0000000000000000000000000000000000000000..bc3823b846550cef1049e461ea35c35b53522146 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "custom_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +template +__global__ void slice_gpt_mask_impl(T* output_mask, + const T* input_mask, + int truncated_seq_len, + int orig_seq_len) +{ + const int in_batch_stride = orig_seq_len * orig_seq_len; + const int out_batch_stride = truncated_seq_len * truncated_seq_len; + + cg::thread_block tb = cg::this_thread_block(); + + const T* input_mask_block = + input_mask + blockIdx.x * in_batch_stride + blockIdx.y * orig_seq_len; + T* output_mask_block = + output_mask + blockIdx.x * out_batch_stride + blockIdx.y * truncated_seq_len; + + for (int i = tb.thread_index().x; i < truncated_seq_len; i += blockDim.x) { + output_mask_block[i] = input_mask_block[i]; + } +} + +template +void launch_slice_gpt_mask(T* output_mask, + const T* input_mask, + int batch_size, + int truncated_seq_len, + int orig_seq_len, + cudaStream_t stream) +{ + const int threads = (truncated_seq_len >= 1024) ? 1024 : truncated_seq_len; + + dim3 block(threads); + dim3 grid(batch_size, truncated_seq_len); + + slice_gpt_mask_impl + <<>>(output_mask, input_mask, truncated_seq_len, orig_seq_len); +} + +template void launch_slice_gpt_mask(float*, const float*, int, int, int, cudaStream_t); + +template void launch_slice_gpt_mask<__half>(__half*, const __half*, int, int, int, cudaStream_t); + +template +__global__ void slice_bert_mask_impl(T* output_mask, + const T* input_mask, + const int32_t* retained_indices, + int32_t truncated_seq_len, + int32_t orig_seq_len) +{ + const int in_batch_stride = orig_seq_len * orig_seq_len; + const int out_batch_stride = truncated_seq_len * truncated_seq_len; + const int out_layer_stride = out_batch_stride * gridDim.y; + + cg::thread_block tb = cg::this_thread_block(); + + const int out_layer_offset = tb.group_index().x * out_layer_stride; + + const int in_batch_offset = tb.group_index().y * in_batch_stride; + const int out_batch_offset = tb.group_index().y * out_batch_stride; + + const int32_t gather_row = + retained_indices[tb.group_index().y * truncated_seq_len + tb.group_index().z]; + const int in_seq_offset = gather_row * orig_seq_len; + const int out_seq_offset = tb.group_index().z * truncated_seq_len; + + const T* in_sequence = input_mask + in_batch_offset + in_seq_offset; + T* out_sequence = output_mask + out_layer_offset + out_batch_offset + out_seq_offset; + const int32_t* gather_data = retained_indices + tb.group_index().y * truncated_seq_len; + + for (int i = tb.thread_index().x; i < truncated_seq_len; i += blockDim.x) { + out_sequence[i] = in_sequence[gather_data[i]]; + } +} + +/* +Since the Bert mask is not causal like GPT, we can't just generate a set of +masks for the entire model based off a single layer sample. + +We map the kernel as follows: +z-dimension: layer +y-dimension: batch +x-dimension: sequence_offset +*/ +template +void launch_slice_bert_mask(T* output_mask, + const T* input_mask, + const int32_t* retained_indices, + int32_t layers, + int32_t batch_size, + int32_t truncated_seq_len, + int32_t orig_seq_len, + cudaStream_t stream) +{ + const int threads = (truncated_seq_len >= 1024) ? 1024 : truncated_seq_len; + dim3 block(threads); + dim3 grid(layers, batch_size, truncated_seq_len); + + slice_bert_mask_impl<<>>( + output_mask, input_mask, retained_indices, truncated_seq_len, orig_seq_len); +} + +template void launch_slice_bert_mask(float*, + const float*, + const int32_t*, + int32_t, + int32_t, + int32_t, + int32_t, + cudaStream_t); + +template void launch_slice_bert_mask<__half>(__half*, + const __half*, + const int32_t*, + int32_t, + int32_t, + int32_t, + int32_t, + cudaStream_t); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/token_sort.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/token_sort.cu new file mode 100644 index 0000000000000000000000000000000000000000..3c1dff49429fc2e0f436ddc9d89a48af490da9c8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/token_sort.cu @@ -0,0 +1,194 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "custom_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +namespace td_sort { +constexpr int threads = 512; +constexpr int granularity = 16; +constexpr int mem_vals = granularity / sizeof(int32_t); +constexpr int max_buffer_size = (threads + 1) * mem_vals; + +#ifdef __HIP_PLATFORM_AMD__ +constexpr int warp_size = ROCM_WAVEFRONT_SIZE; +#else +constexpr int warp_size = 32; +#endif + +constexpr int max_warps = threads / warp_size; +} // namespace td_sort + +template +__global__ void scan_sort(int32_t* data, int reserved_tokens, int original_tokens) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + __shared__ int32_t indices_buffer[td_sort::max_buffer_size]; + __shared__ int32_t intermediate_buffer[td_sort::max_warps]; + __shared__ int32_t sorted_indices_buffer[td_sort::max_buffer_size]; + + for (int i = tb.thread_index().x * td_sort::mem_vals; i < original_tokens + 1; + i += tb.group_dim().x * td_sort::mem_vals) { + uint32_t zeros[td_sort::mem_vals] = {0, 0, 0, 0}; + mem_access::store_shared(indices_buffer + i, zeros); + } + + int32_t local_vals[VALS_PER_THREAD]; + + // We flatten layers/batch into a single indexing dimension + int32_t* data_block = data + tb.group_index().x * reserved_tokens; + + // The next two loops really could be fused for a more logical code layout, but don't want to + // move the barrier forward +#pragma unroll + for (int i = 0; i < VALS_PER_THREAD; i++) { + const int iter_idx = i * td_sort::threads + tb.thread_index().x; + if (iter_idx < reserved_tokens) { + mem_access::load_global(local_vals + i, data_block + iter_idx); + } else { + local_vals[i] = 0; + } + } + + tb.sync(); + +#pragma unroll + for (int i = 0; i < VALS_PER_THREAD; i++) { + const int iter_idx = i * td_sort::threads + tb.thread_index().x; + if (iter_idx < reserved_tokens) { + const int32_t one = 1; + mem_access::store_shared(indices_buffer + local_vals[i], &one); + } + } + + tb.sync(); + + int32_t local_input[td_sort::mem_vals]; + mem_access::load_shared( + local_input, indices_buffer + tb.thread_index().x * td_sort::mem_vals); + + int32_t reduce_vals[td_sort::mem_vals]; + reduce_vals[0] = local_input[0]; + +#pragma unroll + for (int i = 1; i < td_sort::mem_vals; i++) { + reduce_vals[i] = local_input[i] + reduce_vals[i - 1]; + } + + int32_t step_1_val = reduce_vals[td_sort::mem_vals - 1]; + // Short span exclusive scan algorithm (less work efficient) +#pragma unroll + for (int i = 1; i < td_sort::warp_size; i *= 2) { + int32_t step_val = warp.shfl_up(step_1_val, i); + step_1_val = (warp.thread_rank() < i) ? step_1_val : step_1_val + step_val; + } + + if (warp.thread_rank() == td_sort::warp_size - 1) { + mem_access::store_shared(intermediate_buffer + warp.meta_group_rank(), + &step_1_val); + } + + tb.sync(); + + if (warp.meta_group_rank() == 0) { + int32_t step_2_val = 0; + if (warp.thread_rank() < td_sort::max_warps) { + mem_access::load_shared(&step_2_val, + intermediate_buffer + warp.thread_rank()); + } + +#pragma unroll + for (int i = 1; i < td_sort::warp_size; i *= 2) { + int32_t step_val = warp.shfl_up(step_2_val, i); + step_2_val = (warp.thread_rank() < i) ? step_2_val : step_2_val + step_val; + } + + if (warp.thread_rank() < td_sort::max_warps) { + mem_access::store_shared(intermediate_buffer + warp.thread_rank(), + &step_2_val); + } + } + + tb.sync(); + + int step_2_val = 0; + if (warp.meta_group_rank() > 0) { + mem_access::load_shared(&step_2_val, + intermediate_buffer + warp.meta_group_rank() - 1); + } + + const int thread_offset = reduce_vals[td_sort::mem_vals - 1]; + +#pragma unroll + for (int i = 0; i < td_sort::mem_vals; i++) { + reduce_vals[i] += step_1_val + step_2_val - thread_offset; + } + mem_access::store_shared( + indices_buffer + tb.thread_index().x * td_sort::mem_vals, reduce_vals); + + if (tb.thread_index().x == 0) { + indices_buffer[original_tokens] = original_tokens - indices_buffer[original_tokens]; + } + tb.sync(); + + for (int i = 0; i < VALS_PER_THREAD; i++) { + const int iter_idx = i * td_sort::threads + tb.thread_index().x; + if (iter_idx < reserved_tokens) { + if (local_vals[i] == 0) { + int zero = 0; + mem_access::store_shared(sorted_indices_buffer, &zero); + } else { + int sorted_idx; + mem_access::load_shared(&sorted_idx, + indices_buffer + local_vals[i] - 1); + mem_access::store_shared(sorted_indices_buffer + sorted_idx, + local_vals + i); + } + } + } + + tb.sync(); + +#pragma unroll + for (int i = 0; i < VALS_PER_THREAD; i++) { + const int iter_idx = i * td_sort::threads + tb.thread_index().x; + if (iter_idx < reserved_tokens) { + int32_t store_val; + mem_access::load_shared(&store_val, sorted_indices_buffer + iter_idx); + mem_access::store_global(data_block + iter_idx, &store_val); + } + } +} + +void launch_token_sort(int32_t* indices, + int layers, + int batch_size, + int reserved_size, + int original_tokens, + cudaStream_t stream) +{ + // Each sort is completely independent, can flatten this dimension + dim3 grid(layers * batch_size); + dim3 block(td_sort::threads); + + const int vals_per_thread = (reserved_size + td_sort::threads - 1) / td_sort::threads; + + if (vals_per_thread == 1) { + scan_sort<1><<>>(indices, reserved_size, original_tokens); + } else if (vals_per_thread == 2) { + scan_sort<2><<>>(indices, reserved_size, original_tokens); + } else if (vals_per_thread == 3) { + scan_sort<3><<>>(indices, reserved_size, original_tokens); + } else if (vals_per_thread == 4) { + scan_sort<4><<>>(indices, reserved_size, original_tokens); + } else { + assert(false); + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp new file mode 100644 index 0000000000000000000000000000000000000000..352306ba26128b96236f43764c26d2f4191eb391 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp @@ -0,0 +1,127 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/csrc/utils.cpp +*/ + +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif + +typedef std::vector> ret_t; + +void segment_blocks(torch::Tensor layout, + torch::Tensor idx, + torch::Tensor scratch, + int max_width, + ret_t& ret) +{ + size_t H = layout.size(0); + size_t M = layout.size(1); + size_t N = layout.size(2); + torch::Tensor tmp = torch::zeros_like(layout); + + auto _tmp = tmp.accessor(); + auto _layout = layout.accessor(); + auto _idx = idx.accessor(); + auto _scratch = scratch.accessor(); + std::vector current(H, 0); + +#ifdef _OPENMP +#pragma omp parallel for +#endif + for (size_t h = 0; h < H; h++) { + // surrounding indices + std::vector ii_left(max_width, -1); + std::vector> ii_top(max_width, std::vector(N, -1)); + + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++) { + int v = _layout[h][m][n]; + if (v == 0) continue; + int n_left = ii_left[max_width - 1]; + int m_top = ii_top[max_width - 1][n]; + int top = (m_top >= 0) ? _tmp[h][m_top][n] : 0; + int left = (n_left >= 0) ? _tmp[h][m][n_left] : 0; + int topleft = (m_top >= 0 && n_left >= 0) ? _tmp[h][m_top][n_left] : 0; + int width = std::min(left, std::min(top, topleft)) + 1; + + // reset width if blocks cannot be + // packed together (i.e., there's a 1 "in the middle") + for (int nn = n_left + 1; nn < n; nn++) + if (ii_top[max_width - 1][nn] > ii_top[max_width - 1][n]) width = 1; + _tmp[h][m][n] = width; + + // update n_left ring buffer + for (int k = 0; k < max_width - 1; k++) ii_left[k] = ii_left[k + 1]; + ii_left[max_width - 1] = n; + + // update ii_top ring buffer + for (int k = 0; k < max_width - 1; k++) ii_top[k][n] = ii_top[k + 1][n]; + ii_top[max_width - 1][n] = m; + + // block is too small -- skip + if (width != max_width) continue; + + // retained blocks are set to zeros + for (size_t km = 0; km < max_width; km++) + for (size_t kn = 0; kn < max_width; kn++) { + int mm = ii_top[km][n]; + int nn = ii_left[kn]; + if (mm < 0 || nn < 0) continue; + _layout[h][mm][nn] = 0; + _tmp[h][mm][nn] = 0; + _scratch[h][current[h]][0] = (int)h; + _scratch[h][current[h]][1] = (int)mm; + _scratch[h][current[h]][2] = (int)nn; + _scratch[h][current[h]][3] = _idx[h][mm][nn]; + current[h]++; + } + } + } + } + std::vector to_cat; + for (size_t h = 0; h < H; h++) + if (current[h] > 0) to_cat.push_back(scratch[h].slice(0, 0, current[h])); + if (!to_cat.empty()) ret.push_back({max_width, torch::cat(to_cat)}); +} + +ret_t sdd_segment(torch::Tensor layout, int start_width) +{ + ret_t ret; + + // block index + torch::Tensor idx = torch::zeros_like(layout); + int current = 0; + int64_t H = layout.size(0); + int64_t M = layout.size(1); + int64_t N = layout.size(2); + auto _layout = layout.accessor(); + auto _idx = idx.accessor(); + for (int64_t h = 0; h < H; h++) + for (int64_t m = 0; m < M; m++) + for (int64_t n = 0; n < N; n++) { + if (_layout[h][m][n] == 0) continue; + _idx[h][m][n] = current++; + } + + // scratch memory + torch::Tensor scratch = torch::empty({H, layout.sum().item(), 4}, layout.dtype()); + + for (int max_width = start_width; max_width > 0; max_width /= 2) + segment_blocks(layout, idx, scratch, max_width, ret); + return ret; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("sdd_segment", &sdd_segment, "SDD segmentation handler"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu new file mode 100644 index 0000000000000000000000000000000000000000..d831b372b65f398b43c5cd343b2bf2db67562f56 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu @@ -0,0 +1,149 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "memory_access_utils.h" +#include "spatial_cuda_layers.h" + +/* +Fused bias add variants +*/ + +namespace badd_opt { +constexpr int threads = 256; +constexpr int steps = 2; +constexpr int granularity = 16; +constexpr int vals_per_h = granularity / sizeof(__half); +constexpr int vals_per_h2 = granularity / sizeof(__half2); +constexpr int vals_per_block = threads * steps * vals_per_h; +constexpr int stride = vals_per_h * threads; +} // namespace badd_opt + +__global__ void opt_bias_add(__half* result, + const __half* activation, + const __half* bias, + int seq_len, + int channels) +{ + const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; + const int stride = badd_opt::vals_per_h * badd_opt::threads; + + for (int i = 0; i < badd_opt::steps; i++) { + if (id + i * badd_opt::stride < seq_len * channels) { + __half2 act_buffer[badd_opt::vals_per_h2]; + __half2 bias_buffer[badd_opt::vals_per_h2]; + + mem_access::load_global(act_buffer, + activation + id + i * stride); + mem_access::load_global(bias_buffer, + bias + ((id + i * stride) % channels)); + + for (int j = 0; j < badd_opt::vals_per_h2; j++) { act_buffer[j] += bias_buffer[j]; } + + mem_access::store_global(result + id + i * stride, act_buffer); + } + } +} + +__global__ void opt_bias_add_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + int seq_len, + int channels) +{ + const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; + const int stride = badd_opt::vals_per_h * badd_opt::threads; + + for (int i = 0; i < badd_opt::steps; i++) { + if (id + i * badd_opt::stride < seq_len * channels) { + __half2 act_buffer[badd_opt::vals_per_h2]; + __half2 bias_buffer[badd_opt::vals_per_h2]; + __half2 other_buffer[badd_opt::vals_per_h2]; + + mem_access::load_global(act_buffer, + activation + id + i * stride); + mem_access::load_global(bias_buffer, + bias + ((id + i * stride) % channels)); + mem_access::load_global(other_buffer, other + id + i * stride); + + for (int j = 0; j < badd_opt::vals_per_h2; j++) { + act_buffer[j] += bias_buffer[j] + other_buffer[j]; + } + + mem_access::store_global(result + id + i * stride, act_buffer); + } + } +} + +__global__ void opt_bias_add_bias_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + const __half* other_bias, + int seq_len, + int channels) +{ + const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; + const int stride = badd_opt::vals_per_h * badd_opt::threads; + + for (int i = 0; i < badd_opt::steps; i++) { + if (id + i * badd_opt::stride < seq_len * channels) { + __half2 act_buffer[badd_opt::vals_per_h2]; + __half2 bias_buffer[badd_opt::vals_per_h2]; + __half2 other_buffer[badd_opt::vals_per_h2]; + __half2 other_bias_buffer[badd_opt::vals_per_h2]; + + mem_access::load_global(act_buffer, + activation + id + i * stride); + mem_access::load_global(bias_buffer, + bias + ((id + i * stride) % channels)); + mem_access::load_global(other_buffer, other + id + i * stride); + mem_access::load_global( + other_bias_buffer, other_bias + ((id + i * stride) % channels)); + + for (int j = 0; j < badd_opt::vals_per_h2; j++) { + act_buffer[j] = + (act_buffer[j] + bias_buffer[j]) + (other_buffer[j] + other_bias_buffer[j]); + } + + mem_access::store_global(result + id + i * stride, act_buffer); + } + } +} + +void launch_opt_bias_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + const __half* other_bias, + int batch_size, + int seq_len, + int channels, + cudaStream_t stream) +{ + // Should evaluate `true` for reasonable hidden sizes + assert(channels % badd_opt::vals_per_h == 0); + + const int effective_seq_len = batch_size * seq_len; + const int vals = effective_seq_len * channels; + + dim3 block(badd_opt::threads); + dim3 grid((vals + badd_opt::vals_per_block - 1) / badd_opt::vals_per_block); + + if (!other) { + // We shouldn't have a bias if there's no activation + assert(!other_bias); + + opt_bias_add<<>>( + result, activation, bias, effective_seq_len, channels); + } else if (!other_bias) { + opt_bias_add_add<<>>( + result, activation, bias, other, effective_seq_len, channels); + } else { + opt_bias_add_bias_add<<>>( + result, activation, bias, other, other_bias, effective_seq_len, channels); + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..cbf6636a6ee4ac5769a88b7cab7d2a222f00d5cc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp @@ -0,0 +1,112 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "spatial_cuda_layers.h" + +ChannelsLastProblem dimension_problem(at::Tensor& input) +{ + ChannelsLastProblem dims; + + if (input.dim() == 4) { + // In some sense this is unsafe (and a reflection of the assumptions made inside + // the C10 options checker). Basically, there's no great way to be sure that + // a tensor is in channels last because a 1x1 image will appear to be in channels + // last even when it isn't. + assert(input.is_contiguous(at::MemoryFormat::ChannelsLast)); + dims.batch_size = input.size(0); + dims.seq_len = input.size(2) * input.size(3); + dims.channels = input.size(1); + } else { + assert(input.is_contiguous()); + dims.batch_size = input.size(0); + dims.seq_len = input.size(1); + dims.channels = input.size(2); + } + + return dims; +} + +at::Tensor seq_unroll_bias_add(at::Tensor& input, at::Tensor& bias) +{ + assert(input.dtype() == at::kHalf); + + // TODO(cmikeh2): Should probably refactor this into a more portable + // description, since it does generalize for channels-last + ChannelsLastProblem problem = dimension_problem(input); + + auto output = at::empty_like(input); + + launch_opt_bias_add((__half*)output.data_ptr(), + (const __half*)input.data_ptr(), + (const __half*)bias.data_ptr(), + nullptr, + nullptr, + problem.batch_size, + problem.seq_len, + problem.channels, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor seq_bias_add_add(at::Tensor& input, at::Tensor& bias, at::Tensor& other) +{ + assert(input.dtype() == at::kHalf); + + // TODO(cmikeh2): Should probably refactor this into a more portable + // description, since it does generalize for channels-last + ChannelsLastProblem problem = dimension_problem(input); + + auto output = at::empty_like(input); + + launch_opt_bias_add((__half*)output.data_ptr(), + (const __half*)input.data_ptr(), + (const __half*)bias.data_ptr(), + (const __half*)other.data_ptr(), + nullptr, + problem.batch_size, + problem.seq_len, + problem.channels, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor seq_bias_add_bias_add(at::Tensor& input, + at::Tensor& bias, + at::Tensor& other, + at::Tensor& other_bias) +{ + assert(input.dtype() == at::kHalf); + + // TODO(cmikeh2): Should probably refactor this into a more portable + // description, since it does generalize for channels-last + ChannelsLastProblem problem = dimension_problem(input); + + auto output = at::empty_like(input); + + launch_opt_bias_add((__half*)output.data_ptr(), + (const __half*)input.data_ptr(), + (const __half*)bias.data_ptr(), + (const __half*)other.data_ptr(), + (const __half*)other_bias.data_ptr(), + problem.batch_size, + problem.seq_len, + problem.channels, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("nhwc_bias_add", &seq_unroll_bias_add); + m.def("nhwc_bias_add_add", &seq_bias_add_add); + m.def("nhwc_bias_add_bias_add", &seq_bias_add_bias_add); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h new file mode 100644 index 0000000000000000000000000000000000000000..4f56f89f16cd1a200f057ebd651575224be23b7b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#if __CUDA_ARCH__ >= 530 +#define HALF_PRECISION_AVAILABLE = 1 +#endif + +#ifdef __HIP_PLATFORM_AMD__ +#include +#else +#include +#endif + +#include +#include + +/*********** Group Norm Kernels, Structs, and Helpers ************/ + +struct { + int64_t batch_size; + int64_t seq_len; + int64_t channels; +} typedef ChannelsLastProblem; + +void launch_opt_bias_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + const __half* other_bias, + int batch_size, + int seq_len, + int channels, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/cublas_wrappers.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/cublas_wrappers.cu new file mode 100644 index 0000000000000000000000000000000000000000..d982e65b8a81f786b4df550da310725d32ccc826 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/cublas_wrappers.cu @@ -0,0 +1,493 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cublas_wrappers.h" + +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_gemm_ex(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + rocblas_gemm_algo algo) +#else +int cublas_gemm_ex(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + cublasGemmAlgo_t algo) +#endif +{ +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_status status = rocblas_gemm_ex(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, + rocblas_datatype_f32_r, + (transa == rocblas_operation_none) ? m : k, + (const void*)B, + rocblas_datatype_f32_r, + (transb == rocblas_operation_none) ? k : n, + (const void*)beta, + C, + rocblas_datatype_f32_r, + m, + C, + rocblas_datatype_f32_r, + m, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else + cublasStatus_t status = cublasGemmEx(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (transa == CUBLAS_OP_N) ? m : k, + (const void*)B, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (transb == CUBLAS_OP_N) ? k : n, + (const void*)beta, + C, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + m, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + return 0; +} + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_gemm_ex(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const __half* A, + const __half* B, + __half* C, + rocblas_gemm_algo algo) +#else +int cublas_gemm_ex(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const __half* A, + const __half* B, + __half* C, + cublasGemmAlgo_t algo) +#endif +{ +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_status status = rocblas_gemm_ex(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, + rocblas_datatype_f16_r, + (transa == rocblas_operation_none) ? m : k, + (const void*)B, + rocblas_datatype_f16_r, + (transb == rocblas_operation_none) ? k : n, + (const void*)beta, + (void*)C, + rocblas_datatype_f16_r, + m, + (void*)C, + rocblas_datatype_f16_r, + m, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else + cublasStatus_t status = cublasGemmEx(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_16F, +#else + CUDA_R_16F, +#endif + (transa == CUBLAS_OP_N) ? m : k, + (const void*)B, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_16F, +#else + CUDA_R_16F, +#endif + (transb == CUBLAS_OP_N) ? k : n, + (const void*)beta, + (void*)C, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_16F, +#else + CUDA_R_16F, +#endif + m, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + return 0; +} + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_strided_batched_gemm(rocblas_handle handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + rocblas_operation op_A, + rocblas_operation op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + rocblas_gemm_algo algo) +#else +int cublas_strided_batched_gemm(cublasHandle_t handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + cublasOperation_t op_A, + cublasOperation_t op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + cublasGemmAlgo_t algo) +#endif +{ +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_status status = + rocblas_gemm_strided_batched_ex(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, + rocblas_datatype_f32_r, + (op_A == rocblas_operation_none) ? m : k, + stride_A, + B, + rocblas_datatype_f32_r, + (op_B == rocblas_operation_none) ? k : n, + stride_B, + beta, + C, + rocblas_datatype_f32_r, + m, + stride_C, + C, + rocblas_datatype_f32_r, + m, + stride_C, + batch, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else + cublasStatus_t status = cublasGemmStridedBatchedEx(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (op_A == CUBLAS_OP_N) ? m : k, + stride_A, + B, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (op_B == CUBLAS_OP_N) ? k : n, + stride_B, + beta, + C, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + m, + stride_C, + batch, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (batch: %d, m: %d, n: %d, k: %d, error: %d) \n", + batch, + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + return 0; +} + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_strided_batched_gemm(rocblas_handle handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const __half* A, + const __half* B, + __half* C, + rocblas_operation op_A, + rocblas_operation op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + rocblas_gemm_algo algo) +#else +int cublas_strided_batched_gemm(cublasHandle_t handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const __half* A, + const __half* B, + __half* C, + cublasOperation_t op_A, + cublasOperation_t op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + cublasGemmAlgo_t algo) +#endif +{ +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_status status = + rocblas_gemm_strided_batched_ex(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + (op_A == rocblas_operation_none) ? m : k, + stride_A, + B, + rocblas_datatype_f16_r, + (op_B == rocblas_operation_none) ? k : n, + stride_B, + beta, + C, + rocblas_datatype_f16_r, + m, + stride_C, + C, + rocblas_datatype_f16_r, + m, + stride_C, + batch, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else + cublasStatus_t status = cublasGemmStridedBatchedEx(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_16F, +#else + CUDA_R_16F, +#endif + (op_A == CUBLAS_OP_N) ? m : k, + stride_A, + B, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_16F, +#else + CUDA_R_16F, +#endif + (op_B == CUBLAS_OP_N) ? k : n, + stride_B, + beta, + C, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_16F, +#else + CUDA_R_16F, +#endif + m, + stride_C, + batch, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + + return 0; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/dropout_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/dropout_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..97cb3592dcdc4cf0a2dc37e34fc9b13117d9bcc5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/dropout_kernels.cu @@ -0,0 +1,873 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "custom_cuda_layers.h" + +const int unroll_factor = 4; + +__global__ void dropout_kernel(const int N, + const float ratio, + float* out, + const float* Xdata, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + float4 rand = curand_uniform4(&state); + uint8_t m[unroll_factor]; + + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + int i = j * unroll_factor; + + mask[i] = (uint8_t)m[0]; + mask[i + 1] = (uint8_t)m[1]; + mask[i + 2] = (uint8_t)m[2]; + mask[i + 3] = (uint8_t)m[3]; + + out[i] = Xdata[i] * scale * m[0]; + out[i + 1] = Xdata[i + 1] * scale * m[1]; + out[i + 2] = Xdata[i + 2] * scale * m[2]; + out[i + 3] = Xdata[i + 3] * scale * m[3]; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + float4 rand = curand_uniform4(&state); + float* rand_data = &(rand.x); + int k = 0; + for (int i = high_index; i < N; i++) { + uint8_t m = (uint8_t)(rand_data[k++] > ratio); + out[i] = Xdata[i] * scale * m; + mask[i] = m; + } + } +} + +__global__ void dropout_kernel(const int N, + const float ratio, + __half* out, + const __half* Xdata, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + +#ifdef __STOCHASTIC_MODE__ + + const __half2 h_scale = __float2half2_rn(scale); + const float2* x_cast = reinterpret_cast(Xdata); + float2* out_cast = reinterpret_cast(out); + uint32_t* mask_cast = reinterpret_cast(mask); + + uint32_t m_32; + uint8_t* m = reinterpret_cast(&m_32); + + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + __half2 mask_h[2]; + float2 mask_f[2]; + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + float2 x_f = x_cast[j]; + __half2* x_h = reinterpret_cast<__half2*>(&x_f); + + float4 rand = curand_uniform4(&state); + + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + float* mask_f_data = &mask_f[0].x; +#pragma unroll + for (int i = 0; i < unroll_factor; i++) mask_f_data[i] = (float)(m[i]); + + mask_h[0] = __float22half2_rn(mask_f[0]); + mask_h[1] = __float22half2_rn(mask_f[1]); + + result_h[0] = x_h[0] * h_scale * mask_h[0]; + result_h[1] = x_h[1] * h_scale * mask_h[1]; + + out_cast[j] = result_f; + + mask_cast[j] = m_32; + } + +#else + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + int i = j * unroll_factor; + + const __half2* vals_half = reinterpret_cast(Xdata + i); + float2 vals_half_f[2]; + vals_half_f[0] = __half22float2(vals_half[0]); + vals_half_f[1] = __half22float2(vals_half[1]); + + uint8_t m[unroll_factor]; + float4 rand = curand_uniform4(&state); + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + out[i] = __float2half(vals_half_f[0].x * scale * m[0]); + out[i + 1] = __float2half(vals_half_f[0].y * scale * m[1]); + out[i + 2] = __float2half(vals_half_f[1].x * scale * m[2]); + out[i + 3] = __float2half(vals_half_f[1].y * scale * m[3]); + + mask[i] = m[0]; + mask[i + 1] = m[1]; + mask[i + 2] = m[2]; + mask[i + 3] = m[3]; + } + +#endif + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + float4 rand = curand_uniform4(&state); + float* rand_data = &(rand.x); + int k = 0; + for (int i = high_index; i < N; i++) { + uint8_t m = (uint8_t)(rand_data[k++] > ratio); + out[i] = __float2half((float)Xdata[i] * scale * m); + mask[i] = m; + } + } +} + +__global__ void dropout_kernel_bwd(const int N, + const float ratio, + const float* Xdata, + float* out, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + int i = j * unroll_factor; + + out[i] = mask[i] ? Xdata[i] * scale : 0.0; + out[i + 1] = mask[i + 1] ? Xdata[i + 1] * scale : 0.0; + out[i + 2] = mask[i + 2] ? Xdata[i + 2] * scale : 0.0; + out[i + 3] = mask[i + 3] ? Xdata[i + 3] * scale : 0.0; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + for (int i = high_index; i < N; i++) { out[i] = mask[i] ? Xdata[i] * scale : 0.0; } + } +} + +__global__ void dropout_kernel_bwd(const int N, + const float ratio, + const __half* Xdata, + __half* out, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + +#ifdef __STOCHASTIC_MODE__ + + const __half2 h_scale = __float2half2_rn(scale); + + const float2* x_cast = reinterpret_cast(Xdata); + float2* out_cast = reinterpret_cast(out); + uint32_t* mask_cast = reinterpret_cast(mask); + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + float2 x_f = x_cast[j]; + __half2* x_h = reinterpret_cast<__half2*>(&x_f); + + uint32_t m_32 = mask_cast[j]; + uint8_t* m = (uint8_t*)&m_32; + + __half2 mask_h[2]; + float2 mask_f[2]; + + float* mask_f_data = &mask_f[0].x; +#pragma unroll + for (int i = 0; i < unroll_factor; i++) mask_f_data[i] = (float)(m[i]); + +#pragma unroll + for (int i = 0; i < 2; i++) mask_h[i] = __float22half2_rn(mask_f[i]); + + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + + result_h[0] = x_h[0] * h_scale * mask_h[0]; + result_h[1] = x_h[1] * h_scale * mask_h[1]; + + out_cast[j] = result_f; + } + +#else + + const __half h_scale = __float2half(scale); + const __half h_zero = __float2half(0.0); + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + int i = j * unroll_factor; + + const __half2* vals_half = reinterpret_cast(Xdata + i); + + uint8_t* m = mask + i; + + float2 vals_half_f[2]; + + vals_half_f[0] = __half22float2(vals_half[0]); + vals_half_f[1] = __half22float2(vals_half[1]); + + out[i] = __float2half(vals_half_f[0].x * scale * m[0]); + out[i + 1] = __float2half(vals_half_f[0].y * scale * m[1]); + out[i + 2] = __float2half(vals_half_f[1].x * scale * m[2]); + out[i + 3] = __float2half(vals_half_f[1].y * scale * m[3]); + } + +#endif + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + for (int i = high_index; i < N; i++) { + out[i] = __float2half((float)Xdata[i] * scale * mask[i]); + } + } +} + +template +void launch_dropout(T* out, + const T* vals, + uint8_t* mask, + int total_count, + int dim, + float ratio, + cudaStream_t stream, + bool bwd) +{ + assert(unroll_factor == 4); + + dim3 grid_dim = DS_GET_BLOCKS(total_count / unroll_factor); + dim3 block_dim = DS_CUDA_NUM_THREADS; + + if (dim > 512) { + block_dim.x >>= 1; + grid_dim.x <<= 1; + } + uint64_t inc = total_count / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + if (bwd) + dropout_kernel_bwd<<>>( + total_count, ratio, vals, out, mask, seed); + else + dropout_kernel<<>>( + total_count, ratio, out, vals, mask, seed); +} + +template void launch_dropout(float* out, + const float* vals, + uint8_t* mask, + int total_count, + int dim, + float ratio, + cudaStream_t stream, + bool); +template void launch_dropout(__half* out, + const __half* vals, + uint8_t* mask, + int total_count, + int dim, + float ratio, + cudaStream_t stream, + bool); + +__global__ void dropout_grad_kernel(const int N, const float scale, float* Xdata, uint8_t* mask) +{ + CUDA_1D_KERNEL_LOOP(i, N) { Xdata[i] *= scale * mask[i]; } +} + +__global__ void dropout_grad_kernel(const int N, const float scale, __half* Xdata, uint8_t* mask) +{ + const __half2 h_scale = __float2half2_rn(scale); + float2* x_cast = reinterpret_cast(Xdata); + uint32_t* mask_cast = reinterpret_cast(mask); + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + float2 x_data = x_cast[j]; + uint32_t m_32 = mask_cast[j]; + uint8_t* m = (uint8_t*)&m_32; + + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + +#ifdef __STOCHASTIC_MODE__ + + __half2* x_data_h = reinterpret_cast<__half2*>(&x_data); + __half2 mask_h[2]; + float2 mask_f[2]; + + float* mask_f_data = &mask_f[0].x; +#pragma unroll + for (int i = 0; i < unroll_factor; i++) *(mask_f_data++) = (float)(m[i]); + + mask_h[0] = __float22half2_rn(mask_f[0]); + mask_h[1] = __float22half2_rn(mask_f[1]); + + result_h[0] = x_data_h[0] * h_scale * mask_h[0]; + result_h[1] = x_data_h[1] * h_scale * mask_h[1]; + +#else + + __half* x_data_h = reinterpret_cast<__half*>(&x_data); + float2 result[2]; + + result[0].x = (float)x_data_h[0] * scale * m[0]; + result[0].y = (float)x_data_h[1] * scale * m[1]; + result[1].x = (float)x_data_h[2] * scale * m[2]; + result[1].y = (float)x_data_h[3] * scale * m[3]; + + result_h[0] = __float22half2_rn(result[0]); + result_h[1] = __float22half2_rn(result[1]); + +#endif + x_cast[j] = result_f; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + for (int i = high_index; i < N; i++) { + Xdata[i] = __float2half((float)Xdata[i] * scale * mask[i]); + } + } +} + +template +void launch_dropout_grad(T* vals, uint8_t* mask, int total_count, float ratio, cudaStream_t stream) +{ + assert(unroll_factor == 4); + + const float scale = 1. / (1. - ratio); + dropout_grad_kernel<<>>(total_count, scale, vals, mask); +} + +template void launch_dropout_grad(float* vals, + uint8_t* mask, + int total_count, + float ratio, + cudaStream_t stream); +template void launch_dropout_grad(__half* vals, + uint8_t* mask, + int total_count, + float ratio, + cudaStream_t stream); + +__global__ void dropout_grad_kernel(const int N, + const float scale, + const float* Xdata, + float* out, + uint8_t* mask) +{ + CUDA_1D_KERNEL_LOOP(i, N) { out[i] = Xdata[i] * scale * mask[i]; } +} + +__global__ void dropout_grad_kernel(const int N, + const float scale, + const __half* Xdata, + __half* out, + uint8_t* mask) +{ + const float2* x_cast = reinterpret_cast(Xdata); + float2* out_cast = reinterpret_cast(out); + const uint32_t* mask_cast = reinterpret_cast(mask); + + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + + CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) + { + float2 x_data = x_cast[j]; + uint32_t m_32 = mask_cast[j]; + uint8_t* m = (uint8_t*)&m_32; + + __half* x_data_h = reinterpret_cast<__half*>(&x_data); + float2 result[2]; + + result[0].x = (float)x_data_h[0] * scale * m[0]; + result[0].y = (float)x_data_h[1] * scale * m[1]; + result[1].x = (float)x_data_h[2] * scale * m[2]; + result[1].y = (float)x_data_h[3] * scale * m[3]; + + result_h[0] = __float22half2_rn(result[0]); + result_h[1] = __float22half2_rn(result[1]); + + out_cast[j] = result_f; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + for (int i = high_index; i < N; i++) { + out[i] = __float2half((float)Xdata[i] * scale * mask[i]); + } + } +} + +template +void launch_dropout_grad(T* vals_out, + const T* vals, + uint8_t* mask, + int total_count, + float ratio, + cudaStream_t stream) +{ + assert(unroll_factor == 4); + + const float scale = 1. / (1. - ratio); + dropout_grad_kernel<<>>(total_count, scale, vals, vals_out, mask); +} +template void launch_dropout_grad(float*, + const float* vals, + uint8_t* mask, + int total_count, + float ratio, + cudaStream_t stream); +template void launch_dropout_grad(__half*, + const __half* vals, + uint8_t* mask, + int total_count, + float ratio, + cudaStream_t stream); + +__global__ void dropout_kernel(const int N, + const int dim, + const float ratio, + const float* bias, + float* Xdata, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int tid = threadIdx.x % (dim / unroll_factor); + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + float4* Xdata_cast = reinterpret_cast(Xdata); + uint32_t* mask_32 = reinterpret_cast(mask); + const float4* bias_cast = reinterpret_cast(bias); + + CUDA_1D_KERNEL_LOOP(j, N) + { + float4 rand = curand_uniform4(&state); + uint32_t m_32; + uint8_t* m = (uint8_t*)&m_32; + + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + float4 x_data = Xdata_cast[j]; + float4 b_data = bias_cast[j % (dim / unroll_factor)]; + + x_data.x += b_data.x; + x_data.y += b_data.y; + x_data.z += b_data.z; + x_data.w += b_data.w; + + x_data.x = x_data.x * scale * m[0]; + x_data.y = x_data.y * scale * m[1]; + x_data.z = x_data.z * scale * m[2]; + x_data.w = x_data.w * scale * m[3]; + + mask_32[j] = m_32; + Xdata_cast[j] = x_data; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + float4 rand = curand_uniform4(&state); + float* rand_data = &(rand.x); + int k = 0; + for (int i = high_index; i < N; i++) { + float x_data = Xdata[i] + bias[i % dim]; + uint8_t m = (uint8_t)(rand_data[k++] > ratio); + Xdata[i] = x_data * scale * m; + mask[i] = m; + } + } +} + +__global__ void dropout_kernel(const int N, + const int dim, + const float ratio, + const __half* bias, + __half* Xdata, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int tid = threadIdx.x % (dim / unroll_factor); + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + float2* Xdata_cast = reinterpret_cast(Xdata); + uint32_t* mask_32 = reinterpret_cast(mask); + const float2* bias_cast = reinterpret_cast(bias); + + CUDA_1D_KERNEL_LOOP(j, N) + { + float4 rand = curand_uniform4(&state); + + float2 data_f; + __half2* data_h = reinterpret_cast<__half2*>(&data_f); + + float2 bias_f; + __half2* bias_h = reinterpret_cast<__half2*>(&bias_f); + + data_f = Xdata_cast[j]; + bias_f = bias_cast[j % (dim / unroll_factor)]; + + float2 data_h_0 = __half22float2(data_h[0]); + float2 data_h_1 = __half22float2(data_h[1]); + + float2 bias_h_0 = __half22float2(bias_h[0]); + float2 bias_h_1 = __half22float2(bias_h[1]); + + data_h_0.x += bias_h_0.x; + data_h_0.y += bias_h_0.y; + data_h_1.x += bias_h_1.x; + data_h_1.y += bias_h_1.y; + + uint32_t m_32; + uint8_t* m = (uint8_t*)&m_32; + + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + data_h_0.x = __float2half(data_h_0.x * scale * m[0]); + data_h_0.y = __float2half(data_h_0.y * scale * m[1]); + data_h_1.x = __float2half(data_h_1.x * scale * m[2]); + data_h_1.y = __float2half(data_h_1.y * scale * m[3]); + + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + + result_h[0] = __float22half2_rn(data_h_0); + result_h[1] = __float22half2_rn(data_h_1); + + Xdata_cast[j] = result_f; + mask_32[j] = m_32; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + float4 rand = curand_uniform4(&state); + float* rand_data = &(rand.x); + int k = 0; + for (int i = high_index; i < N; i++) { + float x_data = (float)Xdata[i] + (float)bias[i % dim]; + uint8_t m = (uint8_t)(rand_data[k++] > ratio); + Xdata[i] = __float2half(x_data * scale * m); + mask[i] = m; + } + } +} + +template +void launch_dropout(T* out, + const T* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream) +{ + assert(unroll_factor == 4); + + int total_count = batch * dim / unroll_factor; + + dim3 grid_dim = DS_GET_BLOCKS(total_count); + dim3 block_dim = DS_CUDA_NUM_THREADS; + + uint64_t inc = (batch * dim) / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + + dropout_kernel<<>>( + total_count, dim, ratio, bias, out, mask, seed); +} + +template void launch_dropout(float*, + const float* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream); +template void launch_dropout(__half*, + const __half* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream); + +__global__ void dropout_kernel(const int N, + const int dim, + const float ratio, + const float* input, + const float* residual, + const float* bias, + float* out, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int tid = threadIdx.x % (dim / unroll_factor); + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + float4* out_cast = reinterpret_cast(out); + uint32_t* mask_32 = reinterpret_cast(mask); + + const float4* bias_cast = reinterpret_cast(bias); + const float4* residual_cast = reinterpret_cast(residual); + const float4* input_cast = reinterpret_cast(input); + + CUDA_1D_KERNEL_LOOP(j, N) + { + float4 rand = curand_uniform4(&state); + + uint32_t m_32; + uint8_t* m = (uint8_t*)&m_32; + + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + float4 out_data; + float4 b_data = bias_cast[j % (dim / unroll_factor)]; + float4 res_data = residual_cast[j]; + float4 inp_data = input_cast[j]; + + out_data.x = (b_data.x + inp_data.x); + out_data.y = (b_data.y + inp_data.y); + out_data.z = (b_data.z + inp_data.z); + out_data.w = (b_data.w + inp_data.w); + + out_data.x = out_data.x * scale * m[0]; + out_data.y = out_data.y * scale * m[1]; + out_data.z = out_data.z * scale * m[2]; + out_data.w = out_data.w * scale * m[3]; + + out_data.x += res_data.x; + out_data.y += res_data.y; + out_data.z += res_data.z; + out_data.w += res_data.w; + + mask_32[j] = m_32; + out_cast[j] = out_data; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + float4 rand = curand_uniform4(&state); + float* rand_data = &(rand.x); + int k = 0; + for (int i = high_index; i < N; i++) { + float x_data = input[i] + bias[i % dim]; + uint8_t m = (uint8_t)(rand_data[k++] > ratio); + x_data = x_data * scale * m; + x_data += residual[i]; + + out[i] = x_data; + mask[i] = m; + } + } +} + +__global__ void dropout_kernel(const int N, + const int dim, + const float ratio, + const __half* input, + const __half* residual, + const __half* bias, + __half* out, + uint8_t* mask, + std::pair seed) +{ + const float scale = 1. / (1. - ratio); + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int tid = threadIdx.x % (dim / unroll_factor); + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + float2* out_cast = reinterpret_cast(out); + uint32_t* mask_32 = reinterpret_cast(mask); + + const float2* bias_cast = reinterpret_cast(bias); + const float2* residual_cast = reinterpret_cast(residual); + const float2* input_cast = reinterpret_cast(input); + + CUDA_1D_KERNEL_LOOP(j, N) + { + float4 rand = curand_uniform4(&state); + + float2 data_f; + __half2* data_h = reinterpret_cast<__half2*>(&data_f); + + float2 bias_f; + __half2* bias_h = reinterpret_cast<__half2*>(&bias_f); + + float2 residual_f; + __half2* residual_h = reinterpret_cast<__half2*>(&residual_f); + + float2 input_f; + __half2* input_h = reinterpret_cast<__half2*>(&input_f); + + bias_f = bias_cast[j % (dim / unroll_factor)]; + residual_f = residual_cast[j]; + input_f = input_cast[j]; + + float2 data_h_0 = __half22float2(data_h[0]); + float2 data_h_1 = __half22float2(data_h[1]); + + float2 bias_h_0 = __half22float2(bias_h[0]); + float2 bias_h_1 = __half22float2(bias_h[1]); + + float2 residual_h_0 = __half22float2(residual_h[0]); + float2 residual_h_1 = __half22float2(residual_h[1]); + + float2 input_h_0 = __half22float2(input_h[0]); + float2 input_h_1 = __half22float2(input_h[1]); + + data_h_0.x = (bias_h_0.x + input_h_0.x); + data_h_0.y = (bias_h_0.y + input_h_0.y); + data_h_1.x = (bias_h_1.x + input_h_1.x); + data_h_1.y = (bias_h_1.y + input_h_1.y); + + uint32_t m_32; + uint8_t* m = (uint8_t*)&m_32; + + m[0] = (uint8_t)(rand.x > ratio); + m[1] = (uint8_t)(rand.y > ratio); + m[2] = (uint8_t)(rand.z > ratio); + m[3] = (uint8_t)(rand.w > ratio); + + data_h_0.x = __float2half(data_h_0.x * scale * m[0]); + data_h_0.y = __float2half(data_h_0.y * scale * m[1]); + data_h_1.x = __float2half(data_h_1.x * scale * m[2]); + data_h_1.y = __float2half(data_h_1.y * scale * m[3]); + + data_h_0.x += residual_h_0.x; + data_h_0.y += residual_h_0.y; + data_h_1.x += residual_h_1.x; + data_h_1.y += residual_h_1.y; + + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + + result_h[0] = __float22half2_rn(data_h_0); + result_h[1] = __float22half2_rn(data_h_1); + + out_cast[j] = result_f; + mask_32[j] = m_32; + } + int high_index = + ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; + if (N > high_index) { + float4 rand = curand_uniform4(&state); + float* rand_data = &(rand.x); + int k = 0; + for (int i = high_index; i < N; i++) { + float x_data = (float)input[i] + (float)bias[i % dim]; + uint8_t m = (uint8_t)(rand_data[k++] > ratio); + x_data = x_data * scale * m; + x_data += (float)residual[i]; + + out[i] = __float2half(x_data); + mask[i] = m; + } + } +} + +template +void launch_dropout(T* out, + const T* input, + const T* residual, + const T* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream) +{ + assert(unroll_factor == 4); + + int total_count = batch * dim / unroll_factor; + dim3 grid_dim = DS_GET_BLOCKS(total_count); + dim3 block_dim = DS_CUDA_NUM_THREADS; + + uint64_t inc = (batch * dim) / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + + dropout_kernel<<>>( + total_count, dim, ratio, input, residual, bias, out, mask, seed); +} + +template void launch_dropout(float*, + const float*, + const float* residual, + const float* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream); +template void launch_dropout(__half*, + const __half*, + const __half* residual, + const __half* bias, + uint8_t* mask, + int batch, + int dim, + float ratio, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp new file mode 100644 index 0000000000000000000000000000000000000000..b637bb710c67a4f980571cd876c4caa350f4a71b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp @@ -0,0 +1,1055 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include +#include +#include +#include +#include +#include +#include "Timer.h" +#include "context.h" +#include "cublas_wrappers.h" +#include "custom_cuda_layers.h" +#include "ds_transformer_cuda.h" + +static std::unordered_map> s_transformer_layers; + +const int init_seq_length = 128; + +// C++ interface + +template +unsigned get_workspace_size(unsigned maxBatchSize, + unsigned seq_len, + unsigned hidden_size, + unsigned intermediate_size, + unsigned heads, + bool training, + bool gelu_checkpoint) +{ + unsigned workSpacesize = 4 * (size_t(maxBatchSize) * seq_len * hidden_size); + if (training) { + workSpacesize += 2 * (size_t(maxBatchSize) * seq_len * hidden_size); + workSpacesize += ((std::max)((size_t(maxBatchSize) * seq_len * intermediate_size), + 2 * (size_t(maxBatchSize) * heads * seq_len * seq_len))); + if (gelu_checkpoint) + workSpacesize += 2 * (size_t(maxBatchSize) * seq_len * intermediate_size); + } + return workSpacesize; // * sizeof(T); +} + +// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. +#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) + +template +BertTransformerLayer::BertTransformerLayer(unsigned layer_id, + unsigned batch_size, + unsigned hidden_size, + unsigned num_heads, + unsigned intermediate_size, + unsigned seq_length, + float attn_prob_dropout_ratio, + float hidden_output_dropout_ratio, + float layer_norm_eps, + bool pre_or_postLayerNorm, + const std::vector>& gemm_algos, + bool attn_dropout_checkpoint, + bool normalize_invertible, + bool gelu_checkpoint, + bool stochastic_mode) + : _layer_id(layer_id), + _batch_size(batch_size), + _hidden_size(hidden_size), + _heads(num_heads), + _intermediate_size(intermediate_size), + _seq_length(seq_length), + _training(true), + _pre_or_postLayerNorm(pre_or_postLayerNorm), + _attn_dropout_checkpoint(attn_dropout_checkpoint), + _normalize_invertible(normalize_invertible), + _gelu_checkpoint(gelu_checkpoint), + _stochastic_mode(stochastic_mode), + _stream(TrainingContext::Instance().GetCurrentStream()), + _cublasHandle(TrainingContext::Instance().GetCublasHandle()), + _qkv_linear(typename FeedForward::Config(batch_size * seq_length, + 3 * hidden_size, + hidden_size, + gemm_algos[0])), + _attn_out_linear(typename FeedForward::Config(batch_size * seq_length, + hidden_size, + hidden_size, + gemm_algos[0])), + _attn_layer_norm(typename Normalize_Layer::Config(batch_size, + seq_length, + hidden_size, + layer_norm_eps, + true, + !normalize_invertible)), + _layer_norm(typename Normalize_Layer::Config(batch_size, + seq_length, + hidden_size, + layer_norm_eps, + true, + !normalize_invertible)), + _ff1(typename FeedForward::Config(batch_size * seq_length, + _intermediate_size, + hidden_size, + gemm_algos[1])), + _ff2(typename FeedForward::Config(batch_size * seq_length, + hidden_size, + _intermediate_size, + gemm_algos[2])), + _softmax(typename Softmax::Config(batch_size, num_heads, seq_length)), + _gelu(typename Gelu::Config(_intermediate_size)), + _attn_prob_dropout(typename Dropout::Config(attn_prob_dropout_ratio, _seq_length)), + _attn_output_dropout(typename Dropout::Config(hidden_output_dropout_ratio, _hidden_size)), + _layer_output_dropout(typename Dropout::Config(hidden_output_dropout_ratio, _hidden_size)), + _attn_scores(typename StridedBatchGemm::Config(_batch_size * _heads, + _seq_length, + _seq_length, + _hidden_size / _heads, + (T(1.0) / T(sqrt(_hidden_size / _heads))), + T(0.0), + CUBLAS_OP_T, + CUBLAS_OP_N, + gemm_algos[3])), + _attn_context(typename StridedBatchGemm::Config(_batch_size * _heads, + _hidden_size / _heads, + _seq_length, + _seq_length, + T(1.0), + T(0.0), + CUBLAS_OP_N, + CUBLAS_OP_N, + gemm_algos[4])) +{ + assert(_hidden_size % _heads == 0); + + Initialize(); +} + +template +BertTransformerLayer::~BertTransformerLayer() +{ +} + +template +void BertTransformerLayer::Initialize() +{ +#ifndef __HIP_PLATFORM_AMD__ + if (std::is_same::value) cublasSetMathMode(_cublasHandle, CUBLAS_TENSOR_OP_MATH); +#endif +} + +template +void BertTransformerLayer::Forward(unsigned bsz, + const T* input_ptr, + const T* input_mask_ptr, + const T* attn_qkvw_ptr, + const T* attn_qkvb_ptr, + const T* attn_ow_ptr, + const T* attn_ob_ptr, + const T* attn_nw_ptr, + const T* attn_nb_ptr, + const T* inter_w_ptr, + const T* inter_b_ptr, + const T* output_w_ptr, + const T* output_b_ptr, + const T* norm_w_ptr, + const T* norm_b_ptr, + T* out_ptr, + T* inp_norm_ptr, + T* q_tf_ptr, + T* k_tf_ptr, + T* v_tf_ptr, + T* soft_out_ptr, + T* ctx_bufB_ptr, + T* attn_o_inp_ptr, + T* add_res_ptr, + T* ff1_inp_ptr, + T* gelu_inp_ptr, + T* ff2_inp_ptr) +{ + cublasSetStream(_cublasHandle, _stream); + + if (!_stochastic_mode) cudaStreamSynchronize(_stream); + + T* workspace = static_cast(TrainingContext::Instance().GetWorkSpace()); + size_t small_buf_size = bsz * _seq_length * _hidden_size; + T* buf_0 = workspace; + T* buf_1 = buf_0 + small_buf_size; + T* buf_2 = buf_1; + + if (_normalize_invertible) { + add_res_ptr = buf_1 + 3 * small_buf_size; + buf_2 = add_res_ptr; + } + if (_gelu_checkpoint) buf_2 += small_buf_size; + if (_attn_dropout_checkpoint) + ctx_bufB_ptr = + (_gelu_checkpoint ? (buf_2 + (_intermediate_size / _hidden_size) * small_buf_size) + : (buf_1 + 4 * small_buf_size)); + + int bsz_seq = bsz * _seq_length; + + if (_pre_or_postLayerNorm) { + if (_layer_norm.UseMean()) + _layer_norm.ForwardCheckpoint( + bsz_seq, inp_norm_ptr, input_ptr, norm_w_ptr, norm_b_ptr, _stream, true); + + else + _layer_norm.Forward( + bsz_seq, inp_norm_ptr, input_ptr, norm_w_ptr, norm_b_ptr, _stream, true); + } + + if (_pre_or_postLayerNorm) + _qkv_linear.Forward(bsz_seq, inp_norm_ptr, attn_qkvw_ptr, buf_0, _cublasHandle); + else + _qkv_linear.Forward(bsz_seq, input_ptr, attn_qkvw_ptr, buf_0, _cublasHandle); + + launch_bias_add_transform_0213( + q_tf_ptr, buf_0, attn_qkvb_ptr, bsz, _seq_length, _hidden_size, _heads, _stream, 3); + + int bsz_heads = bsz * _heads; + + // attention scores + _attn_scores.Forward(bsz_heads, soft_out_ptr, k_tf_ptr, q_tf_ptr, _cublasHandle); + + // Softmax + Mask + _softmax.Forward(bsz, soft_out_ptr, input_mask_ptr, _stream); + + // attn prob dropout. + _attn_prob_dropout.Forward(bsz_heads * _seq_length, ctx_bufB_ptr, soft_out_ptr, _stream); + + // attention context + _attn_context.Forward(bsz_heads, buf_1, v_tf_ptr, ctx_bufB_ptr, _cublasHandle); + + launch_transform4d_0213( + attn_o_inp_ptr, buf_1, bsz, _heads, _seq_length, _hidden_size, _stream, 1); + + if (_pre_or_postLayerNorm) + _attn_out_linear.Forward(bsz_seq, attn_o_inp_ptr, attn_ow_ptr, buf_1, _cublasHandle); + else + _attn_out_linear.Forward(bsz_seq, attn_o_inp_ptr, attn_ow_ptr, ff1_inp_ptr, _cublasHandle); + + // attn output dropout. + if (_pre_or_postLayerNorm) + _attn_output_dropout.ForwardWithBias( + bsz_seq, add_res_ptr, buf_1, input_ptr, attn_ob_ptr, _stream); + else + _attn_output_dropout.ForwardWithBias( + bsz_seq, add_res_ptr, ff1_inp_ptr, input_ptr, attn_ob_ptr, _stream); + + if (_pre_or_postLayerNorm) { + if (_attn_layer_norm.UseMean()) + _attn_layer_norm.ForwardCheckpoint( + bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); + else + _attn_layer_norm.Forward( + bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); + } else { + if (_attn_layer_norm.UseMean()) + _attn_layer_norm.ForwardCheckpoint( + bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); + else + _attn_layer_norm.Forward( + bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); + } + + _ff1.Forward(bsz_seq, + ff1_inp_ptr, + inter_w_ptr, + (_gelu_checkpoint ? ff2_inp_ptr : gelu_inp_ptr), + _cublasHandle); + + _gelu.ForwardWithBiasAdd(bsz_seq, + (_gelu_checkpoint ? ff2_inp_ptr : gelu_inp_ptr), + inter_b_ptr, + (_gelu_checkpoint ? buf_2 : ff2_inp_ptr), + _stream); + + _ff2.Forward( + bsz_seq, (_gelu_checkpoint ? buf_2 : ff2_inp_ptr), output_w_ptr, out_ptr, _cublasHandle); + + // layer output dropout. + if (_pre_or_postLayerNorm) + _layer_output_dropout.ForwardWithBias( + bsz_seq, out_ptr, out_ptr, add_res_ptr, output_b_ptr, _stream); + else + _layer_output_dropout.ForwardWithBias( + bsz_seq, inp_norm_ptr, out_ptr, ff1_inp_ptr, output_b_ptr, _stream); + + if (!_pre_or_postLayerNorm) { + if (_layer_norm.UseMean()) + _layer_norm.ForwardCheckpoint( + bsz_seq, out_ptr, inp_norm_ptr, norm_w_ptr, norm_b_ptr, _stream, true); + else + _layer_norm.Forward( + bsz_seq, out_ptr, inp_norm_ptr, norm_w_ptr, norm_b_ptr, _stream, true); + } +} + +template +void BertTransformerLayer::Backward(unsigned bsz, + const T* grad_output_ptr, + const T* input_ptr, + const T* output_ptr, + const T* inp_norm_ptr, + const T* q_tf_ptr, + const T* k_tf_ptr, + const T* v_tf_ptr, + const T* soft_out_ptr, + const T* ctx_bufB_ptr, + const T* attn_o_inp_ptr, + const T* add_res_ptr, + const T* ff1_inp_ptr, + const T* gelu_inp_ptr, + const T* ff2_inp_ptr, + const T* input_mask_ptr, + const T* attn_qkvw_ptr, + const T* attn_ow_ptr, + const T* attn_nw_ptr, + const T* attn_nb_ptr, + const T* inter_w_ptr, + const T* inter_b_ptr, + const T* output_w_ptr, + const T* norm_w_ptr, + const T* norm_b_ptr, + + T* grad_input_ptr, + T* grad_attn_qkvw_ptr, + T* grad_attn_qkvb_ptr, + T* grad_attn_ow_ptr, + T* grad_attn_ob_ptr, + T* grad_attn_nw_ptr, + T* grad_attn_nb_ptr, + T* grad_inter_w_ptr, + T* grad_inter_b_ptr, + T* grad_output_w_ptr, + T* grad_output_b_ptr, + T* grad_norm_w_ptr, + T* grad_norm_b_ptr) +{ + cublasSetStream(_cublasHandle, _stream); + + if (!_stochastic_mode) cudaStreamSynchronize(_stream); + + T* workspace = static_cast(TrainingContext::Instance().GetWorkSpace()); + size_t small_buf_size = bsz * _seq_length * _hidden_size; + T* buf_0 = workspace; + T* buf_1 = buf_0 + small_buf_size; + T* buf_2 = buf_1 + small_buf_size; + T* buf_3 = buf_2 + small_buf_size; + + T* ff2_buf = (_gelu_checkpoint ? buf_3 + (bsz * _seq_length * _intermediate_size) + : buf_3 + small_buf_size); + T* ctx_bufB_ptr_recomp = ff2_buf + (_seq_length * _seq_length * bsz * _heads); + + cudaStream_t streams[2] = {_stream, _stream}; + + int bsz_seq = bsz * _seq_length; + int bsz_heads = bsz * _heads; + + if (!_pre_or_postLayerNorm) { + if (_layer_norm.UseMean()) + _layer_norm.Backward(bsz_seq, + grad_output_ptr, + norm_w_ptr, + grad_norm_w_ptr, + grad_norm_b_ptr, + streams, + buf_1, + inp_norm_ptr); + + else + _layer_norm.Backward(bsz_seq, + grad_output_ptr, + norm_w_ptr, + norm_b_ptr, + grad_norm_w_ptr, + grad_norm_b_ptr, + streams, + buf_1, + output_ptr); + } + + if (_pre_or_postLayerNorm) + _layer_output_dropout.Backward(bsz_seq, buf_0, grad_output_ptr, _stream); + else + _layer_output_dropout.Backward(bsz_seq, buf_0, buf_1, _stream); + + const T* layer_dropout_buf = _layer_output_dropout.HasDropout() + ? buf_0 + : (_pre_or_postLayerNorm ? grad_output_ptr : buf_1); + + if (_gelu_checkpoint) + _gelu.ForwardWithBiasAdd(bsz_seq, ff2_inp_ptr, inter_b_ptr, buf_2, _stream); + _ff2.Backward(bsz_seq, + layer_dropout_buf, + (_gelu_checkpoint ? buf_2 : ff2_inp_ptr), + output_w_ptr, + grad_output_w_ptr, + grad_output_b_ptr, + _cublasHandle, + _stream, + ff2_buf); + + _gelu.Backward( + bsz_seq, ff2_buf, (_gelu_checkpoint ? ff2_inp_ptr : gelu_inp_ptr), inter_b_ptr, _stream); + + _ff1.Backward(bsz_seq, + ff2_buf, + ff1_inp_ptr, + inter_w_ptr, + grad_inter_w_ptr, + grad_inter_b_ptr, + _cublasHandle, + _stream, + buf_3); + + if (!_pre_or_postLayerNorm) + launch_fused_add2(buf_2, buf_3, buf_1, bsz, _seq_length, _hidden_size, _stream); + + if (_pre_or_postLayerNorm) { + if (_attn_layer_norm.UseMean()) + _attn_layer_norm.BackwardFusedAdd(bsz_seq, + buf_3, + grad_output_ptr, + attn_nw_ptr, + grad_attn_nw_ptr, + grad_attn_nb_ptr, + streams, + buf_0, + add_res_ptr); + + else + _attn_layer_norm.BackwardFusedAdd(bsz_seq, + buf_3, + grad_output_ptr, + attn_nw_ptr, + attn_nb_ptr, + grad_attn_nw_ptr, + grad_attn_nb_ptr, + streams, + buf_0, + ff1_inp_ptr); + } else { + if (_attn_layer_norm.UseMean()) + _attn_layer_norm.Backward(bsz_seq, + buf_2, + attn_nw_ptr, + grad_attn_nw_ptr, + grad_attn_nb_ptr, + streams, + buf_0, + add_res_ptr); + + else + _attn_layer_norm.Backward(bsz_seq, + buf_2, + attn_nw_ptr, + attn_nb_ptr, + grad_attn_nw_ptr, + grad_attn_nb_ptr, + streams, + buf_0, + ff1_inp_ptr); + } + + _attn_output_dropout.Backward(bsz_seq, buf_2, buf_0, _stream); + + T* attn_output_dropout_buf = _attn_output_dropout.HasDropout() ? buf_2 : buf_0; + + _attn_out_linear.Backward(bsz_seq, + attn_output_dropout_buf, + attn_o_inp_ptr, + attn_ow_ptr, + grad_attn_ow_ptr, + grad_attn_ob_ptr, + _cublasHandle, + _stream, + buf_1); + + launch_transform_0213(buf_2, buf_1, bsz, _seq_length, _hidden_size, _heads, _stream); + + if (_attn_prob_dropout.HasDropout()) { + if (_attn_dropout_checkpoint) + _attn_prob_dropout.Forward( + bsz_heads * _seq_length, ctx_bufB_ptr_recomp, soft_out_ptr, _stream, true); + + _attn_context.Backward(bsz_heads, + buf_2, + v_tf_ptr, + (_attn_dropout_checkpoint ? ctx_bufB_ptr_recomp : ctx_bufB_ptr), + _cublasHandle, + buf_3, + ff2_buf); + } else + _attn_context.Backward( + bsz_heads, buf_2, v_tf_ptr, soft_out_ptr, _cublasHandle, buf_3, ff2_buf); + + _attn_prob_dropout.Backward(bsz_heads * _seq_length, ff2_buf, _stream); + + _softmax.Backward(bsz, ff2_buf, soft_out_ptr, _stream); + + _attn_scores.Backward(bsz_heads, ff2_buf, k_tf_ptr, q_tf_ptr, _cublasHandle, buf_2, buf_1); + + launch_transform4d_0213(ff2_buf, buf_1, bsz, _heads, _seq_length, _hidden_size, _stream, 3); + + if (_pre_or_postLayerNorm) + _qkv_linear.Backward(bsz_seq, + ff2_buf, + inp_norm_ptr, + attn_qkvw_ptr, + grad_attn_qkvw_ptr, + grad_attn_qkvb_ptr, + _cublasHandle, + _stream, + buf_2); + else + _qkv_linear.Backward(bsz_seq, + ff2_buf, + input_ptr, + attn_qkvw_ptr, + grad_attn_qkvw_ptr, + grad_attn_qkvb_ptr, + _cublasHandle, + _stream, + buf_2); + + if (_pre_or_postLayerNorm) { + if (_layer_norm.UseMean()) + _layer_norm.BackwardFusedAdd(bsz_seq, + buf_2, + buf_0, + norm_w_ptr, + grad_norm_w_ptr, + grad_norm_b_ptr, + streams, + grad_input_ptr, + input_ptr); + + else + _layer_norm.BackwardFusedAdd(bsz_seq, + buf_2, + buf_0, + norm_w_ptr, + norm_b_ptr, + grad_norm_w_ptr, + grad_norm_b_ptr, + streams, + grad_input_ptr, + inp_norm_ptr); + } else + launch_fused_add2(grad_input_ptr, buf_2, buf_0, bsz, _seq_length, _hidden_size, _stream); +} + +template +void BertTransformerLayer::SetTrainingMode(bool training) +{ + // Dropout will be skipped when not in training model. + _attn_prob_dropout.SetTrainingMode(training); + _attn_output_dropout.SetTrainingMode(training); + _layer_output_dropout.SetTrainingMode(training); +} + +template +void BertTransformerLayer::SetIntermediateBuffers(uint8_t* attn_prob_dropout_mask_ptr, + uint8_t* attn_output_dropout_mask_ptr, + uint8_t* layer_output_dropout_mask_ptr, + T* attn_layer_norm_var, + T* attn_layer_norm_mean, + T* layer_norm_var, + T* layer_norm_mean) +{ + _attn_prob_dropout.SetMask(attn_prob_dropout_mask_ptr); + _attn_output_dropout.SetMask(attn_output_dropout_mask_ptr); + _layer_output_dropout.SetMask(layer_output_dropout_mask_ptr); + + _attn_layer_norm.SetVar(attn_layer_norm_var); + _attn_layer_norm.SetMean(attn_layer_norm_mean); + _layer_norm.SetVar(layer_norm_var); + _layer_norm.SetMean(layer_norm_mean); +} + +template +void BertTransformerLayer::SetSeqLength(unsigned seq_len) +{ + _seq_length = seq_len; + + _softmax.SetSeqLength(_seq_length); + _attn_prob_dropout.SetDimension(_seq_length); + _attn_scores.SetConfig(_seq_length, _seq_length, _hidden_size / _heads); + _attn_context.SetConfig(_hidden_size / _heads, _seq_length, _seq_length); +} + +template +int create_transformer_layer(unsigned layer_id, + unsigned batch_size, + unsigned hidden_dim, + unsigned num_heads, + unsigned intermediate_size, + float attn_dropout_ratio, + float hidden_dropout_ratio, + float layer_norm_eps, + int seed, + bool pre_or_postLayerNorm, + bool test_gemm, + bool attn_dropout_checkpoint, + bool normalize_invertible, + bool gelu_checkpoint, + bool stochastic_mode) +{ + TrainingContext::Instance().SetSeed(seed); + TrainingContext::Instance().TestGemmFP16( + test_gemm, batch_size, init_seq_length, num_heads, hidden_dim / num_heads); + + auto layer = + std::make_shared>(layer_id, + batch_size, + hidden_dim, + num_heads, + intermediate_size, + init_seq_length, + attn_dropout_ratio, + hidden_dropout_ratio, + layer_norm_eps, + pre_or_postLayerNorm, + TrainingContext::Instance().GetGemmAlgos(), + attn_dropout_checkpoint, + normalize_invertible, + gelu_checkpoint, + stochastic_mode); + + s_transformer_layers[layer_id] = layer; + + std::string dtype = (std::is_same::value) ? "half" : "float"; + + std::cout << "layer #" << layer_id << " is created with date type [" << dtype << "]." + << std::endl; + + return 0; +} + +template +std::vector ds_transformer_forward(unsigned layer_id, + const torch::Tensor& input, + const torch::Tensor& input_mask, + const torch::Tensor& attn_qkvw, + const torch::Tensor& attn_qkvb, + const torch::Tensor& attn_ow, + const torch::Tensor& attn_ob, + const torch::Tensor& attn_nw, + const torch::Tensor& attn_nb, + const torch::Tensor& inter_w, + const torch::Tensor& inter_b, + const torch::Tensor& output_w, + const torch::Tensor& output_b, + const torch::Tensor& norm_w, + const torch::Tensor& norm_b, + bool training_mode, + bool prelayernorm, + bool attn_dropout_checkpoint, + bool normalize_invertible, + bool gelu_checkpoint) +{ + CHECK_INPUT(input); + CHECK_INPUT(input_mask); + CHECK_INPUT(attn_qkvw); + CHECK_INPUT(attn_qkvb); + CHECK_INPUT(attn_ow); + CHECK_INPUT(attn_ob); + CHECK_INPUT(attn_nw); + CHECK_INPUT(attn_nb); + CHECK_INPUT(inter_w); + CHECK_INPUT(inter_b); + CHECK_INPUT(output_w); + CHECK_INPUT(output_b); + CHECK_INPUT(norm_w); + CHECK_INPUT(norm_b); + + unsigned bsz = input.size(0); + + const T* input_ptr = (const T*)input.data_ptr(); + const T* input_mask_ptr = (const T*)input_mask.data_ptr(); + const T* attn_qkvw_ptr = (const T*)attn_qkvw.data_ptr(); + const T* attn_qkvb_ptr = (const T*)attn_qkvb.data_ptr(); + const T* attn_ow_ptr = (const T*)attn_ow.data_ptr(); + const T* attn_ob_ptr = (const T*)attn_ob.data_ptr(); + const T* attn_nw_ptr = (const T*)attn_nw.data_ptr(); + const T* attn_nb_ptr = (const T*)attn_nb.data_ptr(); + const T* inter_w_ptr = (const T*)inter_w.data_ptr(); + const T* inter_b_ptr = (const T*)inter_b.data_ptr(); + const T* output_w_ptr = (const T*)output_w.data_ptr(); + const T* output_b_ptr = (const T*)output_b.data_ptr(); + const T* norm_w_ptr = (const T*)norm_w.data_ptr(); + const T* norm_b_ptr = (const T*)norm_b.data_ptr(); + + auto output = torch::empty_like(input); + T* out_ptr = (T*)output.data_ptr(); + + auto options = torch::TensorOptions() + .dtype(input.options().dtype()) + .layout(torch::kStrided) + .device(torch::kCUDA) + .requires_grad(true); + + auto uint8_options = torch::TensorOptions() + .dtype(torch::kInt8) + .layout(torch::kStrided) + .device(torch::kCUDA) + .requires_grad(false); + + std::shared_ptr> layer = + std::static_pointer_cast>(s_transformer_layers[layer_id]); + + unsigned seq_len = layer->GetSeqLength(); + if (input.size(1) != seq_len) { + seq_len = input.size(1); + layer->SetSeqLength(seq_len); + } + + auto workspace = torch::empty({get_workspace_size(bsz, + seq_len, + layer->GetHiddenSize(), + layer->GetIntermediateSize(), + layer->GetNumHeads(), + layer->IsTrainingMode(), + layer->GeluCheckpoint())}, + options); + TrainingContext::Instance().SetWorkSpace((T*)workspace.data_ptr()); + + auto inp_norm = ((prelayernorm || !normalize_invertible) ? torch::empty_like(input) : output); + auto add_res = (normalize_invertible ? inp_norm : torch::empty_like(input)); + auto attn_o_inp = torch::empty_like(input); + auto qkv_tf = torch::empty({(bsz * seq_len), output_w.size(0) * 3}, options); + + auto attn_prob_dropout_mask = + torch::empty({(bsz * layer->GetNumHeads() * seq_len), seq_len}, uint8_options); + auto attn_output_dropout_mask = + torch::empty({(bsz * seq_len), layer->GetHiddenSize()}, uint8_options); + auto layer_output_dropout_mask = + torch::empty({(bsz * seq_len), layer->GetHiddenSize()}, uint8_options); + + auto attn_layer_norm_var = torch::empty({(bsz * seq_len)}, options); + auto attn_layer_norm_mean = torch::empty({(bsz * seq_len)}, options); + auto layer_norm_var = torch::empty({(bsz * seq_len)}, options); + auto layer_norm_mean = torch::empty({(bsz * seq_len)}, options); + + T* inp_norm_ptr = (T*)inp_norm.data_ptr(); + T* add_res_ptr = (T*)add_res.data_ptr(); + T* q_tf_ptr = (T*)qkv_tf.data_ptr(); + T* k_tf_ptr = q_tf_ptr + (bsz * seq_len * output_w.size(0)); //(T*)k_tf.data_ptr(); + T* v_tf_ptr = k_tf_ptr + (bsz * seq_len * output_w.size(0)); //(T*)v_tf.data_ptr(); + T* attn_o_inp_ptr = (T*)attn_o_inp.data_ptr(); + + torch::Tensor ff2_inp = torch::empty({(bsz * seq_len), output_w.size(1)}, options); + torch::Tensor gelu_inp = + (gelu_checkpoint ? ff2_inp : torch::empty({(bsz * seq_len), output_w.size(1)}, options)); + auto ff1_inp = torch::empty_like(input); + T* ff2_inp_ptr = (T*)ff2_inp.data_ptr(); + T* gelu_inp_ptr = (T*)gelu_inp.data_ptr(); + T* ff1_inp_ptr = (T*)ff1_inp.data_ptr(); + + torch::Tensor soft_out = + torch::empty({(bsz * layer->GetNumHeads() * seq_len), seq_len}, options); + torch::Tensor ctx_bufB = + (attn_dropout_checkpoint + ? soft_out + : torch::empty({(bsz * layer->GetNumHeads() * seq_len), seq_len}, options)); + T* soft_out_ptr = (T*)soft_out.data_ptr(); + T* ctx_bufB_ptr = (T*)ctx_bufB.data_ptr(); + + layer->SetTrainingMode(training_mode); + layer->SetIntermediateBuffers((uint8_t*)attn_prob_dropout_mask.data_ptr(), + (uint8_t*)attn_output_dropout_mask.data_ptr(), + (uint8_t*)layer_output_dropout_mask.data_ptr(), + (T*)attn_layer_norm_var.data_ptr(), + (T*)attn_layer_norm_mean.data_ptr(), + (T*)layer_norm_var.data_ptr(), + (T*)layer_norm_mean.data_ptr()); + + layer->Forward(bsz, + input_ptr, + input_mask_ptr, + attn_qkvw_ptr, + attn_qkvb_ptr, + attn_ow_ptr, + attn_ob_ptr, + attn_nw_ptr, + attn_nb_ptr, + inter_w_ptr, + inter_b_ptr, + output_w_ptr, + output_b_ptr, + norm_w_ptr, + norm_b_ptr, + out_ptr, + inp_norm_ptr, + q_tf_ptr, + k_tf_ptr, + v_tf_ptr, + soft_out_ptr, + ctx_bufB_ptr, + attn_o_inp_ptr, + add_res_ptr, + ff1_inp_ptr, + gelu_inp_ptr, + ff2_inp_ptr); + + return {output, + inp_norm, + qkv_tf, + soft_out, + ctx_bufB, + attn_o_inp, + add_res, + ff1_inp, + gelu_inp, + ff2_inp, + attn_prob_dropout_mask, + attn_output_dropout_mask, + layer_output_dropout_mask, + attn_layer_norm_var, + attn_layer_norm_mean, + layer_norm_var, + layer_norm_mean}; +} + +template +std::vector ds_transformer_backward(unsigned layer_id, + const torch::Tensor& grad_output, + const torch::Tensor& output, + const torch::Tensor& inp_norm, + const torch::Tensor& qkv_tf, + const torch::Tensor& soft_out, + const torch::Tensor& ctx_bufB, + const torch::Tensor& attn_o_inp, + const torch::Tensor& add_res, + const torch::Tensor& ff1_inp, + const torch::Tensor& gelu_inp, + const torch::Tensor& ff2_inp, + const torch::Tensor& attn_prob_dropout_mask, + const torch::Tensor& attn_output_dropout_mask, + const torch::Tensor& layer_output_dropout_mask, + const torch::Tensor& attn_layer_norm_var, + const torch::Tensor& attn_layer_norm_mean, + const torch::Tensor& layer_norm_var, + const torch::Tensor& layer_norm_mean, + const torch::Tensor& input, + const torch::Tensor& input_mask, + const torch::Tensor& attn_qkvw, + const torch::Tensor& attn_qkvb, + const torch::Tensor& attn_ow, + const torch::Tensor& attn_ob, + const torch::Tensor& attn_nw, + const torch::Tensor& attn_nb, + const torch::Tensor& inter_w, + const torch::Tensor& inter_b, + const torch::Tensor& output_w, + const torch::Tensor& output_b, + const torch::Tensor& norm_w, + const torch::Tensor& norm_b) +{ + auto g_output = grad_output.contiguous(); + CHECK_INPUT(g_output); + CHECK_INPUT(output); + CHECK_INPUT(inp_norm); + CHECK_INPUT(qkv_tf); + CHECK_INPUT(add_res); + CHECK_INPUT(soft_out); + CHECK_INPUT(ctx_bufB); + CHECK_INPUT(attn_o_inp); + CHECK_INPUT(ff1_inp); + CHECK_INPUT(gelu_inp); + CHECK_INPUT(ff2_inp); + CHECK_INPUT(input); + CHECK_INPUT(input_mask); + CHECK_INPUT(attn_qkvw); + CHECK_INPUT(attn_qkvb); + CHECK_INPUT(attn_ow); + CHECK_INPUT(attn_ob); + CHECK_INPUT(attn_nw); + CHECK_INPUT(attn_nb); + CHECK_INPUT(inter_w); + CHECK_INPUT(inter_b); + CHECK_INPUT(output_w); + CHECK_INPUT(output_b); + CHECK_INPUT(norm_w); + CHECK_INPUT(norm_b); + + unsigned bsz = g_output.size(0); + + std::shared_ptr> layer = + std::static_pointer_cast>(s_transformer_layers[layer_id]); + + unsigned seq_len = layer->GetSeqLength(); + if (g_output.size(1) != seq_len) { + seq_len = g_output.size(1); + layer->SetSeqLength(seq_len); + } + auto options = torch::TensorOptions() + .dtype(g_output.options().dtype()) + .layout(torch::kStrided) + .device(torch::kCUDA) + .requires_grad(true); + auto workspace = torch::empty({get_workspace_size(bsz, + seq_len, + layer->GetHiddenSize(), + layer->GetIntermediateSize(), + layer->GetNumHeads(), + layer->IsTrainingMode(), + layer->GeluCheckpoint())}, + options); + TrainingContext::Instance().SetWorkSpace((T*)workspace.data_ptr()); + + auto grad_input = torch::empty_like(input); + auto grad_attn_qkvw = torch::empty_like(attn_qkvw); + auto grad_attn_qkvb = torch::empty_like(attn_qkvb); + auto grad_attn_ow = torch::empty_like(attn_ow); + auto grad_attn_ob = torch::empty_like(attn_ob); + auto grad_attn_nw = torch::empty_like(attn_nw); + auto grad_attn_nb = torch::empty_like(attn_nb); + auto grad_inter_w = torch::empty_like(inter_w); + auto grad_inter_b = torch::empty_like(inter_b); + auto grad_output_w = torch::empty_like(output_w); + auto grad_output_b = torch::empty_like(output_b); + auto grad_norm_w = torch::empty_like(norm_w); + auto grad_norm_b = torch::empty_like(norm_b); + + // inputs. + const T* grad_output_ptr = (const T*)g_output.data_ptr(); + const T* input_ptr = (const T*)input.data_ptr(); + const T* output_ptr = (const T*)output.data_ptr(); + const T* inp_norm_ptr = (const T*)inp_norm.data_ptr(); + const T* q_tf_ptr = (const T*)qkv_tf.data_ptr(); + const T* add_res_ptr = (const T*)add_res.data_ptr(); + const T* k_tf_ptr = + q_tf_ptr + (bsz * layer->GetSeqLength() * output_w.size(0)); //(const T*)k_tf.data_ptr(); + const T* v_tf_ptr = + k_tf_ptr + (bsz * layer->GetSeqLength() * output_w.size(0)); //(const T*)v_tf.data_ptr(); + const T* ff1_inp_ptr = (const T*)ff1_inp.data_ptr(); + const T* gelu_inp_ptr = (const T*)gelu_inp.data_ptr(); + const T* ff2_inp_ptr = (const T*)ff2_inp.data_ptr(); + const T* ctx_bufB_ptr = (const T*)ctx_bufB.data_ptr(); + const T* soft_out_ptr = (const T*)soft_out.data_ptr(); + const T* attn_o_inp_ptr = (const T*)attn_o_inp.data_ptr(); + const T* input_mask_ptr = (const T*)input_mask.data_ptr(); + const T* attn_qkvw_ptr = (const T*)attn_qkvw.data_ptr(); + const T* attn_ow_ptr = (const T*)attn_ow.data_ptr(); + const T* attn_nw_ptr = (const T*)attn_nw.data_ptr(); + const T* attn_nb_ptr = (const T*)attn_nb.data_ptr(); + const T* inter_w_ptr = (const T*)inter_w.data_ptr(); + const T* inter_b_ptr = (const T*)inter_b.data_ptr(); + const T* output_w_ptr = (const T*)output_w.data_ptr(); + const T* norm_w_ptr = (const T*)norm_w.data_ptr(); + const T* norm_b_ptr = (const T*)norm_b.data_ptr(); + + // outputs. + T* grad_input_ptr = (T*)grad_input.data_ptr(); + T* grad_attn_qkvw_ptr = (T*)grad_attn_qkvw.data_ptr(); + T* grad_attn_qkvb_ptr = (T*)grad_attn_qkvb.data_ptr(); + T* grad_attn_ow_ptr = (T*)grad_attn_ow.data_ptr(); + T* grad_attn_ob_ptr = (T*)grad_attn_ob.data_ptr(); + T* grad_attn_nw_ptr = (T*)grad_attn_nw.data_ptr(); + T* grad_attn_nb_ptr = (T*)grad_attn_nb.data_ptr(); + T* grad_inter_w_ptr = (T*)grad_inter_w.data_ptr(); + T* grad_inter_b_ptr = (T*)grad_inter_b.data_ptr(); + T* grad_output_w_ptr = (T*)grad_output_w.data_ptr(); + T* grad_output_b_ptr = (T*)grad_output_b.data_ptr(); + T* grad_norm_w_ptr = (T*)grad_norm_w.data_ptr(); + T* grad_norm_b_ptr = (T*)grad_norm_b.data_ptr(); + + layer->SetIntermediateBuffers((uint8_t*)attn_prob_dropout_mask.data_ptr(), + (uint8_t*)attn_output_dropout_mask.data_ptr(), + (uint8_t*)layer_output_dropout_mask.data_ptr(), + (T*)attn_layer_norm_var.data_ptr(), + (T*)attn_layer_norm_mean.data_ptr(), + (T*)layer_norm_var.data_ptr(), + (T*)layer_norm_mean.data_ptr()); + + layer->Backward(bsz, + grad_output_ptr, + input_ptr, + output_ptr, + inp_norm_ptr, + q_tf_ptr, + k_tf_ptr, + v_tf_ptr, + soft_out_ptr, + ctx_bufB_ptr, + attn_o_inp_ptr, + add_res_ptr, + ff1_inp_ptr, + gelu_inp_ptr, + ff2_inp_ptr, + input_mask_ptr, + attn_qkvw_ptr, + attn_ow_ptr, + attn_nw_ptr, + attn_nb_ptr, + inter_w_ptr, + inter_b_ptr, + output_w_ptr, + norm_w_ptr, + norm_b_ptr, + + grad_input_ptr, + grad_attn_qkvw_ptr, + grad_attn_qkvb_ptr, + grad_attn_ow_ptr, + grad_attn_ob_ptr, + grad_attn_nw_ptr, + grad_attn_nb_ptr, + grad_inter_w_ptr, + grad_inter_b_ptr, + grad_output_w_ptr, + grad_output_b_ptr, + grad_norm_w_ptr, + grad_norm_b_ptr); + + return {grad_input, + grad_attn_qkvw, + grad_attn_qkvb, + grad_attn_ow, + grad_attn_ob, + grad_attn_nw, + grad_attn_nb, + grad_inter_w, + grad_inter_b, + grad_output_w, + grad_output_b, + grad_norm_w, + grad_norm_b}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("forward_fp32", + &ds_transformer_forward, + "DeepSpeed Transformer forward with fp32 (CUDA)"); + m.def("forward_fp16", + &ds_transformer_forward<__half>, + "DeepSpeed Transformer forward with fp16 (CUDA)"); + m.def("backward_fp32", + &ds_transformer_backward, + "DeepSpeed Transformer backward with fp32 (CUDA)"); + m.def("backward_fp16", + &ds_transformer_backward<__half>, + "DeepSpeed Transformer backward with fp16 (CUDA)"); + m.def("create_transformer_layer_fp32", + &create_transformer_layer, + "Create DeepSpeed Transformer Transformer Layer with fp32 (CUDA)"); + m.def("create_transformer_layer_fp16", + &create_transformer_layer<__half>, + "Create DeepSpeed Transformer Transformer Layer with fp16 (CUDA)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/gelu_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/gelu_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..273891b91923056fe3468600827f538c9cee5009 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/gelu_kernels.cu @@ -0,0 +1,335 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "custom_cuda_layers.h" + +inline __device__ float gelu(const float x) +{ + const float sqrt_param = 0.79788456080286535587989211986876f; + const float mul_param = 0.044715; + return x * 0.5f * (1.0f + tanhf(sqrt_param * (x + mul_param * x * x * x))); +} + +inline __device__ float d_gelu(const float x) +{ + const float sqrt_param = 0.79788456080286535587989211986876f; + const float mul_param = 0.044715; + + float x2mul = x * x * mul_param; + float tan_h = tanhf(sqrt_param * (x + x * x2mul)); + float dg1 = 0.5f * (1.0f + tan_h); + float dg2 = x * 0.5f * sqrt_param * (1 - tan_h * tan_h); + float dg3 = dg2 * 3 * x2mul; + return (dg1 + dg2 + dg3); +} + +/* +Fused bias add with GELU + +Loads a vector of 4 elements each iteration, for stride +iterations. It was written with the intention to launch 256 thread +threadblocks, so to launch for bert-large, we would set ITERATIONS +to 4. This is currently done automatically as a heuristic, setting +the number of iterations as blocks of 1024. + +For FP16, the values are loaded from memory as __half, but converted +to FP32 for the arithmetic itself, to prevent numerous overflow on +the intermediate hyperbolic tangent, since there's no intrinsic +that computes it directly. +*/ + +__global__ void gelu_kernel(const float* input, float* vals, int row_stride, int iterations) +{ + int row = blockIdx.x; + int id = threadIdx.x; + int loop_stride = blockDim.x; + + const float4* input_cast = reinterpret_cast(input); + float4* vals_cast = reinterpret_cast(vals); + + for (int i = 0; i < iterations; i++) { + if (i * loop_stride + id < row_stride) { + float4 data = input_cast[row * row_stride + i * loop_stride + id]; + + data.x = gelu(data.x); + data.y = gelu(data.y); + data.z = gelu(data.z); + data.w = gelu(data.w); + + vals_cast[row * row_stride + i * loop_stride + id] = data; + } + } +} + +__global__ void gelu_kernel(const __half* input, __half* vals, int row_stride, int iterations) +{ +#ifdef HALF_PRECISION_AVAILABLE + int row = blockIdx.x; + int id = threadIdx.x; + int loop_stride = blockDim.x; + + const float2* input_cast = reinterpret_cast(input); + float2* vals_cast = reinterpret_cast(vals); + + for (int i = 0; i < iterations; i++) { + if (i * loop_stride + id < row_stride) { + float2 vals_vec = input_cast[row * row_stride + i * loop_stride + id]; + + __half2* vals_half = reinterpret_cast<__half2*>(&vals_vec); + + float2 low_data = __half22float2(vals_half[0]); + float2 high_data = __half22float2(vals_half[1]); + + low_data.x = gelu(low_data.x); + low_data.y = gelu(low_data.y); + high_data.x = gelu(high_data.x); + high_data.y = gelu(high_data.y); + + vals_half[0] = __float22half2_rn(low_data); + vals_half[1] = __float22half2_rn(high_data); + + vals_cast[row * row_stride + i * loop_stride + id] = vals_vec; + } + } +#endif +} + +__global__ void fused_bias_gelu(const float* input, + const float* bias, + float* vals, + int row_stride, + int iterations) +{ + int row = blockIdx.x; + int id = threadIdx.x; + int loop_stride = blockDim.x; + + const float4* input_cast = reinterpret_cast(input); + float4* vals_cast = reinterpret_cast(vals); + const float4* bias_cast = reinterpret_cast(bias); + + for (int i = 0; i < iterations; i++) { + if (i * loop_stride + id < row_stride) { + float4 data = input_cast[row * row_stride + i * loop_stride + id]; + float4 bias_data = bias_cast[i * loop_stride + id]; + + data.x += bias_data.x; + data.y += bias_data.y; + data.z += bias_data.z; + data.w += bias_data.w; + + data.x = gelu(data.x); + data.y = gelu(data.y); + data.z = gelu(data.z); + data.w = gelu(data.w); + + vals_cast[row * row_stride + i * loop_stride + id] = data; + } + } +} + +__global__ void fused_bias_gelu(const __half* input, + const __half* bias, + __half* vals, + int row_stride, + int iterations) +{ +#ifdef HALF_PRECISION_AVAILABLE + int row = blockIdx.x; + int id = threadIdx.x; + int loop_stride = blockDim.x; + + const float2* input_cast = reinterpret_cast(input); + float2* vals_cast = reinterpret_cast(vals); + const float2* bias_cast = reinterpret_cast(bias); + + for (int i = 0; i < iterations; i++) { + if (i * loop_stride + id < row_stride) { + float2 vals_vec = input_cast[row * row_stride + i * loop_stride + id]; + float2 bias_vec = bias_cast[i * loop_stride + id]; + + __half2* vals_half = reinterpret_cast<__half2*>(&vals_vec); + __half2* bias_half = reinterpret_cast<__half2*>(&bias_vec); + + float2 low_data = __half22float2(vals_half[0]); + float2 high_data = __half22float2(vals_half[1]); + + float2 low_bias = __half22float2(bias_half[0]); + float2 high_bias = __half22float2(bias_half[1]); + + low_data.x += low_bias.x; + low_data.y += low_bias.y; + high_data.x += high_bias.x; + high_data.y += high_bias.y; + + low_data.x = gelu(low_data.x); + low_data.y = gelu(low_data.y); + high_data.x = gelu(high_data.x); + high_data.y = gelu(high_data.y); + + vals_half[0] = __float22half2_rn(low_data); + vals_half[1] = __float22half2_rn(high_data); + + vals_cast[row * row_stride + i * loop_stride + id] = vals_vec; + } + } +#endif +} + +__global__ void d_gelu_func(float* d_output, + const float* gelu_input, + const float* bias, + int row_stride, + int iterations) +{ + int row = blockIdx.x; + int id = threadIdx.x; + int loop_stride = blockDim.x; + + float4* d_output_cast = reinterpret_cast(d_output); + const float4* gelu_input_cast = reinterpret_cast(gelu_input); + const float4* bias_cast = reinterpret_cast(bias); + + for (int i = 0; i < iterations; i++) { + if (i * loop_stride + id < row_stride) { + float4 output_data = d_output_cast[row * row_stride + i * loop_stride + id]; + float4 gelu_input_data = gelu_input_cast[row * row_stride + i * loop_stride + id]; + float4 bias_data = bias_cast[i * loop_stride + id]; + + gelu_input_data.x += bias_data.x; + gelu_input_data.y += bias_data.y; + gelu_input_data.z += bias_data.z; + gelu_input_data.w += bias_data.w; + + output_data.x *= d_gelu(gelu_input_data.x); + output_data.y *= d_gelu(gelu_input_data.y); + output_data.z *= d_gelu(gelu_input_data.z); + output_data.w *= d_gelu(gelu_input_data.w); + + d_output_cast[row * row_stride + i * loop_stride + id] = output_data; + } + } +} + +__global__ void d_gelu_func(__half* d_output, + const __half* gelu_input, + const __half* bias, + int row_stride, + int iterations) +{ +#ifdef HALF_PRECISION_AVAILABLE + int row = blockIdx.x; + int id = threadIdx.x; + int loop_stride = blockDim.x; + + float2* d_output_cast = reinterpret_cast(d_output); + const float2* gelu_input_cast = reinterpret_cast(gelu_input); + const float2* bias_cast = reinterpret_cast(bias); + +#pragma unroll + for (int i = 0; i < iterations; i++) { + if (i * loop_stride + id < row_stride) { + float2 output_data = d_output_cast[row * row_stride + i * loop_stride + id]; + float2 gelu_input_data = gelu_input_cast[row * row_stride + i * loop_stride + id]; + float2 bias_vec = bias_cast[i * loop_stride + id]; + + __half2* output_data_half = reinterpret_cast<__half2*>(&output_data); + __half2* gelu_input_data_half = reinterpret_cast<__half2*>(&gelu_input_data); + __half2* bias_half = reinterpret_cast<__half2*>(&bias_vec); + + float2 output_half_0 = __half22float2(output_data_half[0]); + float2 output_half_1 = __half22float2(output_data_half[1]); + + float2 gelu_input_half_0 = __half22float2(gelu_input_data_half[0]); + float2 gelu_input_half_1 = __half22float2(gelu_input_data_half[1]); + + float2 bias_half_0 = __half22float2(bias_half[0]); + float2 bias_half_1 = __half22float2(bias_half[1]); + + gelu_input_half_0.x += bias_half_0.x; + gelu_input_half_0.y += bias_half_0.y; + gelu_input_half_1.x += bias_half_1.x; + gelu_input_half_1.y += bias_half_1.y; + + output_half_0.x *= d_gelu(gelu_input_half_0.x); + output_half_0.y *= d_gelu(gelu_input_half_0.y); + output_half_1.x *= d_gelu(gelu_input_half_1.x); + output_half_1.y *= d_gelu(gelu_input_half_1.y); + + float2 result; + __half2* result_half2 = reinterpret_cast<__half2*>(&result); + + result_half2[0] = __float22half2_rn(output_half_0); + result_half2[1] = __float22half2_rn(output_half_1); + + d_output_cast[row * row_stride + i * loop_stride + id] = result; + } + } +#endif +} + +template +void launch_bias_gelu(const T* input, + const T* bias, + T* output, + int intermediate_size, + int batch_size, + cudaStream_t stream) +{ + int iterations = (intermediate_size + 1023) / 1024; + int threads = (intermediate_size - 1) / (iterations * 4) + 1; + dim3 block_dims(threads); + dim3 grid_dims(batch_size); + + fused_bias_gelu<<>>( + input, bias, output, intermediate_size / 4, iterations); +} + +template +void launch_gelu(const T* input, + T* output, + int intermediate_size, + int batch_size, + cudaStream_t stream) +{ + int iterations = (intermediate_size + 1023) / 1024; + int threads = (intermediate_size - 1) / (iterations * 4) + 1; + dim3 block_dims(threads); + dim3 grid_dims(batch_size); + + gelu_kernel<<>>( + input, output, intermediate_size / 4, iterations); +} + +template void launch_bias_gelu(const float*, const float*, float*, int, int, cudaStream_t); +template void launch_bias_gelu<__half>(const __half*, + const __half*, + __half*, + int, + int, + cudaStream_t); + +template void launch_gelu(const float*, float*, int, int, cudaStream_t); +template void launch_gelu<__half>(const __half*, __half*, int, int, cudaStream_t); + +template +void launch_d_gelu(T* d_output, + const T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream) +{ + int iterations = (intermediate_size + 1023) / 1024; + int threads = (intermediate_size - 1) / (iterations * 4) + 1; + dim3 block_dims(threads); + dim3 grid_dims(batch_size); + + d_gelu_func<<>>( + d_output, input, bias, intermediate_size / 4, iterations); +} + +template void launch_d_gelu(float*, const float*, const float*, int, int, cudaStream_t); +template void launch_d_gelu<__half>(__half*, const __half*, const __half*, int, int, cudaStream_t); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/general_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/general_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..a987eec5ef0bdc87153422eae635507f44188802 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/general_kernels.cu @@ -0,0 +1,416 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "general_kernels.h" + +namespace cg = cooperative_groups; + +template +__global__ void column_sum_reduce(const T* __restrict__ inp, + T* __restrict__ out, + int rows, + int width) +{ + __shared__ float tile[TILE_DIM][TILE_DIM + 1]; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int idx = blockDim.x * blockIdx.x + threadIdx.x; + + int y_stride = width * TILE_DIM; + + float localSum = 0; + + // Loop across matrix height + if (idx < width) { + int offset = threadIdx.y * width + idx; + for (int r = threadIdx.y; r < rows; r += TILE_DIM) { + localSum += (float)inp[offset]; + offset += y_stride; + } + } + + tile[threadIdx.x][threadIdx.y] = localSum; + + __syncthreads(); + + // Sum the shared buffer. + float sum = tile[threadIdx.y][threadIdx.x]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < TILE_DIM; i <<= 1) sum += g.shfl_down(sum, i); + + if (threadIdx.x == 0) { + int pos = blockIdx.x * TILE_DIM + threadIdx.y; + if (pos < width) out[pos] = sum; + } +} + +template +void launch_fuse_transpose_bias_kernel(const T* inp, + T* out, + int rows, + int cols, + cudaStream_t stream); + +template <> +void launch_fuse_transpose_bias_kernel(const float* inp, + float* out, + int rows, + int cols, + cudaStream_t stream) +{ + // assert(rows % TILE_DIM == 0); + // assert(cols % TILE_DIM == 0); + + dim3 grid_dim((cols - 1) / TILE_DIM + 1); + dim3 block_dim(TILE_DIM, TILE_DIM); + + column_sum_reduce<<>>(inp, out, rows, cols); +} + +template <> +void launch_fuse_transpose_bias_kernel<__half>(const __half* inp, + __half* out, + int rows, + int cols, + cudaStream_t stream) +{ + // assert(rows % TILE_DIM == 0); + // assert(cols % TILE_DIM == 0); + + dim3 grid_dim((cols - 1) / TILE_DIM + 1); + dim3 block_dim(TILE_DIM, TILE_DIM); + + column_sum_reduce<__half><<>>(inp, out, rows, cols); +} + +__global__ void fused_add2_kernel(const int N, float* out, const float* inp1, const float* inp2) +{ + const float4* inp1_4 = reinterpret_cast(inp1); + const float4* inp2_4 = reinterpret_cast(inp2); + float4* out_4 = reinterpret_cast(out); + + CUDA_1D_KERNEL_LOOP(j, N) + { + float4 val; + float4 inp1_reg = inp1_4[j]; + float4 inp2_reg = inp2_4[j]; + + val.x = inp1_reg.x + inp2_reg.x; + val.y = inp1_reg.y + inp2_reg.y; + val.z = inp1_reg.z + inp2_reg.z; + val.w = inp1_reg.w + inp2_reg.w; + + out_4[j] = val; + } +} + +__global__ void fused_add2_kernel(const int N, __half* out, const __half* inp1, const __half* inp2) +{ + float2 inp1_4; + float2 inp2_4; + + __half2* inp1_h = reinterpret_cast<__half2*>(&inp1_4); + __half2* inp2_h = reinterpret_cast<__half2*>(&inp2_4); + + const float2* inp1_arr = reinterpret_cast(inp1); + const float2* inp2_arr = reinterpret_cast(inp2); + + CUDA_1D_KERNEL_LOOP(j, N) + { + inp1_4 = inp1_arr[j]; + inp2_4 = inp2_arr[j]; + + float2 inp1_h_f_0 = __half22float2(inp1_h[0]); + float2 inp1_h_f_1 = __half22float2(inp1_h[1]); + + float2 inp2_h_f_0 = __half22float2(inp2_h[0]); + float2 inp2_h_f_1 = __half22float2(inp2_h[1]); + + inp1_h_f_0.x += inp2_h_f_0.x; + inp1_h_f_0.y += inp2_h_f_0.y; + inp1_h_f_1.x += inp2_h_f_1.x; + inp1_h_f_1.y += inp2_h_f_1.y; + + float2 val_f; + __half2* val_h = reinterpret_cast<__half2*>(&val_f); + + val_h[0] = __float22half2_rn(inp1_h_f_0); + val_h[1] = __float22half2_rn(inp1_h_f_1); + + float2* out_4 = reinterpret_cast(out); + out_4[j] = val_f; + } +} + +template <> +void launch_fused_add2(float* out, + const float* inp1, + const float* inp2, + int batch_size, + int seq_length, + int hidden_dim, + cudaStream_t& stream) +{ + int total_count = batch_size * seq_length * hidden_dim / 4; + dim3 grid_dim = DS_GET_BLOCKS(total_count); //(batch_size * seq_length); + + dim3 block_dim = DS_CUDA_NUM_THREADS; //(hidden_dim / 4); + + fused_add2_kernel<<>>(total_count, out, inp1, inp2); +} + +template <> +void launch_fused_add2<__half>(__half* out, + const __half* inp1, + const __half* inp2, + int batch_size, + int seq_length, + int hidden_dim, + cudaStream_t& stream) +{ + int total_count = batch_size * seq_length * hidden_dim / 4; + dim3 grid_dim = DS_GET_BLOCKS(total_count); //(batch_size * seq_length); + + dim3 block_dim = DS_CUDA_NUM_THREADS; //(hidden_dim / 4); + + fused_add2_kernel<<>>(total_count, out, inp1, inp2); +} + +__global__ void fused_add3_kernel(float* out, + const float* inp1, + const float* inp2, + const float* inp3, + int size, + int row_stride) +{ + int row = blockIdx.x; + int id = threadIdx.x; + + const float4* inp1_4 = reinterpret_cast(inp1); + const float4* inp2_4 = reinterpret_cast(inp2); + const float4* inp3_4 = reinterpret_cast(inp3); + + float4* out_4 = reinterpret_cast(out); + + float4 val; + float4 inp1_reg = inp1_4[row * row_stride + id]; + float4 inp2_reg = inp2_4[row * row_stride + id]; + float4 inp3_reg = inp3_4[row * row_stride + id]; + + val.x = inp1_reg.x + inp2_reg.x + inp3_reg.x; + val.y = inp1_reg.y + inp2_reg.y + inp3_reg.y; + val.z = inp1_reg.z + inp2_reg.z + inp3_reg.z; + val.w = inp1_reg.w + inp2_reg.w + inp3_reg.w; + + out_4[row * row_stride + id] = val; +} + +__global__ void fused_add3_kernel(__half* out, + const __half* inp1, + const __half* inp2, + const __half* inp3, + int size, + int row_stride) +{ + int row = blockIdx.x; + int id = threadIdx.x; + const float2* inp1_arr = reinterpret_cast(inp1); + const float2* inp2_arr = reinterpret_cast(inp2); + const float2* inp3_arr = reinterpret_cast(inp3); + + float2 inp1_4 = inp1_arr[row * row_stride + id]; + float2 inp2_4 = inp2_arr[row * row_stride + id]; + float2 inp3_4 = inp3_arr[row * row_stride + id]; + + __half2* inp1_h = reinterpret_cast<__half2*>(&inp1_4); + __half2* inp2_h = reinterpret_cast<__half2*>(&inp2_4); + __half2* inp3_h = reinterpret_cast<__half2*>(&inp3_4); + + float2 inp1_h_f_0 = __half22float2(inp1_h[0]); + float2 inp1_h_f_1 = __half22float2(inp1_h[1]); + + float2 inp2_h_f_0 = __half22float2(inp2_h[0]); + float2 inp2_h_f_1 = __half22float2(inp2_h[1]); + + float2 inp3_h_f_0 = __half22float2(inp3_h[0]); + float2 inp3_h_f_1 = __half22float2(inp3_h[1]); + + inp1_h_f_0.x += (inp2_h_f_0.x + inp3_h_f_0.x); + inp1_h_f_0.y += (inp2_h_f_0.y + inp3_h_f_0.y); + inp1_h_f_1.x += (inp2_h_f_1.x + inp3_h_f_1.x); + inp1_h_f_1.y += (inp2_h_f_1.y + inp3_h_f_1.y); + + float2 val_f; + __half2* val_h = reinterpret_cast<__half2*>(&val_f); + + val_h[0] = __float22half2_rn(inp1_h_f_0); + val_h[1] = __float22half2_rn(inp1_h_f_1); + + float2* out_4 = reinterpret_cast(out); + out_4[row * row_stride + id] = val_f; +} + +template <> +void launch_fused_add3(float* out, + const float* inp1, + const float* inp2, + const float* inp3, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream) +{ + dim3 grid_dim(batch_size * seq_length); + + dim3 block_dim(hidden_size / 4); + + fused_add3_kernel<<>>( + out, inp1, inp2, inp3, (batch_size * seq_length * hidden_size), hidden_size / 4); +} + +template <> +void launch_fused_add3<__half>(__half* out, + const __half* inp1, + const __half* inp2, + const __half* inp3, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream) +{ + dim3 grid_dim(batch_size * seq_length); + + dim3 block_dim(hidden_size / 4); + + fused_add3_kernel<<>>( + out, inp1, inp2, inp3, (batch_size * seq_length * hidden_size), hidden_size / 4); +} + +__global__ void fused_add4_kernel(float* out, + const float* inp1, + const float* inp2, + const float* inp3, + const float* inp4, + int size, + int row_stride) +{ + int row = blockIdx.x; + int id = threadIdx.x; + + const float4* inp1_4 = reinterpret_cast(inp1); + const float4* inp2_4 = reinterpret_cast(inp2); + const float4* inp3_4 = reinterpret_cast(inp3); + const float4* inp4_4 = reinterpret_cast(inp4); + float4* out_4 = reinterpret_cast(out); + + float4 val; + float4 inp1_reg = inp1_4[row * row_stride + id]; + float4 inp2_reg = inp2_4[row * row_stride + id]; + float4 inp3_reg = inp3_4[row * row_stride + id]; + float4 inp4_reg = inp4_4[row * row_stride + id]; + + val.x = inp1_reg.x + inp2_reg.x + inp3_reg.x + inp4_reg.x; + val.y = inp1_reg.y + inp2_reg.y + inp3_reg.y + inp4_reg.y; + val.z = inp1_reg.z + inp2_reg.z + inp3_reg.z + inp4_reg.z; + val.w = inp1_reg.w + inp2_reg.w + inp3_reg.w + inp4_reg.w; + + out_4[row * row_stride + id] = val; +} + +__global__ void fused_add4_kernel(__half* out, + const __half* inp1, + const __half* inp2, + const __half* inp3, + const __half* inp4, + int size, + int row_stride) +{ + int row = blockIdx.x; + int id = threadIdx.x; + const float2* inp1_arr = reinterpret_cast(inp1); + const float2* inp2_arr = reinterpret_cast(inp2); + const float2* inp3_arr = reinterpret_cast(inp3); + const float2* inp4_arr = reinterpret_cast(inp4); + + float2 inp1_4 = inp1_arr[row * row_stride + id]; + float2 inp2_4 = inp2_arr[row * row_stride + id]; + float2 inp3_4 = inp3_arr[row * row_stride + id]; + float2 inp4_4 = inp4_arr[row * row_stride + id]; + + __half2* inp1_h = reinterpret_cast<__half2*>(&inp1_4); + __half2* inp2_h = reinterpret_cast<__half2*>(&inp2_4); + __half2* inp3_h = reinterpret_cast<__half2*>(&inp3_4); + __half2* inp4_h = reinterpret_cast<__half2*>(&inp4_4); + + float2 inp1_h_f_0 = __half22float2(inp1_h[0]); + float2 inp1_h_f_1 = __half22float2(inp1_h[1]); + + float2 inp2_h_f_0 = __half22float2(inp2_h[0]); + float2 inp2_h_f_1 = __half22float2(inp2_h[1]); + + float2 inp3_h_f_0 = __half22float2(inp3_h[0]); + float2 inp3_h_f_1 = __half22float2(inp3_h[1]); + + float2 inp4_h_f_0 = __half22float2(inp4_h[0]); + float2 inp4_h_f_1 = __half22float2(inp4_h[1]); + + inp1_h_f_0.x += (inp2_h_f_0.x + inp3_h_f_0.x + inp4_h_f_0.x); + inp1_h_f_0.y += (inp2_h_f_0.y + inp3_h_f_0.y + inp4_h_f_0.y); + inp1_h_f_1.x += (inp2_h_f_1.x + inp3_h_f_1.x + inp4_h_f_1.x); + inp1_h_f_1.y += (inp2_h_f_1.y + inp3_h_f_1.y + inp4_h_f_1.y); + + float2 val_f; + __half2* val_h = reinterpret_cast<__half2*>(&val_f); + + val_h[0] = __float22half2_rn(inp1_h_f_0); + val_h[1] = __float22half2_rn(inp1_h_f_1); + + float2* out_4 = reinterpret_cast(out); + out_4[row * row_stride + id] = val_f; +} + +template <> +void launch_fused_add4(float* out, + const float* inp1, + const float* inp2, + const float* inp3, + const float* inp4, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream) +{ + dim3 grid_dim(batch_size * seq_length); + + dim3 block_dim(hidden_size / 4); + + fused_add4_kernel<<>>( + out, inp1, inp2, inp3, inp4, (batch_size * seq_length * hidden_size), hidden_size / 4); +} + +template <> +void launch_fused_add4<__half>(__half* out, + const __half* inp1, + const __half* inp2, + const __half* inp3, + const __half* inp4, + int batch_size, + int seq_length, + int hidden_size, + cudaStream_t& stream) +{ + dim3 grid_dim(batch_size * seq_length); + + dim3 block_dim(hidden_size / 4); + + fused_add4_kernel<<>>( + out, inp1, inp2, inp3, inp4, (batch_size * seq_length * hidden_size), hidden_size / 4); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu new file mode 100644 index 0000000000000000000000000000000000000000..bbb8a7f00b1f5a0be87a6bce7c4c65ff3ee84145 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#ifdef __HIP_PLATFORM_AMD__ +#include "hip/hip_cooperative_groups.h" +#else +#include "cooperative_groups.h" +#endif +#include "ds_kernel_utils.h" +#include "inference_cuda_layers.h" +#include "memory_access_utils.h" + +#ifndef __HIP_PLATFORM_AMD__ +#include +#endif + +namespace cg = cooperative_groups; + +namespace rot_half { +constexpr int threads = 256; +} // namespace rot_half + +template +__global__ void apply_rotary_pos_half(T* mixed_query, + T* key_layer, + unsigned rotary_dim, + unsigned seq_len, + unsigned seq_offset, + unsigned num_heads, + unsigned head_size, + unsigned total_count, + float rope_theta, + int max_out_tokens) +{ + constexpr int T_per_thread = granularity / sizeof(T); + constexpr int heads_per_block = rot_half::threads / threadsPerHead; + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile head_group = cg::tiled_partition(tb); + + const int head_idx = blockIdx.x * heads_per_block + threadIdx.x / threadsPerHead; + const int cur_seq_idx = head_idx % seq_len; + const int offset = head_idx * head_size; + const int k_offset = (cur_seq_idx + (head_idx / seq_len) * max_out_tokens) * head_size; + + const int seq_idx = cur_seq_idx + seq_offset; + const int half_dim = rotary_dim >> 1; + const int half_dim_threads = half_dim / T_per_thread; + + if (head_idx < total_count) { + const int base_neuron_idx = head_group.thread_rank() * T_per_thread; + + T q[T_per_thread], k[T_per_thread]; + mem_access::load_global(q, mixed_query + offset + base_neuron_idx); + mem_access::load_global(k, key_layer + k_offset + base_neuron_idx); + +#pragma unroll + for (int i = 0; i < T_per_thread; i++) { + const int neuron_idx = base_neuron_idx + i; + if (neuron_idx < rotary_dim) { + float inv_freq = (float)((neuron_idx % half_dim) * 2) / (float)rotary_dim; + inv_freq = 1.0 / powf(rope_theta, inv_freq) * (float)seq_idx; + + float rotary_sign = (neuron_idx > (half_dim - 1) ? -1.0 : 1.0); + float q_rot = conversion::to(q[i]) * rotary_sign; + float k_rot = conversion::to(k[i]) * rotary_sign; + + const int target_lane = (neuron_idx < half_dim) + ? head_group.thread_rank() + half_dim_threads + : head_group.thread_rank() - half_dim_threads; + + const float q_rot_temp = head_group.shfl(q_rot, target_lane); + const float k_rot_temp = head_group.shfl(k_rot, target_lane); + + q[i] = conversion::to(conversion::to(q[i]) * cosf(inv_freq) + + q_rot_temp * sinf(inv_freq)); + k[i] = conversion::to(conversion::to(k[i]) * cosf(inv_freq) + + k_rot_temp * sinf(inv_freq)); + } + } + + mem_access::store_global(mixed_query + offset + base_neuron_idx, q); + mem_access::store_global(key_layer + k_offset + base_neuron_idx, k); + } +} + +#define LAUNCH_ROT_POS_EMB_HALF(HEAD_THREADS, ALIGNMENT) \ + apply_rotary_pos_half<<>>(mixed_query, \ + key_layer, \ + rotary_dim, \ + seq_len, \ + offset, \ + num_heads, \ + head_size, \ + total_count, \ + rope_theta, \ + max_out_tokens); + +#if defined(__HIP_PLATFORM_AMD__) and ROCM_WAVEFRONT_SIZE == 64 +#define LAUNCH_FOR_ALIGNMENT(ALIGNMENT) \ + if (threads_per_head == 4) { \ + LAUNCH_ROT_POS_EMB_HALF(4, ALIGNMENT); \ + } else if (threads_per_head == 8) { \ + LAUNCH_ROT_POS_EMB_HALF(8, ALIGNMENT); \ + } else if (threads_per_head == 16) { \ + LAUNCH_ROT_POS_EMB_HALF(16, ALIGNMENT); \ + } else if (threads_per_head == 32) { \ + LAUNCH_ROT_POS_EMB_HALF(32, ALIGNMENT); \ + } else if (threads_per_head == 64) { \ + LAUNCH_ROT_POS_EMB_HALF(64, ALIGNMENT); \ + } else { \ + assert(false); \ + } +#else +#define LAUNCH_FOR_ALIGNMENT(ALIGNMENT) \ + if (threads_per_head == 4) { \ + LAUNCH_ROT_POS_EMB_HALF(4, ALIGNMENT); \ + } else if (threads_per_head == 8) { \ + LAUNCH_ROT_POS_EMB_HALF(8, ALIGNMENT); \ + } else if (threads_per_head == 16) { \ + LAUNCH_ROT_POS_EMB_HALF(16, ALIGNMENT); \ + } else if (threads_per_head == 32) { \ + LAUNCH_ROT_POS_EMB_HALF(32, ALIGNMENT); \ + } else { \ + assert(false); \ + } +#endif + +template +void launch_apply_rotary_pos_emb(T* mixed_query, + T* key_layer, + unsigned head_size, + unsigned seq_len, + unsigned rotary_dim, + unsigned offset, + unsigned num_heads, + unsigned batch, + float rope_theta, + cudaStream_t stream, + int max_out_tokens) +{ + const int half_dim = rotary_dim >> 1; + + int alignment = sizeof(T); + if (half_dim % (16 / sizeof(T)) == 0) { + alignment = 16; + } else if (half_dim % (8 / sizeof(T)) == 0) { + alignment = 8; + } else if (half_dim % (4 / sizeof(T)) == 0) { + alignment = 4; + } else { + assert(false); + } + const int T_per_elem = alignment / sizeof(T); + + int total_count = batch * num_heads * seq_len; + + const int padded_head_size = next_pow2(head_size); + + assert(padded_head_size <= hw_warp_size * T_per_elem); + + const int threads_per_head = padded_head_size / T_per_elem; + const int heads_per_block = rot_half::threads / threads_per_head; + + dim3 block(rot_half::threads); + dim3 grid((total_count + heads_per_block - 1) / heads_per_block); + + if (alignment == 4) { + LAUNCH_FOR_ALIGNMENT(4); + } else if (alignment == 8) { + LAUNCH_FOR_ALIGNMENT(8); + } else if (alignment == 16) { + LAUNCH_FOR_ALIGNMENT(16); + } else { + assert(false); + } +} + +#define INSTANTIATE_LAUNCH_ROTARY_POS_EMB(T) \ + template void launch_apply_rotary_pos_emb(T*, \ + T*, \ + unsigned, \ + unsigned, \ + unsigned, \ + unsigned, \ + unsigned, \ + unsigned, \ + float, \ + cudaStream_t, \ + int); + +INSTANTIATE_LAUNCH_ROTARY_POS_EMB(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_ROTARY_POS_EMB(__nv_bfloat16); +#endif +INSTANTIATE_LAUNCH_ROTARY_POS_EMB(__half); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..7a8e7ca446b0fbc4154fd0198c8e301f77cc7815 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu @@ -0,0 +1,153 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "inference_cuda_layers.h" + +#define MAX_QUANTIZE_GROUPING 1024 + +#define loop_unroll 1 +#define loop_unroll_bits 1 + +template +__global__ void dequantize_kernel(T* output, + const int8_t* input, + const float* qscale, + int output_size, + int hidden_dim, + int groups, + int merge_count) +{ + unsigned merge_hidden = hidden_dim >> merge_count; + unsigned quantization_stride = (merge_hidden * output_size) / groups; + + unsigned bid = blockIdx.x; + unsigned tid = threadIdx.x; + + while (tid < output_size) { + unsigned w_index = bid / merge_hidden; + unsigned q_index = tid + bid * output_size; + + auto q = input[q_index]; + + unsigned merge_hidden_total = w_index * merge_hidden; + unsigned scale_index = + ((((bid - merge_hidden_total) + tid * merge_hidden) / quantization_stride) + << merge_count) + + w_index; + + float scale_data = qscale[scale_index]; + + output[q_index] = conversion::to(scale_data * (float)q); + tid += blockDim.x; + } +} + +template +void launch_dequantize(T* output, + const int8_t* input, + const float* qscale, + unsigned output_size, + unsigned hidden_dim, + unsigned groups, + unsigned merge_count, + cudaStream_t stream) +{ + unsigned threads = 1024; + dim3 block_dims(threads); + dim3 grid_dims(hidden_dim); + + dequantize_kernel<<>>( + output, input, qscale, output_size, hidden_dim, groups, merge_count); +} + +#define INSTANTIATE_DEQUANTIZE_MERGE(T) \ + template void launch_dequantize( \ + T*, const int8_t*, const float*, unsigned, unsigned, unsigned, unsigned, cudaStream_t); + +INSTANTIATE_DEQUANTIZE_MERGE(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_DEQUANTIZE_MERGE(__nv_bfloat16); +#endif +INSTANTIATE_DEQUANTIZE_MERGE(__half); + +__global__ void dequantize_kernel(float* output, + const int8_t* input, + const float* qscale, + int hidden_dim, + unsigned merge_hidden, + int cnt) +{ +} + +template +__global__ void dequantize_kernel(T* output, + const int8_t* input, + const float* qscale, + unsigned hidden_dim, + unsigned merge_hidden, + int cnt) +{ + unsigned bid = blockIdx.x * gridDim.y + blockIdx.y; + unsigned tid = threadIdx.x; + + float local_scale = qscale[blockIdx.x]; + + const float* input_cast = reinterpret_cast(input); + float2* output_cast = reinterpret_cast(output); + + input_cast += bid * merge_hidden; + output_cast += bid * merge_hidden; + + for (int c = 0; c < cnt; c++) { + if (tid < merge_hidden) { + float q = input_cast[tid]; + int8_t* q_int8 = (int8_t*)&q; + + float2 q_f; + T* q_h = (T*)&q_f; + + q_h[0] = conversion::to(local_scale * (float)q_int8[0]); + q_h[1] = conversion::to(local_scale * (float)q_int8[1]); + q_h[2] = conversion::to(local_scale * (float)q_int8[2]); + q_h[3] = conversion::to(local_scale * (float)q_int8[3]); + output_cast[tid] = q_f; + tid += blockDim.x; + } + } +} + +template +void launch_dequantize(T* output, + const int8_t* input, + const float* qscale, + unsigned output_size, + unsigned hidden_dim, + unsigned groups, + cudaStream_t stream) +{ + unsigned threads = 1024; + hidden_dim /= 4; + unsigned thd_cnt = (hidden_dim - 1) / threads + 1; + + assert(output_size % groups == 0); + unsigned blocks = output_size / groups; + + dim3 block_dims(threads); + dim3 grid_dims(groups, blocks); + + dequantize_kernel<<>>( + output, input, qscale, hidden_dim, hidden_dim, thd_cnt); +} + +#define INSTANTIATE_DEQUANTIZE_NO_MERGE(T) \ + template void launch_dequantize( \ + T*, const int8_t*, const float*, unsigned, unsigned, unsigned, cudaStream_t); + +INSTANTIATE_DEQUANTIZE_NO_MERGE(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_DEQUANTIZE_NO_MERGE(__nv_bfloat16); +#endif +INSTANTIATE_DEQUANTIZE_NO_MERGE(__half); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu new file mode 100644 index 0000000000000000000000000000000000000000..97857bc3f70bfa044cbf26e4a82bf6dd781f75b9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu @@ -0,0 +1,716 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "inference_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; +#define MAX_CAP 4 +#define MAX_SEQ 2048 + +// only used to avoid compilation error due to lack of definition. +#ifndef BF16_AVAILABLE +#if defined(__CUDA_BF16_H__) +static_assert(sizeof(__nv_bfloat162) == sizeof(__half2), + "CUDA's __nv_bfloat162 doesn't match __half2 size"); +#else +// Fallback to simple typedef only if CUDA doesn't provide it +using __nv_bfloat162 = __half2; +#endif +#endif + +inline __device__ float gelu(const float x) +{ + constexpr float sqrt_param = 0.79788456080286535587989211986876f; + constexpr float mul_param = 0.044715; + return x * 0.5f * (1.0f + tanhf(sqrt_param * (x + mul_param * x * x * x))); +} + +/* +In-place gelu(biasAdd(x)) for channels last +*/ +template +__global__ void fused_bias_gelu(T* input, const T* bias, int total_count, int intermediate_size) +{ + // Input restriction: intermediate_size % vals_per_access == 0 + constexpr int granularity = 16; + constexpr int values_per_access = granularity / sizeof(T); + const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * values_per_access; + + if (offset < total_count) { + T data[values_per_access]; + T data_bias[values_per_access]; + mem_access::load_global(data, input + offset); + mem_access::load_global( + data_bias, bias + (offset % intermediate_size), bias != nullptr); + +#pragma unroll + for (int i = 0; i < values_per_access; i++) { + float data_f = conversion::to(data[i]); + float bias_f = conversion::to(data_bias[i]); + data[i] = conversion::to(gelu(data_f + bias_f)); + } + + mem_access::store_global(input + offset, data); + } +} + +template +void launch_bias_gelu(T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream) +{ + constexpr int threads = 1024; + constexpr int granularity = 16; + + const int total_count = batch_size * intermediate_size; + const int elems_per_block = threads * (granularity / sizeof(T)); + dim3 block_dims(threads); + dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); + + fused_bias_gelu<<>>( + input, bias, total_count, intermediate_size); +} + +#define INSTANTIATE_LAUNCH_BIAS_GELU(T) \ + template void launch_bias_gelu(T*, const T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_GELU(float) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_BIAS_GELU(__nv_bfloat16) +#endif +INSTANTIATE_LAUNCH_BIAS_GELU(__half) + +/* +In-place channels-last bias add +*/ +template +__global__ void fused_bias_add(T* input, const T* bias, int total_count, int intermediate_size) +{ + // Input restriction: intermediate_size % vals_per_access == 0 + constexpr int granularity = 16; + constexpr int values_per_access = granularity / sizeof(T); + const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * values_per_access; + + if (offset < total_count) { + T data[values_per_access]; + T data_bias[values_per_access]; + mem_access::load_global(data, input + offset); + mem_access::load_global( + data_bias, bias + (offset % intermediate_size), bias != nullptr); + +#pragma unroll + for (int i = 0; i < values_per_access; i++) { + float data_f = conversion::to(data[i]); + float bias_f = conversion::to(data_bias[i]); + data[i] = conversion::to(data_f + bias_f); + } + + mem_access::store_global(input + offset, data); + } +} + +template +void launch_bias_add(T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream) +{ + constexpr int threads = 1024; + constexpr int granularity = 16; + + const int total_count = batch_size * intermediate_size; + const int elems_per_block = threads * (granularity / sizeof(T)); + dim3 block_dims(threads); + dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); + + fused_bias_add<<>>( + input, bias, total_count, intermediate_size); +} + +#define INSTANTIATE_LAUNCH_BIAS_ADD(T) \ + template void launch_bias_add(T*, const T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_ADD(float) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_BIAS_ADD(__nv_bfloat16) +#endif +INSTANTIATE_LAUNCH_BIAS_ADD(__half) + +__global__ void fused_bias_residual(float* residual, + const float* hidden_state, + const float* attn, + const float* bias, + const float* attn_bias, + const int total_count, + const int intermediate_size, + const float mp_scale, + const bool preln) +{ + float4* res_fl4_ptr = reinterpret_cast(residual); + const float4* hs_fl4_ptr = reinterpret_cast(hidden_state); + const float4* attn_fl4_ptr = reinterpret_cast(attn); + const float4* bias_fl4_ptr = reinterpret_cast(bias); + const float4* attn_bias_fl4_ptr = reinterpret_cast(attn_bias); + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + + if (offset < total_count) { + float4 res_fl4 = res_fl4_ptr[offset]; + const float4 hs_fl4 = hs_fl4_ptr[offset]; + const float4 attn_fl4 = attn_fl4_ptr[offset]; + const float4 bias_fl4 = bias_fl4_ptr[offset % intermediate_size]; + const float4 attn_bias_fl4 = attn_bias_fl4_ptr[offset % intermediate_size]; + if (preln) { + // residual = (residual + attention + bias + attention_bias) * + // mp_scale + hidden_state + res_fl4.x = + (res_fl4.x + attn_fl4.x + bias_fl4.x + attn_bias_fl4.x) * mp_scale + (hs_fl4.x); + res_fl4.y = + (res_fl4.y + attn_fl4.y + bias_fl4.y + attn_bias_fl4.y) * mp_scale + (hs_fl4.y); + res_fl4.z = + (res_fl4.z + attn_fl4.z + bias_fl4.z + attn_bias_fl4.z) * mp_scale + (hs_fl4.z); + res_fl4.w = + (res_fl4.w + attn_fl4.w + bias_fl4.w + attn_bias_fl4.w) * mp_scale + (hs_fl4.w); + } else { + // residual += hidden_state + bias + res_fl4.x = res_fl4.x + hs_fl4.x + bias_fl4.x; + res_fl4.y = res_fl4.y + hs_fl4.y + bias_fl4.y; + res_fl4.z = res_fl4.z + hs_fl4.z + bias_fl4.z; + res_fl4.w = res_fl4.w + hs_fl4.w + bias_fl4.w; + } + res_fl4_ptr[offset] = res_fl4; + } +} + +template +__global__ void fused_bias_residual(T* residual, + const T* hidden_state, + const T* attn, + const T* bias, + const T* attn_bias, + const int total_count, + const int intermediate_size, + const float mp_scale, + const bool preln) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + float2* res_fl2_ptr = reinterpret_cast(residual); + const float2* hs_fl2_ptr = reinterpret_cast(hidden_state); + const float2* attn_fl2_ptr = reinterpret_cast(attn); + const float2* bias_fl2_ptr = reinterpret_cast(bias); + const float2* attn_bias_fl2_ptr = reinterpret_cast(attn_bias); + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + + if (offset < total_count) { + float2 res_fl2 = res_fl2_ptr[offset]; + const float2 hs_fl2 = hs_fl2_ptr[offset]; + const float2 attn_fl2 = attn_fl2_ptr[offset]; + const float2 bias_fl2 = bias_fl2_ptr[offset % intermediate_size]; + const float2 attn_bias_fl2 = attn_bias_fl2_ptr[offset % intermediate_size]; + + T2* res_half2 = reinterpret_cast(&res_fl2); + const T2* hs_half2 = reinterpret_cast(&hs_fl2); + const T2* attn_half2 = reinterpret_cast(&attn_fl2); + const T2* bias_half2 = reinterpret_cast(&bias_fl2); + const T2* attn_bias_half2 = reinterpret_cast(&attn_bias_fl2); + + float2 res_low = conversion::to(res_half2[0]); + float2 res_high = conversion::to(res_half2[1]); + + const float2 hs_low = conversion::to(hs_half2[0]); + const float2 hs_high = conversion::to(hs_half2[1]); + + const float2 attn_low = conversion::to(attn_half2[0]); + const float2 attn_high = conversion::to(attn_half2[1]); + + const float2 bias_low = conversion::to(bias_half2[0]); + const float2 bias_high = conversion::to(bias_half2[1]); + + const float2 attn_bias_low = conversion::to(attn_bias_half2[0]); + const float2 attn_bias_high = conversion::to(attn_bias_half2[1]); + + if (preln) { + // residual = (residual + attention + bias + attention_bias) * + // mp_scale + hidden_state + res_low.x = + (res_low.x + attn_low.x + bias_low.x + attn_bias_low.x) * mp_scale + hs_low.x; + res_low.y = + (res_low.y + attn_low.y + bias_low.y + attn_bias_low.y) * mp_scale + hs_low.y; + res_high.x = + (res_high.x + attn_high.x + bias_high.x + attn_bias_high.x) * mp_scale + hs_high.x; + res_high.y = + (res_high.y + attn_high.y + bias_high.y + attn_bias_high.y) * mp_scale + hs_high.y; + } else { + // residual += hidden_state + bias + res_low.x = (res_low.x + hs_low.x + bias_low.x); + res_low.y = (res_low.y + hs_low.y + bias_low.y); + res_high.x = (res_high.x + hs_high.x + bias_high.x); + res_high.y = (res_high.y + hs_high.y + bias_high.y); + } + res_half2[0] = conversion::to(res_low); + res_half2[1] = conversion::to(res_high); + + res_fl2_ptr[offset] = res_fl2; + } +} + +template +void launch_bias_residual(T* residual, + T* hidden_state, + T* attn, + T* bias, + T* attn_bias, + int batch, + int hidden_dim, + int mp_size, + bool preln, + cudaStream_t stream) +{ + int total_count = batch * hidden_dim / 4; + dim3 block_dims(1024); + dim3 grid_dims((total_count - 1) / 1024 + 1); // (batch_size); + + fused_bias_residual<<>>(residual, + hidden_state, + attn, + bias, + attn_bias, + total_count, + hidden_dim / 4, + 1.0 / mp_size, + preln); +} + +#define INSTANTIATE_LAUNCH_BIAS_RESIDUAL(T) \ + template void launch_bias_residual(T*, T*, T*, T*, T*, int, int, int, bool, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_RESIDUAL(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_BIAS_RESIDUAL(__nv_bfloat16); +#endif +INSTANTIATE_LAUNCH_BIAS_RESIDUAL(__half); + +__global__ void gptj_residual_add(float* residual, + const float* hidden_state, + const float* attn, + const float* bias, + const float* attn_bias, + const int total_count, + const int intermediate_size, + const float mp_scale) +{ + float4* res_fl4_ptr = reinterpret_cast(residual); + const float4* hs_fl4_ptr = reinterpret_cast(hidden_state); + const float4* attn_fl4_ptr = reinterpret_cast(attn); + const float4* bias_fl4_ptr = reinterpret_cast(bias); + const float4* attn_bias_fl4_ptr = reinterpret_cast(attn_bias); + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + + if (offset < total_count) { + float4 res_fl4 = res_fl4_ptr[offset]; + const float4 hs_fl4 = hs_fl4_ptr[offset]; + const float4 attn_fl4 = attn_fl4_ptr[offset]; + const float4 bias_fl4 = bias_fl4_ptr[offset % intermediate_size]; + + if (attn_bias) { + float4 attn_bias_fl4 = attn_bias_fl4_ptr[offset % intermediate_size]; + // residual += attention_bias + res_fl4.x += attn_bias_fl4.x; + res_fl4.y += attn_bias_fl4.y; + res_fl4.z += attn_bias_fl4.z; + res_fl4.w += attn_bias_fl4.w; + } + // residual = hidden_state + attention + (residual + bias) * mp_scale + res_fl4.x = hs_fl4.x + attn_fl4.x + (res_fl4.x + bias_fl4.x) * mp_scale; + res_fl4.y = hs_fl4.y + attn_fl4.y + (res_fl4.y + bias_fl4.y) * mp_scale; + res_fl4.z = hs_fl4.z + attn_fl4.z + (res_fl4.z + bias_fl4.z) * mp_scale; + res_fl4.w = hs_fl4.w + attn_fl4.w + (res_fl4.w + bias_fl4.w) * mp_scale; + + res_fl4_ptr[offset] = res_fl4; + } +} + +template +__global__ void gptj_residual_add(T* residual, + const T* hidden_state, + const T* attn, + const T* bias, + const T* attn_bias, + const int total_count, + const int intermediate_size, + const float mp_scale) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + float2* res_fl2_ptr = reinterpret_cast(residual); + const float2* hs_fl2_ptr = reinterpret_cast(hidden_state); + const float2* attn_fl2_ptr = reinterpret_cast(attn); + const float2* bias_fl2_ptr = reinterpret_cast(bias); + const float2* attn_bias_fl2_ptr = reinterpret_cast(attn_bias); + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + + if (offset < total_count) { + float2 res_fl2 = res_fl2_ptr[offset]; + const float2 hs_fl2 = hs_fl2_ptr[offset]; + const float2 attn_fl2 = attn_fl2_ptr[offset]; + const float2 bias_fl2 = bias_fl2_ptr[offset % intermediate_size]; + + T2* res_half2 = reinterpret_cast(&res_fl2); + const T2* hs_half2 = reinterpret_cast(&hs_fl2); + const T2* attn_half2 = reinterpret_cast(&attn_fl2); + const T2* bias_half2 = reinterpret_cast(&bias_fl2); + + float2 res_low = conversion::to(res_half2[0]); + float2 res_high = conversion::to(res_half2[1]); + + const float2 hs_low = conversion::to(hs_half2[0]); + const float2 hs_high = conversion::to(hs_half2[1]); + + const float2 attn_low = conversion::to(attn_half2[0]); + const float2 attn_high = conversion::to(attn_half2[1]); + + const float2 bias_low = conversion::to(bias_half2[0]); + const float2 bias_high = conversion::to(bias_half2[1]); + + if (attn_bias) { + const float2 attn_bias_fl2 = attn_bias_fl2_ptr[offset % intermediate_size]; + const T2* attn_bias_half2 = reinterpret_cast(&attn_bias_fl2); + const float2 attn_bias_low = conversion::to(attn_bias_half2[0]); + const float2 attn_bias_high = conversion::to(attn_bias_half2[1]); + // residual += attention_bias + res_low.x += attn_bias_low.x; + res_low.y += attn_bias_low.y; + res_high.x += attn_bias_high.x; + res_high.y += attn_bias_high.y; + } + // residual = hidden_state + attention + (residual + bias) * mp_scale + res_low.x = attn_low.x + hs_low.x + (res_low.x + bias_low.x) * mp_scale; + res_low.y = attn_low.y + hs_low.y + (res_low.y + bias_low.y) * mp_scale; + res_high.x = attn_high.x + hs_high.x + (res_high.x + bias_high.x) * mp_scale; + res_high.y = attn_high.y + hs_high.y + (res_high.y + bias_high.y) * mp_scale; + + res_half2[0] = conversion::to(res_low); + res_half2[1] = conversion::to(res_high); + + res_fl2_ptr[offset] = res_fl2; + } +} + +template +void launch_gptj_residual_add(T* residual, + T* hidden_state, + T* attn, + T* bias, + T* attn_bias, + int hidden_dim, + int batch, + int mp_size, + cudaStream_t stream) +{ + int total_count = batch * hidden_dim / 4; + dim3 block_dims(1024); + dim3 grid_dims((total_count - 1) / 1024 + 1); // (batch_size); + + gptj_residual_add<<>>( + residual, hidden_state, attn, bias, attn_bias, total_count, hidden_dim / 4, 1.0 / mp_size); +} + +#define INSTANTIATE_GPT_RES_ADD(T) \ + template void launch_gptj_residual_add(T*, T*, T*, T*, T*, int, int, int, cudaStream_t); + +INSTANTIATE_GPT_RES_ADD(float); +INSTANTIATE_GPT_RES_ADD(__half); +#ifdef BF16_AVAILABLE +INSTANTIATE_GPT_RES_ADD(__nv_bfloat16); +#endif + +template +__global__ void moe_res_matmul(T* residual, T* coef, T* mlp_out, int seq_len, int hidden_dim) +{ + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(T); + + T* residual_seq = residual + blockIdx.x * hidden_dim; + T* mlp_out_seq = mlp_out + blockIdx.x * hidden_dim; + + for (unsigned tid = threadIdx.x * vals_per_access; tid < hidden_dim; + tid += blockDim.x * vals_per_access) { + T mlp[vals_per_access]; + T res[vals_per_access]; + T coef1[vals_per_access]; + T coef2[vals_per_access]; + + mem_access::load_global(mlp, mlp_out_seq + tid); + mem_access::load_global(res, residual_seq + tid); + mem_access::load_global(coef1, coef + tid); + mem_access::load_global(coef2, coef + tid + hidden_dim); + +#pragma unroll + for (int idx = 0; idx < vals_per_access; idx++) { + mlp[idx] = mlp[idx] * coef2[idx] + res[idx] * coef1[idx]; + } + + mem_access::store_global(mlp_out_seq + tid, mlp); + } +} + +template +void launch_moe_res_matmul(T* residual, + T* coef, + T* mlp_out, + int seq_len, + int hidden_dim, + cudaStream_t stream) +{ + dim3 grid_dim(seq_len); + dim3 block_dim(1024); + moe_res_matmul<<>>( + residual, coef, mlp_out, seq_len, hidden_dim); +} + +#define INSTANTIATE_LAUNCH_MOE_RES_MATMUL(T) \ + template void launch_moe_res_matmul(T*, T*, T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_MOE_RES_MATMUL(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_MOE_RES_MATMUL(__nv_bfloat16); +#endif +INSTANTIATE_LAUNCH_MOE_RES_MATMUL(__half); + +template +__global__ void pad_data_kernel(T* padded_output, T* output, int head_size, int padded_head_size) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + float4* padded_output_cast = reinterpret_cast(padded_output); + float4* output_cast = reinterpret_cast(output); + int bid = blockIdx.x * (blockDim.y) + threadIdx.y; + int idx = threadIdx.x; + padded_output_cast += (bid * padded_head_size); + output_cast += (bid * head_size); + float4 ZERO; + const T2 zero_h = conversion::to(0.f); + T2* ZERO_h = reinterpret_cast(&ZERO); +#pragma unroll + for (int i = 0; i < 4; i++) ZERO_h[i] = zero_h; + if (idx < head_size) + padded_output_cast[idx] = output_cast[idx]; + else + padded_output_cast[idx] = ZERO; +} + +__global__ void pad_data_kernel(float* padded_output, + float* output, + int head_size, + int padded_head_size) +{ +} + +template +void pad_data(T* padded_output, + T* output, + int bsz, + int head_size, + int padded_head_size, + cudaStream_t stream) +{ + dim3 grid_dim((bsz - 1) / 16 + 1); + dim3 block_dim(padded_head_size / 8, 16); + pad_data_kernel<<>>( + padded_output, output, head_size / 8, padded_head_size / 8); +} + +#define INSTANTIATE_PAD_DATA(T) template void pad_data(T*, T*, int, int, int, cudaStream_t stream); + +INSTANTIATE_PAD_DATA(float); +INSTANTIATE_PAD_DATA(__half); +#ifdef BF16_AVAILABLE +INSTANTIATE_PAD_DATA(__nv_bfloat16); +#endif + +template +__global__ void pad_head_seq_kernel(T* padded_output, + T* output, + int seq_len, + int padded_seq_len, + int head_size, + int padded_head_size) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + float4* padded_output_cast = reinterpret_cast(padded_output); + float4* output_cast = reinterpret_cast(output); + int bsz = blockIdx.x; + int bid = blockIdx.y * (blockDim.y) + threadIdx.y; + int idx = threadIdx.x; + padded_output_cast += (bsz * padded_seq_len + bid) * padded_head_size; + output_cast += (bsz * seq_len + bid) * head_size; + float4 ZERO; + const T2 zero_h = conversion::to(0.f); + T2* ZERO_h = reinterpret_cast(&ZERO); +#pragma unroll + for (int i = 0; i < 4; i++) ZERO_h[i] = zero_h; + + if (idx < head_size && bid < seq_len) + padded_output_cast[idx] = output_cast[idx]; + else + padded_output_cast[idx] = ZERO; +} + +__global__ void pad_head_seq_kernel(float* padded_output, + float* output, + int seq_len, + int padded_seq_len, + int head_size, + int padded_head_size) +{ +} + +template +void pad_head_seq(T* padded_output, + T* output, + int bsz, + int seq_len, + int padded_seq_len, + int head_size, + int padded_head_size, + cudaStream_t stream) +{ + dim3 grid_dim(bsz, padded_seq_len / 16); + dim3 block_dim(padded_head_size / 8, 16); + pad_head_seq_kernel<<>>( + padded_output, output, seq_len, padded_seq_len, head_size / 8, padded_head_size / 8); +} + +#define INSTANTIATE_PAD_HEAD_SEQ(T) \ + template void pad_head_seq(T*, T*, int, int, int, int, int, cudaStream_t); + +INSTANTIATE_PAD_HEAD_SEQ(__half); +#ifdef BF16_AVAILABLE +INSTANTIATE_PAD_HEAD_SEQ(__nv_bfloat16); +#endif +INSTANTIATE_PAD_HEAD_SEQ(float); + +// TODO(cmikeh2): evaluate different GeLU performance +__device__ __forceinline__ float old_gelu(float val) +{ + // 1 / sqrt(2) + constexpr float rsqrt_2 = 0.707106769084930419922; + return val * 0.5f * (1.0f + erff(val * rsqrt_2)); +} + +namespace fused_geglu { +constexpr int threads = 256; +constexpr int steps = 2; +constexpr int granularity = 16; +} // namespace fused_geglu + +__device__ __forceinline__ float silu(float val) { return val / (1.0f + expf(-val)); } + +template +__global__ void fused_gate_activation(T* output, + const T* activation, + const T* bias, + int base_channels, + int output_stride, + int total_elems) +{ + constexpr int T_per_access = fused_geglu::granularity / sizeof(T); + constexpr int T_per_step = T_per_access * fused_geglu::threads; + constexpr int T_per_block = T_per_step * fused_geglu::steps; + + const int id = blockIdx.x * T_per_block + threadIdx.x * T_per_access; + +#pragma unroll + for (int i = 0; i < fused_geglu::steps; i++) { + T activation_buffer_1[T_per_access]; + T activation_buffer_2[T_per_access]; + T bias_buffer_1[T_per_access]; + T bias_buffer_2[T_per_access]; + + const int iter_id = id + T_per_step * i; + if (iter_id < total_elems) { + const int channel_id = iter_id % base_channels; + const int seq_id = iter_id / base_channels; + const int seq_offset = seq_id * base_channels * 2; + + mem_access::load_global(activation_buffer_1, + activation + seq_offset + channel_id); + mem_access::load_global( + activation_buffer_2, activation + seq_offset + channel_id + base_channels); + mem_access::load_global( + bias_buffer_1, bias + channel_id, bias != nullptr); + mem_access::load_global( + bias_buffer_2, bias + channel_id + base_channels, bias != nullptr); + + // Since the GeLU is going to happen at float, might as well + // convert +#pragma unroll + for (int v = 0; v < T_per_access; v++) { + T hidden_state = activation_buffer_1[v] + bias_buffer_1[v]; + T pre_gate = activation_buffer_2[v] + bias_buffer_2[v]; + float pre_gate_f = conversion::to(pre_gate); + float gate_f = (useGelu) ? old_gelu(pre_gate_f) : silu(pre_gate_f); + T gate = conversion::to(gate_f); + activation_buffer_1[v] = hidden_state * gate; + } + + mem_access::store_global( + output + seq_id * output_stride + channel_id, activation_buffer_1); + } + } +} + +template +void launch_gated_activation(T* output, + const T* activation, + const T* bias, + int rows, + int output_stride, + int elems_per_row, + bool use_gelu, + cudaStream_t stream) +{ + /* + Fused bias GEGLU is a variant of the gated activation functions. + The input here is a matrix of [batch, seq_len, 2 * intermediate_dim] + where the second half of the channels act as GeLU gates for the first + half. + */ + + // Re-derive the above figures + constexpr int T_per_access = fused_geglu::granularity / sizeof(T); + constexpr int T_per_step = T_per_access * fused_geglu::threads; + constexpr int T_per_block = T_per_step * fused_geglu::steps; + + const int base_channels = elems_per_row / 2; + const int total_elems = base_channels * rows; + + dim3 block(fused_geglu::threads); + dim3 grid((total_elems + T_per_block - 1) / T_per_block); + + if (use_gelu) { + fused_gate_activation<<>>( + output, activation, bias, base_channels, output_stride, total_elems); + } else { + fused_gate_activation<<>>( + output, activation, bias, base_channels, output_stride, total_elems); + } +} + +#define INSTANTIATE_LAUNCH_GATED_ACTIVATION(T) \ + template void launch_gated_activation( \ + T*, const T*, const T*, int, int, int, bool, cudaStream_t); + +INSTANTIATE_LAUNCH_GATED_ACTIVATION(__half); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_GATED_ACTIVATION(__nv_bfloat16); +#endif +INSTANTIATE_LAUNCH_GATED_ACTIVATION(float); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu new file mode 100644 index 0000000000000000000000000000000000000000..e5e7e89c9d2069678bc445004053f8ad5f9809b1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu @@ -0,0 +1,503 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "inference_cuda_layers.h" +#include "memory_access_utils.h" +#include "reduction_utils.h" + +namespace cg = cooperative_groups; +using rop = reduce::ROpType; + +namespace ln { +constexpr int granularity = 16; +} // namespace ln + +/* +Primary layer norm implementation. Assumes elems_per_row % 8 +is equal to 0. + +Args: + output: buffer for output data + vals: buffer for input data + gamma: gain for normalization + beta: bias for normalization + epsilon: numeric stability + elems_per_row: number of elements each block will normalize +*/ +template +__global__ void fused_ln(T* output, + const T* vals, + const T* gamma, + const T* beta, + float epsilon, + int elems_per_row) +{ + constexpr int T_per_load = ln::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // X-dimension of the block + const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + + (tb.thread_index().y * elems_per_row); + const int thread_offset = tb.thread_index().x * T_per_load; + const int base_offset = block_offset + thread_offset; + const int stride = blockDim.x * T_per_load; + + float sum = reduce::init(); + + const T* input_base = vals + base_offset; + + T local_buffer[unRoll * T_per_load]; + +#pragma unRoll + for (int i = 0; i < unRoll; i++) { + T* iteration_buffer = local_buffer + i * T_per_load; + + mem_access::load_global( + iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row); + +#pragma unRoll + for (int j = 0; j < T_per_load; j++) { + float vals_up_cast = conversion::to(iteration_buffer[j]); + sum = reduce::element(sum, vals_up_cast); + } + } + + reduce::partitioned_block(tb, warp, sum); + const float mean = sum / elems_per_row; + + float mean_diff = reduce::init(); + +#pragma unRoll + for (int i = 0; i < unRoll; i++) { +#pragma unRoll + for (int j = 0; j < T_per_load; j++) { + // Using a 0 value here skews the variance, have to if-guard + if (thread_offset + i * stride < elems_per_row) { + float diff = (conversion::to(local_buffer[i * T_per_load + j]) - mean); + mean_diff = reduce::element(mean_diff, diff * diff); + } + } + } + + reduce::partitioned_block(tb, warp, mean_diff); + const float variance = mean_diff / elems_per_row; + const float denom = __frsqrt_rn(variance + epsilon); + + // const T mean_compute = conversion::to(mean); + // const T denom_compute = conversion::to(denom); + + T* block_output = output + block_offset; + +#pragma unRoll + for (int i = 0; i < unRoll; i++) { + T* iteration_buffer = local_buffer + i * T_per_load; + const int iter_idx = i * stride + thread_offset; + const bool do_loads = iter_idx < elems_per_row; + + T gamma_local[T_per_load], beta_local[T_per_load]; + + mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); + mem_access::load_global(beta_local, beta + iter_idx, do_loads); + +#pragma unRoll + for (int j = 0; j < T_per_load; j++) { + float val = conversion::to(iteration_buffer[j]); + val = (val - mean) * denom; + val = + val * conversion::to(gamma_local[j]) + conversion::to(beta_local[j]); + iteration_buffer[j] = conversion::to(val); + } + + if (do_loads) { + mem_access::store_global(block_output + iter_idx, iteration_buffer); + } + } +} + +#define LAUNCH_FUSED_LN(unRollFactor, threadsPerGroup, maxThreads) \ + fused_ln \ + <<>>(output, vals, gamma, beta, epsilon, elems_per_row); + +template +void launch_fused_ln(T* output, + const T* vals, + const T* gamma, + const T* beta, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream) +{ + // 8 for __half, 4 for float + constexpr int T_per_load = ln::granularity / sizeof(T); + + constexpr int maxThreads = 256; + + // For Flaoat, unRoll 4, for __half, unRoll 2 + constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2; + + const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); + const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; + + const int groups_per_block_max = + is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1; + const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; + const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; + + dim3 block(threadsPerGroup, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threadsPerGroup * h_per_step; + const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; + + if (is_subblock_schedule) { + // <=128 + if (threadsPerGroup == 1) { + LAUNCH_FUSED_LN(1, 1, maxThreads); + } else if (threadsPerGroup == 2) { + LAUNCH_FUSED_LN(1, 2, maxThreads); + } else if (threadsPerGroup == 4) { + LAUNCH_FUSED_LN(1, 4, maxThreads); + } else if (threadsPerGroup == 8) { + LAUNCH_FUSED_LN(1, 8, maxThreads); + } else if (threadsPerGroup == 16) { + LAUNCH_FUSED_LN(1, 16, maxThreads); + } + } else if (external_unRoll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_FUSED_LN(1 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 2) { + // 4097 - 8192 elems + LAUNCH_FUSED_LN(2 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 3) { + // 8193 - 12288 elems + LAUNCH_FUSED_LN(3 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 4) { + // 12289 - 16384 elems + LAUNCH_FUSED_LN(4 * internal_unRoll, maxThreads, maxThreads); + } +} + +#define INSTANTIATE_FUSED_LN(T) \ + template void launch_fused_ln(T*, const T*, const T*, const T*, float, int, int, cudaStream_t); + +INSTANTIATE_FUSED_LN(__half); +#ifdef BF16_AVAILABLE +INSTANTIATE_FUSED_LN(__nv_bfloat16); +#endif +INSTANTIATE_FUSED_LN(float); + +/* +Fused resiual + bias + layer norm implementation. Assumes elems_per_row % 8 +is equal to 0. + +TODO(cmikeh2): Goal is to deprecate this implementation. The bias + residual +need to be fused into compute-bound producer operations. + +Args: + output: buffer for output data + res_output: output of residual addition + vals: buffer for input data + residual: residual data + bias: bias of of input data + gamma: gain for normalization + beta: bias for normalization + epsilon: numeric stability + elems_per_row: number of elements each block will normalize +Template arg: + StoreResidual: controls whether the residual calculation is stored + or not. When set to false, the input `res_output` is unused. +*/ +template +__global__ void fused_residual_ln(T* output, + T* res_output, + const T* vals, + const T* residual, + const T* bias, + const T* gamma, + const T* beta, + float epsilon, + int elems_per_row) +{ + constexpr int T_per_load = ln::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // X-dimension of the block + const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + + (tb.thread_index().y * elems_per_row); + const int thread_offset = tb.thread_index().x * T_per_load; + const int base_offset = block_offset + thread_offset; + const int stride = tb.size() * T_per_load; + + float sum = reduce::init(); + + const T* input_base = vals + base_offset; + const T* residual_base = residual + base_offset; + const T* bias_base = bias + thread_offset; + + T local_buffer[unRoll * T_per_load]; + + // Unlike a vanilla layernorm, since we're fusing the two adds as well + // an inner unRoll seems to be less valuable. If anything, a double unRoll + // makes the most sense if we find we are having performance issues. +#pragma unRoll + for (int i = 0; i < unRoll; i++) { + T* iteration_buffer = local_buffer + i * T_per_load; + T residual_buffer[T_per_load]; + T bias_buffer[T_per_load]; + + mem_access::load_global( + iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row); + mem_access::load_global(residual_buffer, + residual_base + i * stride, + thread_offset + i * stride < elems_per_row); + mem_access::load_global( + bias_buffer, bias_base + i * stride, thread_offset + i * stride < elems_per_row); + +#pragma unRoll + for (int j = 0; j < T_per_load; j++) { + float vals_up_cast = conversion::to(iteration_buffer[j]); + float res_up_cast = conversion::to(residual_buffer[j]); + float bias_up_cast = conversion::to(bias_buffer[j]); + vals_up_cast = vals_up_cast + bias_up_cast + res_up_cast; + sum = reduce::element(sum, vals_up_cast); + iteration_buffer[j] = conversion::to(vals_up_cast); + } + + if (preLnResidual && (thread_offset + i * stride < elems_per_row)) { + mem_access::store_global(res_output + base_offset + i * stride, + iteration_buffer); + } + } + + reduce::partitioned_block(tb, warp, sum); + const float mean = sum / elems_per_row; + + float mean_diff = reduce::init(); +#pragma unRoll + for (int i = 0; i < unRoll; i++) { +#pragma unRoll + for (int j = 0; j < T_per_load; j++) { + // Using a 0 value here skews the variance, have to if-guard + if (thread_offset + i * stride < elems_per_row) { + float diff = (conversion::to(local_buffer[i * T_per_load + j]) - mean); + mean_diff = reduce::element(mean_diff, diff * diff); + } + } + } + + reduce::partitioned_block(tb, warp, mean_diff); + const float variance = mean_diff / elems_per_row; + const float denom = __frsqrt_rn(variance + epsilon); + + T* block_output = output + block_offset; + +#pragma unRoll + for (int i = 0; i < unRoll; i++) { + T* iteration_buffer = local_buffer + i * T_per_load; + const int iter_idx = i * stride + thread_offset; + const bool do_loads = iter_idx < elems_per_row; + + T gamma_local[T_per_load], beta_local[T_per_load]; + + mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); + mem_access::load_global(beta_local, beta + iter_idx, do_loads); + +#pragma unRoll + for (int j = 0; j < T_per_load; j++) { + // iteration_buffer[j] = (iteration_buffer[j] - mean_compute) * denom_compute; + // iteration_buffer[j] = iteration_buffer[j] * gamma_local[j] + beta_local[j]; + float val = conversion::to(iteration_buffer[j]); + val = (val - mean) * denom; + val = + val * conversion::to(gamma_local[j]) + conversion::to(beta_local[j]); + iteration_buffer[j] = conversion::to(val); + } + + if (do_loads) { + mem_access::store_global(block_output + iter_idx, iteration_buffer); + } + } +} + +// TODO(cmikeh2): There's a bunch of redundancy here that needs to be removed/simplified. +#define LAUNCH_FUSED_RES_LN(unRollFactor, threadsPerGroup, maxThreads) \ + fused_residual_ln \ + <<>>( \ + output, nullptr, vals, residual, bias, gamma, beta, epsilon, elems_per_row); + +template +void launch_fused_residual_ln(T* output, + const T* vals, + const T* residual, + const T* bias, + const T* gamma, + const T* beta, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream) +{ + // 8 for __half, 4 for float + constexpr int T_per_load = ln::granularity / sizeof(T); + + constexpr int maxThreads = 256; + + // For Flaoat, unRoll 4, for __half, unRoll 2 + constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2; + + const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); + const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; + + const int groups_per_block_max = + is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1; + const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; + const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; + + dim3 block(threadsPerGroup, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threadsPerGroup * h_per_step; + const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; + + if (is_subblock_schedule) { + // <=128 + if (threadsPerGroup == 1) { + LAUNCH_FUSED_RES_LN(1, 1, maxThreads); + } else if (threadsPerGroup == 2) { + LAUNCH_FUSED_RES_LN(1, 2, maxThreads); + } else if (threadsPerGroup == 4) { + LAUNCH_FUSED_RES_LN(1, 4, maxThreads); + } else if (threadsPerGroup == 8) { + LAUNCH_FUSED_RES_LN(1, 8, maxThreads); + } else if (threadsPerGroup == 16) { + LAUNCH_FUSED_RES_LN(1, 16, maxThreads); + } + } else if (external_unRoll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_FUSED_RES_LN(1 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 2) { + // 4097 - 8192 elems + LAUNCH_FUSED_RES_LN(2 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 3) { + // 8193 - 12288 elems + LAUNCH_FUSED_RES_LN(3 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 4) { + // 12289 - 16384 elems + LAUNCH_FUSED_RES_LN(4 * internal_unRoll, maxThreads, maxThreads); + } +} + +#define LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(unRollFactor, threadsPerGroup, maxThreads) \ + fused_residual_ln \ + <<>>( \ + norm_output, res_output, vals, residual, bias, gamma, beta, epsilon, elems_per_row); + +template +void launch_fused_residual_ln_store_pre_ln_res(T* norm_output, + T* res_output, + const T* vals, + const T* residual, + const T* bias, + const T* gamma, + const T* beta, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream) +{ + // 8 for __half, 4 for float + constexpr int T_per_load = ln::granularity / sizeof(T); + + constexpr int maxThreads = 256; + + // For Flaoat, unRoll 4, for __half, unRoll 2 + constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2; + + const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); + const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; + + const int groups_per_block_max = + is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1; + const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; + const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; + + dim3 block(threadsPerGroup, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threadsPerGroup * h_per_step; + const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; + + if (is_subblock_schedule) { + // <=128 + if (threadsPerGroup == 1) { + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 1, maxThreads); + } else if (threadsPerGroup == 2) { + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 2, maxThreads); + } else if (threadsPerGroup == 4) { + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 4, maxThreads); + } else if (threadsPerGroup == 8) { + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 8, maxThreads); + } else if (threadsPerGroup == 16) { + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 16, maxThreads); + } + } else if (external_unRoll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 2) { + // 4097 - 8192 elems + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(2 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 3) { + // 8193 - 12288 elems + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(3 * internal_unRoll, maxThreads, maxThreads); + } else if (external_unRoll == 4) { + // 12289 - 16384 elems + LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(4 * internal_unRoll, maxThreads, maxThreads); + } +} + +#define INSTANTIATE_RES_LN(T) \ + template void launch_fused_residual_ln( \ + T*, const T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t); + +#define INSTANTIATE_PRE_LN_RES(T) \ + template void launch_fused_residual_ln_store_pre_ln_res( \ + T*, T*, const T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t); + +INSTANTIATE_RES_LN(__half); +INSTANTIATE_RES_LN(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_RES_LN(__nv_bfloat16); +#endif + +INSTANTIATE_PRE_LN_RES(__half); +INSTANTIATE_PRE_LN_RES(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_PRE_LN_RES(__nv_bfloat16); +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu new file mode 100644 index 0000000000000000000000000000000000000000..0301ff7770426994fc76cf3e6087586287c96c0f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu @@ -0,0 +1,74 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" + +namespace pwise { +constexpr int granularity = 16; +constexpr int unroll = 4; +constexpr int threads = 256; +} // namespace pwise + +template +__global__ void vector_add_kernel(T* out, const T* a, const T* b, float gamma, int num_elems) +{ + constexpr int T_per_access = pwise::granularity / sizeof(T); + + const int block_offset = blockIdx.x * pwise::threads * pwise::unroll * T_per_access; + const int thread_offset = threadIdx.x * T_per_access; + const int total_offset = block_offset + thread_offset; + constexpr int stride = pwise::threads * T_per_access; + +#pragma unroll + for (int i = 0; i < pwise::unroll; i++) { + T temp_buf_a[T_per_access], temp_buf_b[T_per_access]; + + const int iter_idx = total_offset + i * stride; + + mem_access::load_global(temp_buf_a, a + iter_idx, iter_idx < num_elems); + mem_access::load_global(temp_buf_b, b + iter_idx, iter_idx < num_elems); + +#pragma unroll + for (int j = 0; j < T_per_access; j++) { + float up_cast_a = conversion::to(temp_buf_a[j]); + float up_cast_b = conversion::to(temp_buf_b[j]); + temp_buf_a[j] = conversion::to((gamma * up_cast_a) + up_cast_b); + } + + if (iter_idx < num_elems) { + mem_access::store_global(out + iter_idx, temp_buf_a); + } + } +} + +template +void launch_vector_add(T* out, + const T* a, + const T* b, + float gamma, + int num_elems, + cudaStream_t stream) +{ + constexpr int T_per_access = pwise::granularity / sizeof(T); + constexpr int T_per_block = pwise::threads * T_per_access * pwise::unroll; + + dim3 block(pwise::threads); + dim3 grid((num_elems + T_per_block - 1) / T_per_block); + + vector_add_kernel<<>>(out, a, b, gamma, num_elems); +} + +#define INSTANTIATE_VECTOR_ADD(T) \ + template void launch_vector_add( \ + T * out, const T* a, const T* b, float gamma, int num_elems, cudaStream_t stream); + +INSTANTIATE_VECTOR_ADD(float) +INSTANTIATE_VECTOR_ADD(__half) +#ifdef BF16_AVAILABLE +INSTANTIATE_VECTOR_ADD(__nv_bfloat16) +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..19dbe73726f74bbcdcc62c6db1537a009c872b71 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp @@ -0,0 +1,2044 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "inference_context.h" +#include "inference_cublas_wrappers.h" +#include "inference_cuda_layers.h" + +std::array gemm_algos = std::array({99, 99, 99}); + +// NOTE: This activation function type enum should be always in sync +// with the python counterpart, otherwise the casting from python binding +// will be incorrect. +enum class ActivationFuncType { UNKNOWN = 0, GELU = 1, ReLU = 2, GATED_GELU = 3, GATED_SILU = 4 }; + +enum class NormType { UNKNOWN = 0, LayerNorm = 1, GroupNorm = 2, RMSNorm = 3 }; + +enum class TransformerType : uint8_t { UNKNOWN = 0, GPTType = 1, BERTType = 2 }; + +// NOTE: this is a temporary and dodgy solution to distinguish GPT and BERT style models +// based on the dimensions of the corresponding attention mask. +inline auto infer_transformer_type(at::Tensor& attn_mask) -> TransformerType +{ + auto attn_mask_num_dims = attn_mask.sizes().size(); + + if (attn_mask_num_dims > 2) { + return TransformerType::GPTType; + } else if (attn_mask_num_dims == 2) { + return TransformerType::BERTType; + } else { + return TransformerType::UNKNOWN; + } +} + +// infer stride of attention mask memory layout based on the model type. +inline auto get_attn_mask_stride(at::Tensor& attn_mask) -> int +{ + auto trnsfrmr_type = infer_transformer_type(attn_mask); + + if (trnsfrmr_type == TransformerType::GPTType) { + return attn_mask.size(2); + } else if (trnsfrmr_type == TransformerType::BERTType) { + // Bert style models have always a mask stride of 1. + return 1; + } else if (trnsfrmr_type == TransformerType::UNKNOWN) { + return 0; + } + + // this is just to make the compiler happy. + return 0; +} + +template +at::Tensor ds_softmax(at::Tensor& attn_scores, + at::Tensor& attn_mask, + at::Tensor& alibi, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + bool async_op, + float layer_scale, + int head_offset, + int mp_size) +{ + auto attn_scores_c = attn_scores.contiguous(); + int bsz = attn_scores_c.size(0); + + int seq_len = attn_scores_c.size(1); + int len = attn_scores_c.sizes().size(); + if (len > 2) seq_len = attn_scores_c.size(2); + + int soft_len = attn_scores_c.size(2); + if (len > 3) soft_len = attn_scores_c.size(3); + + int heads = 1; + if (len > 1) heads = attn_scores_c.size(1); + + auto mask_stride = get_attn_mask_stride(attn_mask); + + launch_attn_softmax_v2((T*)attn_scores_c.data_ptr(), + (attn_mask.sizes().size() > 1 ? (T*)attn_mask.data_ptr() : nullptr), + (alibi.sizes().size() > 1 ? (T*)alibi.data_ptr() : nullptr), + layer_scale, + triangular, + recompute, + local_attention, + window_size, + bsz, + heads, + seq_len, + soft_len, + head_offset, + mask_stride, + mp_size, + InferenceContext::Instance().GetCurrentStream(async_op)); + + return attn_scores_c; +} + +template +void allocate_workspace(unsigned hidden_dim, + unsigned num_heads, + unsigned prompt_length, + unsigned batch_size, + unsigned num_layers, + unsigned mp_size = 1, + bool external_cache = false, + unsigned rank = 0, + unsigned max_out_tokens = 1024, + unsigned min_out_tokens = 1) +{ + InferenceContext::Instance().GenWorkSpace(num_layers, + num_heads, + batch_size, + prompt_length, + hidden_dim, + mp_size, + external_cache, + sizeof(T), + rank, + max_out_tokens, + min_out_tokens); +} + +template +at::Tensor einsum_sec_sm_ecm(at::Tensor& Q, at::Tensor& W) +{ + auto options = at::TensorOptions() + .dtype(Q.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + float alpha = 1; + float gemm_beta = 0.0; + + /* + // Reallocate memory if we received a new prompt + if (!workspace || input.size(1) != 1) { + allocate_workspace(W.size(1), InferenceContext::Instance().GetMaxTokenLength(), + Q.size(0), 1, head_size); workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + } + */ + + auto O = at::from_blob(workspace, {Q.size(1), Q.size(2), W.size(1)}, options); + unsigned m = W.size(1); + unsigned n = Q.size(1) * Q.size(2); + unsigned k = Q.size(0); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + CUBLAS_OP_N, + CUBLAS_OP_T, + m, + n, + k, + &alpha, + &gemm_beta, + (T*)W.data_ptr(), + (T*)Q.data_ptr(), + (T*)O.data_ptr(), +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + return O; +} + +template +void attention_unfused(at::Tensor& prev_key_cont, + at::Tensor& query_cont, + at::Tensor& attn_mask, + at::Tensor& prev_value_cont, + at::Tensor& output, + int& bsz, + int& seq_len, + int& soft_len, + int& heads, + float& norm_factor, + bool triangular, + bool recompute, + bool local_attention, + int window_size) +{ + auto options = at::TensorOptions() + .dtype(query_cont.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + float alpha = norm_factor; + float gemm_beta = 0.0; + auto attn_score = at::empty({bsz, heads, seq_len, soft_len}, options); + int k = prev_value_cont.size(2) / heads; + + auto mask_stride = get_attn_mask_stride(attn_mask); + + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), + soft_len, + seq_len, + k, + &alpha, + &gemm_beta, + (T*)prev_key_cont.data_ptr(), + (T*)query_cont.data_ptr(), + (T*)attn_score.data_ptr(), + CUBLAS_OP_N, + CUBLAS_OP_N, + soft_len * k, + seq_len * k, + seq_len * soft_len, + bsz * heads, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + launch_attn_softmax_v2((T*)attn_score.data_ptr(), + (T*)(attn_mask.sizes().size() > 1 ? attn_mask.data_ptr() : nullptr), + (T*)nullptr, + 1.0, + triangular, + recompute, + local_attention, + window_size, + bsz, + heads, + seq_len, + soft_len, + 0, + mask_stride, + 1, + InferenceContext::Instance().GetCurrentStream(false)); + alpha = 1.0; + cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), + k, + seq_len, + soft_len, + &alpha, + &gemm_beta, + (T*)prev_value_cont.data_ptr(), + (T*)attn_score.data_ptr(), + (T*)output.data_ptr(), + CUBLAS_OP_N, + CUBLAS_OP_N, + soft_len * k, + seq_len * soft_len, + seq_len * k, + bsz * heads, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif +} + +template +std::vector ds_softmax_context1(at::Tensor& query, + at::Tensor& prev_key, + at::Tensor& new_key, + at::Tensor& attn_mask, + at::Tensor& prev_value, + at::Tensor& new_value, + int heads, + float norm_factor, + bool merging, + bool triangular, + bool local_attention, + int window_size, + bool no_masking) +{ + auto query_cont = query.contiguous(); + auto prev_key_cont = prev_key.contiguous(); + auto prev_value_cont = prev_value.contiguous(); + + int new_size = (new_value.sizes().size() > 1 ? new_value.size(1) : 0); + + // Attn_Score [ batch Head Sequence-length Softmax-length] + + int bsz = query_cont.size(0); + int seq_len = query_cont.size(1); + int soft_len = prev_value.size(1); + + auto options = at::TensorOptions() + .dtype(query_cont.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output = + at::empty({prev_value.size(0), heads, seq_len, prev_value.size(2) / heads}, options); + attention_unfused(prev_key_cont, + query_cont, + attn_mask, //(no_masking ? nullptr : (T*)attn_mask.data_ptr()), + prev_value_cont, + output, + bsz, + seq_len, + soft_len, + heads, + norm_factor, + (triangular && (new_size == 0)), + (new_size == 0), + local_attention, + window_size); + + return {output, prev_key, prev_value}; +} + +template +void ds_softmax_internal(T* attn_scores, + at::Tensor& attn_mask, + at::Tensor& alibi, + float& layer_scale, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + int bsz, + int seq_len, + int soft_len, + int heads) +{ + auto mask_stride = get_attn_mask_stride(attn_mask); + + launch_attn_softmax_v2((T*)attn_scores, + (attn_mask.sizes().size() > 1 ? (T*)attn_mask.data_ptr() : nullptr), + (alibi.sizes().size() > 1 ? (T*)alibi.data_ptr() : nullptr), + layer_scale, + triangular, + recompute, + local_attention, + window_size, + bsz, + heads, + seq_len, + soft_len, + 0, + mask_stride, + 1, + at::cuda::getCurrentCUDAStream()); +} + +template +void attention_unfused(T* prev_key_cont, + T* query_cont, + at::Tensor& attn_mask, + T* prev_value_cont, + T* output, + unsigned& bsz, + int& k, + unsigned& seq_len, + unsigned& soft_len, + int& heads, + float& norm_factor, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + at::Tensor& alibi, + int layer_id) +{ + float layer_scale = alibi.sizes().size() > 1 ? std::max(1, layer_id) : 1.0; + float alpha = norm_factor * norm_factor / layer_scale; + float gemm_beta = 0.0; + T* workspace = (T*)InferenceContext::Instance().GetAttentionUnfusedWorkspace(); + + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), + soft_len, + seq_len, + k, + &alpha, + &gemm_beta, + (T*)prev_key_cont, + (T*)query_cont, + workspace, + CUBLAS_OP_T, + CUBLAS_OP_N, + InferenceContext::Instance().GetMaxTokenLength() * k, + seq_len * k, + seq_len * soft_len, + bsz * heads, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + ds_softmax_internal(workspace, + attn_mask, + alibi, + layer_scale, + triangular, + recompute, + local_attention, + window_size, + bsz, + seq_len, + soft_len, + heads); + alpha = 1.0; + cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), + k, + seq_len, + soft_len, + &alpha, + &gemm_beta, + (T*)prev_value_cont, + workspace, + (T*)output, + CUBLAS_OP_N, + CUBLAS_OP_N, + InferenceContext::Instance().GetMaxTokenLength() * k, + seq_len * soft_len, + seq_len * k, + bsz * heads, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif +} + +void reset_cache() { InferenceContext::Instance().reset_tokens(); } + +template +std::vector ds_softmax_context(at::Tensor& query_key_value, + at::Tensor& attn_mask, + int rotary_dim, + bool rotate_half, + bool rotate_every_two, + int heads, + int num_kv, + float norm_factor, + bool triangular, + bool local_attention, + int window_size, + bool no_masking, + unsigned layer_id, + unsigned num_layers, + at::Tensor& alibi, + float rope_theta, + bool is_prompt, + std::optional token_idx, + std::optional position_ids) +{ + unsigned bsz = query_key_value.size(0); + unsigned seq_len = query_key_value.size(1); + int k = query_key_value.size(2) / (heads + 2 * (num_kv > 0 ? num_kv : heads)); + unsigned hidden_dim = heads * k; + + is_prompt = (seq_len > 1); + + if (is_prompt) InferenceContext::Instance().reset_tokens(seq_len); + unsigned soft_len = InferenceContext::Instance().current_tokens(); + + auto options = at::TensorOptions() + .dtype(query_key_value.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + size_t buf_size = bsz * seq_len * hidden_dim; + auto output = torch::from_blob(workspace + 4 * buf_size, {bsz, seq_len, hidden_dim}, options); + + auto query_cont = workspace + 5 * buf_size; + size_t offset = + 10 * (hidden_dim * bsz * InferenceContext::Instance().GetMaxTokenLength()) + + layer_id * 2 * bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; + unsigned all_tokens = soft_len; + auto kv_cache = workspace + offset + (hidden_dim / heads) * (is_prompt ? 0 : soft_len - 1); + size_t value_offset = bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; + + T* temp_buf = (T*)output.data_ptr() + at::numel(output); + launch_bias_add_transform_0213((T*)query_cont, + kv_cache, + kv_cache + value_offset, + (T*)query_key_value.data_ptr(), + nullptr, + bsz, + seq_len, + (is_prompt ? 0 : soft_len - 1), + soft_len, + hidden_dim, + heads, + (num_kv > 0 ? num_kv : heads), + rotary_dim, + rotate_half, + rotate_every_two, + InferenceContext::Instance().GetCurrentStream(), + 3, + InferenceContext::Instance().GetMaxTokenLength(), + rope_theta); + if (rotary_dim > 0 && rotate_half) + launch_apply_rotary_pos_emb(query_cont, + kv_cache, + k, + seq_len, + rotary_dim, + (is_prompt ? 0 : soft_len - 1), + heads, + bsz, + rope_theta, + InferenceContext::Instance().GetCurrentStream(), + InferenceContext::Instance().GetMaxTokenLength()); + + attention_unfused(workspace + offset, + (T*)query_cont, + attn_mask, + workspace + offset + value_offset, + temp_buf, + bsz, + k, + seq_len, + all_tokens, + heads, + norm_factor, + (triangular && is_prompt), + is_prompt, + local_attention, + window_size, + alibi, + layer_id); + launch_transform4d_0213((T*)output.data_ptr(), + temp_buf, + bsz, + heads, + seq_len, + output.size(2), + InferenceContext::Instance().GetCurrentStream(false), + 1); + + if (layer_id == num_layers - 1) InferenceContext::Instance().advance_tokens(); + auto prev_key = torch::from_blob( + workspace + offset, + {bsz, heads, all_tokens, k}, + {hidden_dim * static_cast(InferenceContext::Instance().GetMaxTokenLength()), + k * static_cast(InferenceContext::Instance().GetMaxTokenLength()), + k, + 1}, + options); + + auto prev_value = torch::from_blob( + workspace + offset + value_offset, + {bsz, heads, all_tokens, k}, + {hidden_dim * static_cast(InferenceContext::Instance().GetMaxTokenLength()), + k * static_cast(InferenceContext::Instance().GetMaxTokenLength()), + k, + 1}, + options); + + return {output, prev_key, prev_value}; +} + +template +at::Tensor ds_bias_gelu(at::Tensor& input, at::Tensor& bias) +{ + auto input_cont = input.contiguous(); + + int bsz = input_cont.size(0) * input_cont.size(1); + int intermediate_size = input_cont.size(2); + + launch_bias_gelu((T*)input_cont.data_ptr(), + (T*)bias.data_ptr(), + intermediate_size, + bsz, + InferenceContext::Instance().GetCurrentStream()); + return input_cont; +} + +#define DISPATCH_GATED_ACT(T_TYPE, C_TYPE) \ + if (activation.options().dtype() == torch::T_TYPE) { \ + launch_gated_activation((C_TYPE*)output.data_ptr(), \ + (const C_TYPE*)activation.data_ptr(), \ + (const C_TYPE*)bias.data_ptr(), \ + rows, \ + out_channels, \ + channels, \ + activation_type == ActivationFuncType::GATED_GELU, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor ds_gated_activation(at::Tensor& activation, at::Tensor& bias, int actFun) +{ + /* + Used in FF of Stable diffusion + */ + + const ActivationFuncType activation_type = static_cast(actFun); + + assert(activation_type == ActivationFuncType::GATED_GELU || + activation_type == ActivationFuncType::GATED_SILU); + + const int batch_size = activation.size(0); + const int seq_len = activation.size(1); + const int channels = activation.size(2); + + const int rows = batch_size * seq_len; + // Dimensionality is cut in half + const int out_channels = channels / 2; + + auto output = at::empty({batch_size, seq_len, out_channels}, activation.options()); + + DISPATCH_GATED_ACT(kFloat, float); + DISPATCH_GATED_ACT(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_GATED_ACT(kBFloat16, __nv_bfloat16); +#endif + + return output; +} + +template +at::Tensor ds_bias_relu(at::Tensor& input, at::Tensor& bias) +{ + auto input_cont = input.contiguous(); + + int bsz = input_cont.size(0) * input_cont.size(1); + int intermediate_size = input_cont.size(2); + + launch_bias_relu((T*)input_cont.data_ptr(), + (T*)bias.data_ptr(), + intermediate_size, + bsz, + InferenceContext::Instance().GetCurrentStream()); + return input_cont; +} + +template +at::Tensor ds_bias_add(at::Tensor& input, at::Tensor& bias) +{ + auto input_cont = input.contiguous(); + + int bsz = input_cont.size(0) * input_cont.size(1); + int hidden_size = input_cont.size(2); + + launch_bias_add((T*)input_cont.data_ptr(), + (T*)bias.data_ptr(), + hidden_size, + bsz, + InferenceContext::Instance().GetCurrentStream()); + return input_cont; +} + +template +at::Tensor ds_bias_residual(at::Tensor& input, at::Tensor& residual, at::Tensor& bias) +{ + auto input_cont = input.contiguous(); + auto residual_cont = residual.contiguous(); + + int bsz = input_cont.size(0) * input_cont.size(1); + // launch_bias_residual((T*)input_cont.data_ptr(), + // (T*)residual_cont.data_ptr(), + // (T*)bias.data_ptr(), + // bsz, + // input_cont.size(2), + // (bias.size(0) > 1), + // InferenceContext::Instance().GetCurrentStream()); + return input_cont; +} + +#define DISPATCH_LAYER_NORM(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_fused_ln((C_TYPE*)output.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + (const C_TYPE*)beta.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor ds_layer_norm(at::Tensor& input, at::Tensor& gamma, at::Tensor& beta, float epsilon) +{ + const int rows = input.size(0) * input.size(1); + const int elems_per_row = input.size(2); + auto output = at::empty_like(input); + + DISPATCH_LAYER_NORM(kFloat, float); + DISPATCH_LAYER_NORM(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_LAYER_NORM(kBFloat16, __nv_bfloat16); +#endif + + return output; +} + +#define DISPATCH_RMS_NORM(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_rms_norm((C_TYPE*)output.data_ptr(), \ + (C_TYPE*)nullptr, \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)nullptr, \ + (const C_TYPE*)gamma.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor ds_rms_norm(at::Tensor& input, at::Tensor& gamma, float epsilon) +{ + // Get number of dims of tensor + int num_dims = input.dim(); + const int rows = (num_dims == 2) ? input.size(0) : input.size(0) * input.size(1); + const int elems_per_row = (num_dims == 2) ? input.size(1) : input.size(2); + + auto output = at::empty_like(input); + + DISPATCH_RMS_NORM(kFloat, float); + DISPATCH_RMS_NORM(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_RMS_NORM(kBFloat16, __nv_bfloat16); +#endif + + return output; +} + +#define DISPATCH_PRE_RMS_NORM(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_rms_norm((C_TYPE*)output.data_ptr(), \ + (C_TYPE*)res_out.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)residual.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +std::vector ds_pre_rms_norm(at::Tensor& input, + at::Tensor& residual, + at::Tensor& gamma, + float epsilon) +{ + // Get number of dims of tensor + int num_dims = input.dim(); + const int rows = (num_dims == 2) ? input.size(0) : input.size(0) * input.size(1); + const int elems_per_row = (num_dims == 2) ? input.size(1) : input.size(2); + + auto output = at::empty_like(input); + auto res_out = at::empty_like(residual); + + DISPATCH_PRE_RMS_NORM(kFloat, float); + DISPATCH_PRE_RMS_NORM(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_PRE_RMS_NORM(kBFloat16, __nv_bfloat16); +#endif + + return {output, res_out}; +} + +template +void ds_layer_norm_internal(T* workspace, + at::Tensor& input, + at::Tensor& gamma, + at::Tensor& beta, + float epsilon) +{ + int bsz = input.size(0) * input.size(1); + launch_fused_ln(workspace, + (const T*)input.data_ptr(), + (const T*)gamma.data_ptr(), + (const T*)beta.data_ptr(), + epsilon, + bsz, + input.size(2), + InferenceContext::Instance().GetCurrentStream()); +} + +#define DISPATCH_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_fused_residual_ln((C_TYPE*)output.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)residual.data_ptr(), \ + (const C_TYPE*)bias.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + (const C_TYPE*)beta.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +/* Currently only used in unit testing */ +at::Tensor ds_layer_norm_residual(at::Tensor& input, + at::Tensor& bias, + at::Tensor& residual, + at::Tensor& gamma, + at::Tensor& beta, + float epsilon) +{ + const int rows = input.size(0) * input.size(1); + const int elems_per_row = input.size(2); + auto output = at::empty_like(input); + + DISPATCH_LAYER_NORM_RESIDUAL(kFloat, float); + DISPATCH_LAYER_NORM_RESIDUAL(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16); +#endif + + return output; +} + +#define DISPATCH_PRE_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_fused_residual_ln_store_pre_ln_res( \ + (C_TYPE*)norm_output.data_ptr(), \ + (C_TYPE*)res_output.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)residual.data_ptr(), \ + (const C_TYPE*)bias.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + (const C_TYPE*)beta.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +/* Currently only used in unit testing */ +std::vector ds_layer_norm_residual_store_pre_ln_res(at::Tensor& input, + at::Tensor& bias, + at::Tensor& residual, + at::Tensor& gamma, + at::Tensor& beta, + float epsilon) +{ + const int rows = input.size(0) * input.size(1); + const int elems_per_row = input.size(2); + auto norm_output = at::empty_like(input); + auto res_output = at::empty_like(input); + + DISPATCH_PRE_LAYER_NORM_RESIDUAL(kFloat, float); + DISPATCH_PRE_LAYER_NORM_RESIDUAL(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_PRE_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16); +#endif + + return {norm_output, res_output}; +} + +template +void quantized_gemm(void* output, + T* input, + at::Tensor& weight, + at::Tensor& qscale, + int groups, + int bsz, + int hidden_size) +{ + // T* weight16 = (T*)InferenceContext::Instance().GetWorkSpace() + 12 * hidden_size * bsz; + + auto options = at::TensorOptions() + .dtype(at::kHalf) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto tmp = torch::empty(weight.sizes(), options); + T* weight16 = (T*)tmp.data_ptr(); + launch_dequantize(weight16, + (int8_t*)weight.data_ptr(), + (float*)qscale.data_ptr(), + weight.size(0), + weight.size(1), + groups, + InferenceContext::Instance().GetCurrentStream()); + + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + CUBLAS_OP_T, + CUBLAS_OP_N, + weight.size(0), + bsz, + weight.size(1), + &alpha, + &gemm_beta, + weight16, + (T*)input, + (T*)output, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif +} + +template +at::Tensor qkv_unfused_cublas(at::Tensor& output, + at::Tensor& input, + at::Tensor& weight, + at::Tensor& q_scale, + at::Tensor& bias, + at::Tensor& gamma, + at::Tensor& beta, + const float epsilon, + bool add_bias, + bool q_int8, + bool transposed_mode) +{ + int bsz = input.size(0) * input.size(1); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + workspace += (3 * bsz * input.size(2)); + ds_layer_norm_internal(workspace, input, gamma, beta, epsilon); + + if (q_int8) { + quantized_gemm( + output.data_ptr(), workspace, weight, q_scale, q_scale.size(0), bsz, input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight.size(transposed_mode ? 0 : 1), + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + workspace, + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + if (add_bias) + launch_bias_add((T*)output.data_ptr(), + (T*)bias.data_ptr(), + (transposed_mode || q_int8) ? weight.size(0) : weight.size(1), + bsz, + InferenceContext::Instance().GetCurrentStream()); + return torch::from_blob(workspace, input.sizes(), input.options()); +} + +template +std::vector ds_rms_qkv(at::Tensor& input, + at::Tensor& weight, + at::Tensor& q_scale, + at::Tensor& gamma, + const float epsilon, + bool q_int8, + bool transposed_mode) +{ + const int bsz = input.size(0) * input.size(1); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + T* rms_norm_ptr = workspace + (3 * bsz * input.size(2)); + int out_size = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); + + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto rms_norm = at::from_blob(rms_norm_ptr, input.sizes(), options); + auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); + + launch_rms_norm((T*)rms_norm.data_ptr(), + (T*)nullptr, + (const T*)input.data_ptr(), + (const T*)nullptr, + (const T*)gamma.data_ptr(), + epsilon, + bsz, + input.size(2), + InferenceContext::Instance().GetCurrentStream()); + + if (q_int8) { + quantized_gemm((T*)output.data_ptr(), + (T*)rms_norm.data_ptr(), + weight, + q_scale, + q_scale.size(0), + bsz, + input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight.size(transposed_mode ? 0 : 1), + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + (T*)rms_norm.data_ptr(), + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + + return {output, rms_norm}; +} + +template +std::vector ds_qkv_gemm(at::Tensor& input, + at::Tensor& weight, + at::Tensor& q_scale, + at::Tensor& bias, + at::Tensor& gamma, + at::Tensor& beta, + const float epsilon, + bool add_bias, + bool q_int8, + bool transposed_mode) +{ + int bsz = input.size(0) * input.size(1); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + int out_size = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); + + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); + auto inp_norm = qkv_unfused_cublas(output, + input, + weight, + q_scale, + bias, + gamma, + beta, + epsilon, + add_bias, + q_int8, + transposed_mode); + + return {output, inp_norm}; +} + +template +void quantized_gemm(at::Tensor& output, + at::Tensor& input, + at::Tensor& weight, + at::Tensor& qscale, + int groups, + int merge_count) +{ + int bsz = input.size(0) * input.size(1); + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto weight16 = at::empty({weight.size(0), weight.size(1)}, options); + + launch_dequantize((T*)weight16.data_ptr(), + (int8_t*)weight.data_ptr(), + (float*)qscale.data_ptr(), + weight.size(0), + weight.size(1), + groups, + merge_count, + InferenceContext::Instance().GetCurrentStream()); + + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + CUBLAS_OP_T, + CUBLAS_OP_N, + weight.size(0), + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight16.data_ptr(), + (T*)input.data_ptr(), + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif +} + +template +at::Tensor ds_linear_layer(at::Tensor& input, + at::Tensor& weight, + at::Tensor& bias, + bool add_bias, + bool do_flash_attn, + int num_heads, + bool transposed_mode, + float rope_theta) +{ + auto input_cont = input.contiguous(); + auto options = at::TensorOptions() + .dtype(input_cont.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + int head_size = input_cont.size(2) / num_heads; + int bsz = input.size(0) * input.size(1); + int out_size = transposed_mode ? weight.size(0) : weight.size(1); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); + + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight.size(transposed_mode ? 0 : 1), + bsz, + input_cont.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + (T*)input_cont.data_ptr(), + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + if (add_bias) + launch_bias_add((T*)output.data_ptr(), + (T*)bias.data_ptr(), + weight.size(transposed_mode ? 0 : 1), + bsz, + InferenceContext::Instance().GetCurrentStream()); + bool add_padding = (head_size % 32 != 0 && head_size < 64) || (head_size % 64 != 0); + if (do_flash_attn) { + if (add_padding) { + int padded_head_size = head_size < 32 ? 32 : (head_size < 64 ? 64 : 128); + auto padded_output = workspace + output.numel(); + auto final_output = + padded_output + (input.size(0) * input.size(1) * 3 * num_heads * padded_head_size); + pad_data(padded_output, + workspace, + 3 * bsz * num_heads, + head_size, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + + launch_bias_add_transform_0213( + final_output, + final_output + (input.size(0) * input.size(1) * num_heads * padded_head_size), + final_output + (input.size(0) * input.size(1) * 2 * num_heads * padded_head_size), + padded_output, + nullptr, + input.size(0), + input.size(1), + 0, + input.size(1), + (num_heads * padded_head_size), + num_heads, + -1, + -1, + false, + false, + InferenceContext::Instance().GetCurrentStream(), + 3, + input.size(1), + rope_theta); + return at::from_blob(final_output, + {3, input.size(0), num_heads, input.size(1), padded_head_size}, + options); + // return at::from_blob(padded_output, {input.size(0) * input.size(1), 3, num_heads, + // padded_head_size}, options); + } else { + auto final_output = workspace + output.numel(); + launch_bias_add_transform_0213( + final_output, + final_output + (input.size(0) * input.size(1) * input_cont.size(2)), + final_output + (input.size(0) * input.size(1) * 2 * input_cont.size(2)), + workspace, + nullptr, + input.size(0), + input.size(1), + 0, + input.size(1), + input_cont.size(2), + num_heads, + -1, + -1, + false, + false, + InferenceContext::Instance().GetCurrentStream(), + 3, + input.size(1), + rope_theta); + return at::from_blob( + final_output, {3, input.size(0), num_heads, input.size(1), head_size}, options); + // return at::from_blob(workspace, {input.size(0) * input.size(1), 3, num_heads, + // head_size}, options); + } + + } else + return output; +} + +template +std::vector add_padding(at::Tensor& query, at::Tensor& key, at::Tensor& value) +{ + int head_size = query.size(3); + int padded_head_size = head_size < 32 ? 32 : (head_size < 64 ? 64 : 128); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + T* key_pad_ptr = workspace + padded_head_size * query.size(0) * query.size(1) * query.size(2); + T* value_pad_ptr = key_pad_ptr + padded_head_size * query.size(0) * query.size(1) * 128; + pad_head_seq(workspace, + (T*)query.data_ptr(), + query.size(0) * query.size(1), + query.size(2), + query.size(2), + head_size, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + pad_head_seq(key_pad_ptr, + (T*)key.data_ptr(), + query.size(0) * query.size(1), + key.size(2), + 128, + head_size, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + pad_head_seq(value_pad_ptr, + (T*)value.data_ptr(), + query.size(0) * query.size(1), + key.size(2), + 128, + head_size, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + return { + at::from_blob(workspace, + {query.size(0), query.size(1), query.size(2), padded_head_size}, + query.options()), + at::from_blob( + key_pad_ptr, {query.size(0), query.size(1), 128, padded_head_size}, query.options()), + at::from_blob( + value_pad_ptr, {query.size(0), query.size(1), 128, padded_head_size}, query.options())}; +} + +template +std::vector padd_add_transform(at::Tensor& query, + at::Tensor& key, + at::Tensor& value, + int heads, + bool add_padding) +{ + int head_size = query.size(2) / heads; + int key_value_length = add_padding ? 128 : key.size(1); + int padded_head_size = add_padding ? (head_size < 32 ? 32 : (head_size < 64 ? 64 : 128)) + : head_size; + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + T* key_pad_ptr = workspace + padded_head_size * query.size(0) * heads * query.size(1); + T* value_pad_ptr = key_pad_ptr + padded_head_size * query.size(0) * heads * key_value_length; + launch_pad_add_transform_0213(workspace, + (T*)query.data_ptr(), + query.size(0), + query.size(2), + query.size(1), + query.size(1), + heads, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + launch_pad_add_transform_0213(key_pad_ptr, + (T*)key.data_ptr(), + key.size(0), + key.size(2), + key.size(1), + key_value_length, + heads, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + launch_pad_add_transform_0213(value_pad_ptr, + (T*)value.data_ptr(), + value.size(0), + value.size(2), + value.size(1), + key_value_length, + heads, + padded_head_size, + InferenceContext::Instance().GetCurrentStream()); + return { + at::from_blob( + workspace, {query.size(0), heads, query.size(1), padded_head_size}, query.options()), + at::from_blob(key_pad_ptr, + {query.size(0), heads, key_value_length, padded_head_size}, + query.options()), + at::from_blob(value_pad_ptr, + {query.size(0), heads, key_value_length, padded_head_size}, + query.options())}; +} + +template +at::Tensor ds_vector_matmul(at::Tensor& input, + at::Tensor& weight, + bool async_op, + at::Tensor& q_scale, + bool q_int8, + bool transposed_mode) +{ + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + int out_size = (q_int8 || transposed_mode) ? weight.size(0) : weight.size(1); + int bsz = input.size(0) * input.size(1); + + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); + if (q_int8) { + quantized_gemm(output.data_ptr(), + (T*)input.data_ptr(), + weight, + q_scale, + q_scale.size(0), + bsz, + input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream(async_op)); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight.size(transposed_mode ? 0 : 1), + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + (T*)input.data_ptr(), + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + return output; +} + +template +at::Tensor ds_vector_matmul_int8(at::Tensor& input, + at::Tensor& weight, + at::Tensor& q_scale, + int groups, + int merge_count) +{ + auto input_cont = input.contiguous(); + auto options = at::TensorOptions() + .dtype(input_cont.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output = at::empty({input_cont.size(0), input_cont.size(1), weight.size(1)}, options); + + quantized_gemm(output, input_cont, weight, q_scale, groups, merge_count); + return output; +} + +template +at::Tensor mlp_unfused_cublas(at::Tensor& output, + at::Tensor& input, + at::Tensor& residual, + at::Tensor& input_bias, + at::Tensor& weight, + at::Tensor& weight1, + at::Tensor& bias, + at::Tensor& gamma, + at::Tensor& beta, + const float epsilon, + bool preLayerNorm, + bool mlp_after_attn, + at::Tensor& q_scale, + at::Tensor& q_scale1, + bool q_int8, + ActivationFuncType act_func_type, + bool transposed_mode) +{ + int bsz = input.size(0) * input.size(1); + T* inp_norm = (T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input) + + torch::numel(output); + T* intermediate = inp_norm + torch::numel(input); + + if (mlp_after_attn) { + launch_fused_residual_ln((T*)inp_norm, + (const T*)input.data_ptr(), + (const T*)residual.data_ptr(), + (const T*)input_bias.data_ptr(), + (const T*)gamma.data_ptr(), + (const T*)beta.data_ptr(), + epsilon, + bsz, + input.size(2), + InferenceContext::Instance().GetCurrentStream()); + } else { + ds_layer_norm_internal(inp_norm, input, gamma, beta, epsilon); + } + if (q_int8) { + quantized_gemm( + intermediate, inp_norm, weight, q_scale, q_scale.size(0), bsz, input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight.size(transposed_mode ? 0 : 1), + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + inp_norm, + intermediate, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + if (act_func_type == ActivationFuncType::GELU) { + launch_bias_gelu(intermediate, + (T*)bias.data_ptr(), + (transposed_mode || q_int8) ? weight.size(0) : weight.size(1), + bsz, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::ReLU) { + launch_bias_relu(intermediate, + (T*)bias.data_ptr(), + (transposed_mode || q_int8) ? weight.size(0) : weight.size(1), + bsz, + InferenceContext::Instance().GetCurrentStream()); + } + + if (q_int8) { + quantized_gemm(output.data_ptr(), + intermediate, + weight1, + q_scale1, + q_scale1.size(0), + bsz, + input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight1.size(transposed_mode ? 0 : 1), + bsz, + weight1.size(transposed_mode ? 1 : 0), + &alpha, + &gemm_beta, + (T*)weight1.data_ptr(), + intermediate, + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + + return torch::from_blob(inp_norm, input.sizes(), input.options()); +} + +template +std::vector ds_mlp_gemm(at::Tensor& input, + at::Tensor& residual, + at::Tensor& input_bias, + at::Tensor& weight_interm, + at::Tensor& weight_out, + at::Tensor& bias, + at::Tensor& gamma, + at::Tensor& beta, + const float epsilon, + bool preLayerNorm, + bool mlp_after_attn, + at::Tensor& q_scale, + at::Tensor& q_scale1, + bool q_int8, + int activation_type, + bool transposed_mode) +{ + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + int out_size = (q_int8 || transposed_mode) ? weight_out.size(0) : weight_out.size(1); + auto output = + at::from_blob((T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input), + {input.size(0), input.size(1), out_size}, + options); + int bsz = input.size(0) * input.size(1); + + auto act_func_type = static_cast(activation_type); + auto res_add = mlp_unfused_cublas(output, + mlp_after_attn ? input : residual, + residual, + input_bias, + weight_interm, + weight_out, + bias, + gamma, + beta, + epsilon, + preLayerNorm, + mlp_after_attn, + q_scale, + q_scale1, + q_int8, + act_func_type, + transposed_mode); + + return {output, res_add}; +} + +template +std::vector ds_rms_mlp_gemm(at::Tensor& input, + at::Tensor& residual, + at::Tensor& weight_interm, + at::Tensor& weight_out, + at::Tensor& gamma, + const float epsilon, + at::Tensor& q_scale, + at::Tensor& q_scale1, + bool q_int8, + int activation_type, + bool transposed_mode) +{ + const int bsz = input.size(0) * input.size(1); + const size_t input_neurons = input.size(2); + const size_t mlp_1_out_neurons = transposed_mode ? weight_interm.size(0) + : weight_interm.size(1); + const size_t mlp_2_in_neurons = transposed_mode ? weight_out.size(1) : weight_out.size(0); + + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + T* output_ptr = (T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input); + T* inp_norm_ptr = output_ptr + torch::numel(input); + T* intermediate_ptr = inp_norm_ptr + torch::numel(input); + + auto output = at::from_blob(output_ptr, input.sizes(), options); + auto inp_norm = at::from_blob(inp_norm_ptr, input.sizes(), options); + auto intermediate_gemm = + at::from_blob(intermediate_ptr, + {input.size(0), input.size(1), static_cast(mlp_1_out_neurons)}, + options); + + auto act_func_type = static_cast(activation_type); + + // RMS Norm, we'll update the residual in-place + launch_rms_norm((T*)inp_norm.data_ptr(), + (T*)residual.data_ptr(), + (const T*)input.data_ptr(), + (const T*)residual.data_ptr(), + (const T*)gamma.data_ptr(), + epsilon, + bsz, + input_neurons, + InferenceContext::Instance().GetCurrentStream()); + + if (q_int8) { + quantized_gemm(intermediate_ptr, + (T*)inp_norm.data_ptr(), + weight_interm, + q_scale, + q_scale.size(0), + bsz, + input_neurons); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + mlp_1_out_neurons, + bsz, + input_neurons, + &alpha, + &gemm_beta, + (T*)weight_interm.data_ptr(), + (T*)inp_norm.data_ptr(), + intermediate_ptr, +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + + if (act_func_type == ActivationFuncType::GELU) { + launch_bias_gelu(intermediate_ptr, + (T*)nullptr, + mlp_1_out_neurons, + bsz, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::ReLU) { + launch_bias_relu(intermediate_ptr, + (T*)nullptr, + mlp_1_out_neurons, + bsz, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::GATED_GELU) { + launch_gated_activation(intermediate_ptr, + (const T*)intermediate_ptr, + (const T*)nullptr, + bsz, + mlp_1_out_neurons, + mlp_1_out_neurons, + true, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::GATED_SILU) { + launch_gated_activation(intermediate_ptr, + (const T*)intermediate_ptr, + (const T*)nullptr, + bsz, + mlp_1_out_neurons, + mlp_1_out_neurons, + false, + InferenceContext::Instance().GetCurrentStream()); + } + + if (q_int8) { + quantized_gemm(output.data_ptr(), + intermediate_ptr, + weight_out, + q_scale1, + q_scale1.size(0), + bsz, + input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + input_neurons, + bsz, + mlp_2_in_neurons, + &alpha, + &gemm_beta, + (T*)weight_out.data_ptr(), + intermediate_ptr, + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard, +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP, +#endif + mlp_1_out_neurons); + } + + return {output, residual}; +} + +template +at::Tensor fused_gemm_gelu(at::Tensor& input, + at::Tensor& weight, + at::Tensor& weight_scale, + at::Tensor& bias, + at::Tensor& weight_out, + at::Tensor& weight_out_scale, + bool q_int8, + bool transposed_mode) +{ + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + int intm_dim = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); + + // auto output = at::from_blob((T*)InferenceContext::Instance().GetWorkSpace() + + // torch::numel(input), + // {input.size(0), input.size(1), out_size}, + // options); + // T* intermediate = (T*)input.data_ptr() + torch::numel(input); + auto intermediate = at::empty({input.size(0), input.size(1), intm_dim}, options); + + int bsz = input.size(0) * input.size(1); + + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + if (q_int8) { + quantized_gemm(intermediate.data_ptr(), + (T*)input.data_ptr(), + weight, + weight_scale, + weight_scale.size(0), + bsz, + input.size(2)); + } else { + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + intm_dim, + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + (T*)input.data_ptr(), + (T*)intermediate.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + launch_bias_gelu((T*)intermediate.data_ptr(), + (T*)bias.data_ptr(), + intm_dim, + bsz, + InferenceContext::Instance().GetCurrentStream()); + + int out_size = (transposed_mode || q_int8) ? weight_out.size(0) : weight_out.size(1); + auto output = at::empty({input.size(0), input.size(1), out_size}, options); + if (q_int8) { + quantized_gemm(output.data_ptr(), + (T*)intermediate.data_ptr(), + weight_out, + weight_out_scale, + weight_out_scale.size(0), + bsz, + input.size(2)); + } else { + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + out_size, + bsz, + intm_dim, + &alpha, + &gemm_beta, + (T*)weight_out.data_ptr(), + (T*)intermediate.data_ptr(), + (T*)output.data_ptr(), +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + // cudaEventRecord(InferenceContext::Instance().GetCompEvent(2), + // InferenceContext::Instance().GetCurrentStream(true)); + return output; +} + +template +at::Tensor& residual_add_bias(at::Tensor& hidden_state, + at::Tensor& residual, + const at::Tensor& attention_output, + const at::Tensor& attention_bias, + const at::Tensor& final_bias, + const int mp_size, + const bool mlp_after_attn, + const bool add_bias, + const bool preln) +{ + int bsz = residual.size(0) * residual.size(1); + int hidden_size = residual.size(2); + if (mlp_after_attn) + launch_bias_residual(static_cast(residual.data_ptr()), + static_cast(hidden_state.data_ptr()), + static_cast(attention_output.data_ptr()), + static_cast(final_bias.data_ptr()), + static_cast(attention_bias.data_ptr()), + bsz, + hidden_size, + mp_size, + preln, + InferenceContext::Instance().GetCurrentStream()); + else + launch_gptj_residual_add( + static_cast(residual.data_ptr()), + static_cast(hidden_state.data_ptr()), + static_cast(attention_output.data_ptr()), + static_cast(final_bias.data_ptr()), + static_cast((add_bias ? attention_bias.data_ptr() : nullptr)), + hidden_size, + bsz, + mp_size, + InferenceContext::Instance().GetCurrentStream()); + return residual; +} + +#define DISPATCH_VECTOR_ADD(T_TYPE, C_TYPE) \ + if (a.scalar_type() == at::k##T_TYPE) { \ + launch_vector_add((C_TYPE*)(a.data_ptr()), \ + (const C_TYPE*)(a.data_ptr()), \ + (const C_TYPE*)(b.data_ptr()), \ + gamma, \ + total_elems, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor& _vector_add(at::Tensor& a, at::Tensor& b, float gamma) +{ + const int total_elems = a.numel(); + + DISPATCH_VECTOR_ADD(Float, float) + DISPATCH_VECTOR_ADD(Half, __half) +#ifdef BF16_AVAILABLE + DISPATCH_VECTOR_ADD(BFloat16, __nv_bfloat16) +#endif + + return a; +} + +std::vector apply_rotary_pos_emb(at::Tensor& mixed_query, + at::Tensor& key_layer, + unsigned rotary_dim, + unsigned offset, + unsigned num_heads, + bool rotate_half, + float rope_theta) +{ + auto query_cont = mixed_query.contiguous(); + auto key_cont = key_layer.contiguous(); + + unsigned bsz = mixed_query.size(0); + unsigned head_size = mixed_query.size(2) / num_heads; + unsigned seq_len = mixed_query.size(1); + + if (mixed_query.scalar_type() == at::kFloat) + launch_apply_rotary_pos_emb((float*)query_cont.data_ptr(), + (float*)key_cont.data_ptr(), + head_size, + seq_len, + rotary_dim, + offset, + num_heads, + bsz, + rope_theta, + InferenceContext::Instance().GetCurrentStream(), + InferenceContext::Instance().GetMaxTokenLength()); + else + launch_apply_rotary_pos_emb<__half>((__half*)query_cont.data_ptr(), + (__half*)key_cont.data_ptr(), + head_size, + seq_len, + rotary_dim, + offset, + num_heads, + bsz, + rope_theta, + InferenceContext::Instance().GetCurrentStream(), + InferenceContext::Instance().GetMaxTokenLength()); + return {query_cont, key_cont}; +} + +#define DISPATCH_MOE_RESIDUAL(T_TYPE, C_TYPE) \ + if (moe_res.scalar_type() == torch::T_TYPE) { \ + launch_moe_res_matmul((C_TYPE*)moe_res.data_ptr(), \ + (C_TYPE*)coef.data_ptr(), \ + (C_TYPE*)output.data_ptr(), \ + M, \ + N, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor moe_res_matmul(at::Tensor& moe_res, at::Tensor& coef, at::Tensor& output) +{ + int M = moe_res.size(0) * moe_res.size(1); + int N = moe_res.size(2); + InferenceContext::Instance().SynchComm(); + + DISPATCH_MOE_RESIDUAL(kFloat, float) + DISPATCH_MOE_RESIDUAL(kHalf, __half) +#ifdef BF16_AVAILABLE + DISPATCH_MOE_RESIDUAL(kBFloat16, __nv_bfloat16) +#endif + + return output; +} + +void ds_release_workspace() { InferenceContext::Instance().release_workspace(); } + +bool ds_retake_workspace() { return InferenceContext::Instance().retake_workspace(); } + +template +at::Tensor ds_dequantize(at::Tensor& weight, at::Tensor& qscale, int groups) +{ + auto options = at::TensorOptions() + .dtype(torch::kFloat16) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto weight16 = at::empty({weight.size(0), weight.size(1)}, options); + + launch_dequantize((T*)weight16.data_ptr(), + (int8_t*)weight.data_ptr(), + (float*)qscale.data_ptr(), + weight.size(0), + weight.size(1), + groups, + InferenceContext::Instance().GetCurrentStream()); + + return weight16; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("softmax_context_int8", + &ds_softmax_context1<__half>, + "DeepSpeed attention with int8 (CUDA)"); + + // The following functions handle type dispatching internally + m.def("gated_activation", &ds_gated_activation, "DeepSpeed Bias GEGLU (CUDA)"); + m.def("layer_norm", &ds_layer_norm, "DeepSpeed layer norm (CUDA)"); + m.def( + "_layer_norm_residual", &ds_layer_norm_residual, "DeepSpeed layer norm + residual (CUDA)"); + m.def("layer_norm_residual_store_pre_ln_res", + &ds_layer_norm_residual_store_pre_ln_res, + "DeepSpeed layer norm + store pre Layernorm residual (CUDA)"); + m.def("rms_norm", &ds_rms_norm, "DeepSpeed rms norm (CUDA)"); + m.def("pre_rms_norm", &ds_pre_rms_norm, "DeepSpeed pre rms norm (CUDA)"); + m.def("_vector_add", &_vector_add, "DeepSpeed vector add (CUDA)"); + m.def("apply_rotary_pos_emb", &apply_rotary_pos_emb, "DeepSpeed mlp with fp16 (CUDA)"); + m.def("moe_res_matmul", &moe_res_matmul, "DeepSpeed moe residual matmul (CUDA)"); + m.def("reset_cache", &reset_cache, "Reset Cache for generation tasks"); + m.def("release_workspace", &ds_release_workspace, "DeepSpeed Release Workspace"); + m.def("retake_workspace", &ds_retake_workspace, "DeepSpeed Retake Workspace"); + + // The following functions are templated and need to be explicitly instantiated and bound + // to different python methods +#define DEF_OPS(_name, _dtype) \ + m.def("softmax_" #_name, &ds_softmax<_dtype>, "DeepSpeed SoftMax with " #_name " (CUDA)"); \ + m.def("softmax_context_" #_name, \ + &ds_softmax_context<_dtype>, \ + "DeepSpeed attention with " #_name " (CUDA)"); \ + m.def("bias_gelu_" #_name, &ds_bias_gelu<_dtype>, "DeepSpeed Gelu with " #_name " (CUDA)"); \ + m.def("bias_add_" #_name, &ds_bias_add<_dtype>, "DeepSpeed Bias Add with " #_name " (CUDA)"); \ + m.def("bias_relu_" #_name, &ds_bias_relu<_dtype>, "DeepSpeed ReLU with " #_name " (CUDA)"); \ + m.def("bias_residual_" #_name, \ + &ds_bias_residual<_dtype>, \ + "DeepSpeed residual-bias add with " #_name " (CUDA)"); \ + m.def("qkv_gemm_" #_name, &ds_qkv_gemm<_dtype>, "DeepSpeed qkv gemm with " #_name " (CUDA)"); \ + m.def("rms_qkv_gemm_" #_name, \ + &ds_rms_qkv<_dtype>, \ + "DeepSpeed rms qkv gemm with " #_name " (CUDA)"); \ + m.def("mlp_gemm_" #_name, &ds_mlp_gemm<_dtype>, "DeepSpeed mlp with " #_name " (CUDA)"); \ + m.def("rms_mlp_gemm_" #_name, \ + &ds_rms_mlp_gemm<_dtype>, \ + "DeepSpeed rms mlp gemm with " #_name " (CUDA)"); \ + m.def("vector_matmul_" #_name, \ + &ds_vector_matmul<_dtype>, \ + "DeepSpeed vector-MM with " #_name " (CUDA)"); \ + m.def("linear_layer_" #_name, \ + &ds_linear_layer<_dtype>, \ + "DeepSpeed linear_layer with " #_name " (CUDA)"); \ + m.def("fused_gemm_gelu_" #_name, \ + &fused_gemm_gelu<_dtype>, \ + "DeepSpeed mlp with " #_name " (CUDA)"); \ + m.def("residual_add_bias_" #_name, \ + &residual_add_bias<_dtype>, \ + "DeepSpeed residual add with " #_name " (CUDA)"); \ + m.def("einsum_sec_sm_ecm_" #_name, \ + &einsum_sec_sm_ecm<_dtype>, \ + "DeepSpeed vector-MM with " #_name " (CUDA)"); \ + m.def("add_padding_" #_name, \ + &add_padding<_dtype>, \ + "DeepSpeed residual add with " #_name " (CUDA)"); \ + m.def("pad_transform_" #_name, \ + &padd_add_transform<_dtype>, \ + "DeepSpeed residual add with " #_name " (CUDA)"); \ + m.def("allocate_workspace_" #_name, \ + &allocate_workspace<_dtype>, \ + "DeepSpeed memory allocation for GPT inference with " #_name " (CUDA)"); \ + m.def("dequantize_" #_name, \ + &ds_dequantize<_dtype>, \ + "DeepSpeed dequantize with " #_name " (CUDA)"); + + DEF_OPS(fp32, float); + DEF_OPS(fp16, __half); +#ifdef BF16_AVAILABLE + DEF_OPS(bf16, __nv_bfloat16); +#endif +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu new file mode 100644 index 0000000000000000000000000000000000000000..40926b776cf2f0bd3c46f44c6eda2ed3ec6bfce5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu @@ -0,0 +1,71 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "inference_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; +#define MAX_CAP 4 +#define MAX_SEQ 2048 + +inline __device__ float relu(const float x) { return x < 0 ? 0 : x; } + +/* +In-place relu(biasAdd(x)) for channels last +*/ +template +__global__ void fused_bias_relu(T* input, const T* bias, int total_count, int intermediate_size) +{ + // Input restriction: intermediate_size % vals_per_access == 0 + constexpr int granularity = 16; + constexpr int values_per_access = granularity / sizeof(T); + const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * values_per_access; + + if (offset < total_count) { + T data[values_per_access]; + T data_bias[values_per_access]; + mem_access::load_global(data, input + offset); + mem_access::load_global( + data_bias, bias + (offset % intermediate_size), bias != nullptr); + +#pragma unroll + for (int i = 0; i < values_per_access; i++) { + float data_f = conversion::to(data[i]); + float bias_f = conversion::to(data_bias[i]); + data[i] = conversion::to(relu(data_f + bias_f)); + } + + mem_access::store_global(input + offset, data); + } +} + +template +void launch_bias_relu(T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream) +{ + constexpr int threads = 1024; + constexpr int granularity = 16; + + const int total_count = batch_size * intermediate_size; + const int elems_per_block = threads * (granularity / sizeof(T)); + dim3 block_dims(threads); + dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); + + fused_bias_relu<<>>( + input, bias, total_count, intermediate_size); +} + +#define INSTANTIATE_LAUNCH_BIAS_RELU(T) \ + template void launch_bias_relu(T*, const T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_RELU(float) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_BIAS_RELU(__nv_bfloat16) +#endif +INSTANTIATE_LAUNCH_BIAS_RELU(__half) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu new file mode 100644 index 0000000000000000000000000000000000000000..5f72a4193752cc54b8ee4e1b375e19ef5c3d1214 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu @@ -0,0 +1,263 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "inference_cuda_layers.h" +#include "memory_access_utils.h" +#include "reduction_utils.h" + +namespace cg = cooperative_groups; +using rop = reduce::ROpType; + +namespace rms { +constexpr int granularity = 16; +} // namespace rms + +template +__global__ void rms_norm(T* output, const T* vals, const T* gamma, float epsilon, int elems_per_row) +{ + constexpr int T_per_load = rms::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // X-dimension of the block + const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + + (tb.thread_index().y * elems_per_row); + const int thread_offset = tb.thread_index().x * T_per_load; + const int base_offset = block_offset + thread_offset; + const int stride = blockDim.x * T_per_load; + + float var_sum = reduce::init(); + + const T* input_base = vals + base_offset; + + T local_buffer[UNROLL * T_per_load]; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + + mem_access::load_global(iteration_buffer, + input_base + (i * stride), + thread_offset + (i * stride) < elems_per_row); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + float up_cast = conversion::to(iteration_buffer[j]); + float sq_val = up_cast * up_cast; + var_sum = reduce::element(var_sum, sq_val); + } + } + + reduce::partitioned_block(tb, warp, var_sum); + const float var = var_sum / elems_per_row; + const T denom = conversion::to(__frsqrt_rn(var + epsilon)); + + T* block_output = output + block_offset; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + const int iter_idx = i * stride + thread_offset; + const bool do_loads = (iter_idx < elems_per_row); + + T gamma_local[T_per_load]; + + mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + iteration_buffer[j] *= denom; + iteration_buffer[j] *= gamma_local[j]; + } + + if (do_loads) { + mem_access::store_global(block_output + iter_idx, iteration_buffer); + } + } +} + +template +__global__ void pre_rms_norm(T* output, + T* res_out, + const T* vals, + const T* residual, + const T* gamma, + float epsilon, + int elems_per_row) +{ + constexpr int T_per_load = rms::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // X-dimension of the block + const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + + (tb.thread_index().y * elems_per_row); + const int thread_offset = tb.thread_index().x * T_per_load; + const int base_offset = block_offset + thread_offset; + const int stride = blockDim.x * T_per_load; + + float var_sum = reduce::init(); + + const T* input_base = vals + base_offset; + const T* residual_base = residual + base_offset; + T* res_output = res_out + base_offset; + + T local_buffer[UNROLL * T_per_load]; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + T residual_buffer[T_per_load]; + + const int iter_offset = i * stride + thread_offset; + const bool do_loads = (iter_offset < elems_per_row); + + mem_access::load_global( + iteration_buffer, input_base + (i * stride), do_loads); + mem_access::load_global( + residual_buffer, residual_base + (i * stride), do_loads); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + iteration_buffer[j] += residual_buffer[j]; + float vals_up_cast = conversion::to(iteration_buffer[j]); + + var_sum = reduce::element(var_sum, vals_up_cast * vals_up_cast); + } + + if (do_loads) { + mem_access::store_global(res_output + i * stride, iteration_buffer); + } + } + + reduce::partitioned_block(tb, warp, var_sum); + const float var = var_sum / elems_per_row; + const T denom = conversion::to(__frsqrt_rn(var + epsilon)); + + T* block_output = output + block_offset; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + const int iter_idx = i * stride + thread_offset; + const bool do_loads = (iter_idx < elems_per_row); + + T gamma_local[T_per_load]; + + mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + iteration_buffer[j] *= denom; + iteration_buffer[j] *= gamma_local[j]; + } + + if (do_loads) { + mem_access::store_global(block_output + iter_idx, iteration_buffer); + } + } +} + +#define LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + rms_norm \ + <<>>(norm_output, vals, gamma, epsilon, elems_per_row); + +#define LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + pre_rms_norm<<>>( \ + norm_output, res_output, vals, residual, gamma, epsilon, elems_per_row); + +#define LAUNCH_ALL_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + if (pre_norm) { \ + LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + } else { \ + LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + } + +template +void launch_rms_norm(T* norm_output, + T* res_output, + const T* vals, + const T* residual, + const T* gamma, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream) +{ + // 8 for __half, 4 for float + constexpr int T_per_load = rms::granularity / sizeof(T); + constexpr int maxThreads = 256; + constexpr int internalUnroll = sizeof(T) == 4 ? 4 : 2; + + const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internalUnroll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); + const int threads_per_group = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; + + const int groups_per_block_max = + is_subblock_schedule ? (maxThreads + threads_per_group - 1) / threads_per_group : 1; + const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; + const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; + + dim3 block(threads_per_group, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threads_per_group * h_per_step; + const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; + + bool pre_norm = (residual == nullptr) ? false : true; + + if (is_subblock_schedule) { + // <=128 + if (threads_per_group == 1) { + LAUNCH_ALL_RMS_NORM(1, 1, maxThreads); + } else if (threads_per_group == 2) { + LAUNCH_ALL_RMS_NORM(1, 2, maxThreads); + } else if (threads_per_group == 4) { + LAUNCH_ALL_RMS_NORM(1, 4, maxThreads); + } else if (threads_per_group == 8) { + LAUNCH_ALL_RMS_NORM(1, 8, maxThreads); + } else if (threads_per_group == 16) { + LAUNCH_ALL_RMS_NORM(1, 16, maxThreads); + } + } else if (external_unRoll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_ALL_RMS_NORM(1 * internalUnroll, maxThreads, maxThreads); + } else if (external_unRoll == 2) { + // 4097 - 8192 elems + LAUNCH_ALL_RMS_NORM(2 * internalUnroll, maxThreads, maxThreads); + } else if (external_unRoll == 3) { + // 8193 - 12288 elems + LAUNCH_ALL_RMS_NORM(3 * internalUnroll, maxThreads, maxThreads); + } else if (external_unRoll == 4) { + // 12289 - 16384 elems + LAUNCH_ALL_RMS_NORM(4 * internalUnroll, maxThreads, maxThreads); + } +} + +#define INSTANTIATE_LAUNCH_RMS_NORM(T) \ + template void launch_rms_norm(T * norm_output, \ + T * res_output, \ + const T* vals, \ + const T* residual, \ + const T* gamma, \ + float epsilon, \ + int rows, \ + int elems_per_row, \ + cudaStream_t stream); + +INSTANTIATE_LAUNCH_RMS_NORM(float) +INSTANTIATE_LAUNCH_RMS_NORM(__half) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_RMS_NORM(__nv_bfloat16) +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu new file mode 100644 index 0000000000000000000000000000000000000000..bb06cc149ef4fb29dae84a7b16173d1babfa5e50 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu @@ -0,0 +1,562 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "conversion_utils.h" +#include "inference_cuda_layers.h" + +#ifndef __HIP_PLATFORM_AMD__ +#include +#endif +#include +#include +#include + +#define MAX_REG_SIZE 8 + +#define minus_infinity -10000.0 + +void CheckCudaErrorAux(const char* file, unsigned line) +{ + cudaError_t err = cudaGetLastError(); + if (err == cudaSuccess) return; + std::cerr << cudaGetErrorString(err) << "(" << err << ") at " << file << ":" << line + << std::endl; + throw std::runtime_error("CUDA ERROR!!!\n"); +} + +#define CUDA_CHECK_ERROR() CheckCudaErrorAux(__FILE__, __LINE__) + +namespace cg = cooperative_groups; + +template +__global__ void attn_softmax_v2(T* vals, + T* mask, + T* alibi, + float layer_scale, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + int total_count, + int heads, + int sequence_length, + int num_seq, + int head_offset, + int mask_stride, + int mp_size, + int reduceWidth) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + float2 low_data[MAX_REG_SIZE]; + float2 high_data[MAX_REG_SIZE]; + const T zero_h = conversion::to(0.f); + + int wid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + + int reduce_blocks = reduceWidth >> 5; + int seq_lane = threadIdx.x % reduceWidth; + + __shared__ float partialSum[MAX_WARP_NUM]; + + int iter_offset = blockIdx.x * (warp_num / reduce_blocks) + (wid / reduce_blocks); + int batch_idx = iter_offset / (num_seq * heads); + int alibi_offset = batch_idx * heads * mp_size + head_offset; + int mask_offset = batch_idx * mask_stride + (iter_offset % mask_stride); + + if (iter_offset < total_count) { + vals += (iter_offset * sequence_length); + + alibi_offset = (alibi_offset + ((iter_offset / num_seq) % heads)) * sequence_length; + mask_offset = mask_offset * sequence_length; + int seq_id = iter_offset % num_seq; + + int real_seq_id = seq_id + (num_seq == sequence_length ? 0 : sequence_length); + int window_stride4 = (local_attention && (real_seq_id >> 2) > (window_size >> 2)) + ? (real_seq_id >> 2) - (window_size >> 2) + : 0; + int window_stride = + (local_attention && real_seq_id >= window_size) ? real_seq_id - window_size : -1; + + float max_val = minus_infinity; + // if (lane == 0) printf("%d, %d: %d \n", wid, blockIdx.x, mask_offset); + for (int i = 0; i < iterations; i++) { + int data_id = i * (reduceWidth << 2) + (seq_lane); + bool check = (data_id >> 2) >= window_stride4; + bool low_x_check = check && (data_id < sequence_length) && + (!triangular || (data_id <= seq_id)) && (data_id > window_stride); + bool low_y_check = check && ((data_id + reduceWidth) < sequence_length) && + (!triangular || ((data_id + reduceWidth) <= seq_id)) && + ((data_id + reduceWidth) > window_stride); + bool high_x_check = check && ((data_id + reduceWidth * 2) < sequence_length) && + (!triangular || ((data_id + reduceWidth * 2) <= seq_id)) && + ((data_id + reduceWidth * 2) > window_stride); + bool high_y_check = check && ((data_id + reduceWidth * 3) < sequence_length) && + (!triangular || ((data_id + reduceWidth * 3) <= seq_id)) && + ((data_id + reduceWidth * 3) > window_stride); + + if (mask && alibi) { + low_data[i].x = low_x_check + ? conversion::to(vals[data_id]) * layer_scale + + (conversion::to(alibi[data_id + alibi_offset])) + + (conversion::to(mask[data_id + mask_offset])) + : minus_infinity; + low_data[i].y = + low_y_check + ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + + (conversion::to(alibi[data_id + alibi_offset + reduceWidth])) + + (conversion::to(mask[data_id + mask_offset + reduceWidth])) + : minus_infinity; + high_data[i].x = + high_x_check + ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + + (conversion::to( + alibi[data_id + alibi_offset + reduceWidth * 2])) + + (conversion::to(mask[data_id + mask_offset + reduceWidth * 2])) + : minus_infinity; + high_data[i].y = + high_y_check + ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + + (conversion::to( + alibi[data_id + alibi_offset + reduceWidth * 3])) + + (conversion::to(mask[data_id + mask_offset + reduceWidth * 3])) + : minus_infinity; + } else if (mask) { + low_data[i].x = low_x_check + ? conversion::to(vals[data_id]) * layer_scale + + (conversion::to(mask[data_id + mask_offset])) + : minus_infinity; + low_data[i].y = + low_y_check + ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + + (conversion::to(mask[data_id + mask_offset + reduceWidth])) + : minus_infinity; + high_data[i].x = + high_x_check + ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + + (conversion::to(mask[data_id + mask_offset + reduceWidth * 2])) + : minus_infinity; + high_data[i].y = + high_y_check + ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + + (conversion::to(mask[data_id + mask_offset + reduceWidth * 3])) + : minus_infinity; + } else if (alibi) { + low_data[i].x = low_x_check + ? conversion::to(vals[data_id]) * layer_scale + + (conversion::to(alibi[data_id + alibi_offset])) + : minus_infinity; + low_data[i].y = + low_y_check + ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + + (conversion::to(alibi[data_id + alibi_offset + reduceWidth])) + : minus_infinity; + high_data[i].x = + high_x_check + ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + + (conversion::to( + alibi[data_id + alibi_offset + reduceWidth * 2])) + : minus_infinity; + high_data[i].y = + high_y_check + ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + + (conversion::to( + alibi[data_id + alibi_offset + reduceWidth * 3])) + : minus_infinity; + } else { + low_data[i].x = low_x_check ? conversion::to(vals[data_id]) * layer_scale + : minus_infinity; + low_data[i].y = + low_y_check ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + : minus_infinity; + high_data[i].x = + high_x_check + ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + : minus_infinity; + high_data[i].y = + high_y_check + ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + : minus_infinity; + } + + // if(lane == 0) printf("%f , %d, %d \n", low_data[i].x, data_id, seq_id); + max_val = (low_data[i].x > max_val ? low_data[i].x : max_val); + max_val = (low_data[i].y > max_val ? low_data[i].y : max_val); + max_val = (high_data[i].x > max_val ? high_data[i].x : max_val); + max_val = (high_data[i].y > max_val ? high_data[i].y : max_val); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + if (reduceWidth > WARP_SIZE) { + if (lane == 0) partialSum[wid] = max_val; + b.sync(); + + if (lane < warp_num) max_val = partialSum[lane]; + + b.sync(); + + for (int i = 1; i < reduce_blocks; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + max_val = g.shfl(max_val, threadIdx.x / WARP_SIZE); + } + float sum = 0; + for (int i = 0; i < iterations; i++) { + low_data[i].x = __expf(low_data[i].x - max_val); + low_data[i].y = __expf(low_data[i].y - max_val); + high_data[i].x = __expf(high_data[i].x - max_val); + high_data[i].y = __expf(high_data[i].y - max_val); + + sum += (low_data[i].x + low_data[i].y + high_data[i].x + high_data[i].y); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) sum += g.shfl_xor(sum, i); + + if (reduceWidth > WARP_SIZE) { + if (lane == 0) partialSum[wid] = sum; + b.sync(); + + if (lane < warp_num) sum = partialSum[lane]; + + b.sync(); + + for (int i = 1; i < reduce_blocks; i *= 2) { sum += g.shfl_xor(sum, i); } + + sum = g.shfl(sum, threadIdx.x / WARP_SIZE); + } + sum += 1e-6; + for (int i = 0; i < iterations; i++) { + int data_id = i * (reduceWidth << 2) + (seq_lane); + if (data_id < sequence_length) { + vals[data_id] = conversion::to(low_data[i].x / sum); + if ((data_id + reduceWidth) < sequence_length) + vals[data_id + reduceWidth] = conversion::to(low_data[i].y / sum); + if ((data_id + reduceWidth * 2) < sequence_length) + vals[data_id + reduceWidth * 2] = conversion::to(high_data[i].x / sum); + if ((data_id + reduceWidth * 3) < sequence_length) + vals[data_id + reduceWidth * 3] = conversion::to(high_data[i].y / sum); + } + } + } +} + +template +__global__ void attn_softmax_v2(float* vals, + float* attn_mask, + float* alibi, + float layer_scale, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + int total_count, + int heads, + int sequence_length, + int num_seq, + int head_offset, + int mask_stride, + int mp_size, + int reduceWidth) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + float4 data[MAX_REG_SIZE]; + + int wid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + + int reduce_blocks = reduceWidth >> 5; + int seq_lane = threadIdx.x % reduceWidth; + + __shared__ float partialSum[MAX_WARP_NUM]; + + int iter_offset = blockIdx.x * (warp_num / reduce_blocks) + (wid / reduce_blocks); + if (iter_offset < total_count) { + vals += (iter_offset * sequence_length); + + int batch_idx = iter_offset / (num_seq * heads); + int mask_offset = batch_idx * mask_stride + (iter_offset % mask_stride); + mask_offset = mask_offset * sequence_length; + int seq_id = iter_offset % num_seq; + + int real_seq_id = seq_id + (num_seq == sequence_length ? 0 : sequence_length); + int window_stride4 = (local_attention && (real_seq_id >> 2) > (window_size >> 2)) + ? (real_seq_id >> 2) - (window_size >> 2) + : 0; + int window_stride = + (local_attention && real_seq_id >= window_size) ? real_seq_id - window_size : -1; + + float max_val = minus_infinity; + + for (int i = 0; i < iterations; i++) { + int data_id = i * (reduceWidth << 2) + (seq_lane); + bool check = (data_id >> 2) >= window_stride4; + bool x_check = check && (data_id < sequence_length) && + (!triangular || (data_id <= seq_id)) && (data_id > window_stride); + bool y_check = check && ((data_id + reduceWidth) < sequence_length) && + (!triangular || ((data_id + reduceWidth) <= seq_id)) && + ((data_id + reduceWidth) > window_stride); + bool z_check = check && ((data_id + reduceWidth * 2) < sequence_length) && + (!triangular || ((data_id + reduceWidth * 2) <= seq_id)) && + ((data_id + reduceWidth * 2) > window_stride); + bool w_check = check && ((data_id + reduceWidth * 3) < sequence_length) && + (!triangular || ((data_id + reduceWidth * 3) <= seq_id)) && + ((data_id + reduceWidth * 3) > window_stride); + + if (attn_mask) { + data[i].x = x_check ? vals[data_id] + attn_mask[data_id + mask_offset] + : minus_infinity; + data[i].y = y_check ? vals[data_id + reduceWidth] + + attn_mask[data_id + mask_offset + reduceWidth] + : minus_infinity; + data[i].z = z_check ? vals[data_id + reduceWidth * 2] + + attn_mask[data_id + mask_offset + reduceWidth * 2] + : minus_infinity; + data[i].w = w_check ? vals[data_id + reduceWidth * 3] + + attn_mask[data_id + mask_offset + reduceWidth * 3] + : minus_infinity; + } else { + data[i].x = x_check ? vals[data_id] : minus_infinity; + data[i].y = y_check ? vals[data_id + reduceWidth] : minus_infinity; + data[i].z = z_check ? vals[data_id + reduceWidth * 2] : minus_infinity; + data[i].w = w_check ? vals[data_id + reduceWidth * 3] : minus_infinity; + } + + max_val = (data[i].x > max_val ? data[i].x : max_val); + max_val = (data[i].y > max_val ? data[i].y : max_val); + max_val = (data[i].z > max_val ? data[i].z : max_val); + max_val = (data[i].w > max_val ? data[i].w : max_val); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + if (reduceWidth > WARP_SIZE) { + if (lane == 0) partialSum[wid] = max_val; + b.sync(); + + if (lane < warp_num) max_val = partialSum[lane]; + + b.sync(); + + for (int i = 1; i < reduce_blocks; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + max_val = g.shfl(max_val, threadIdx.x / WARP_SIZE); + } + + float sum = 0; + for (int i = 0; i < iterations; i++) { + data[i].x = __expf(data[i].x - max_val); + data[i].y = __expf(data[i].y - max_val); + data[i].z = __expf(data[i].z - max_val); + data[i].w = __expf(data[i].w - max_val); + + sum += (data[i].x + data[i].y + data[i].z + data[i].w); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) sum += g.shfl_xor(sum, i); + + if (reduceWidth > WARP_SIZE) { + if (lane == 0) partialSum[wid] = sum; + b.sync(); + + if (lane < warp_num) sum = partialSum[lane]; + + b.sync(); + + for (int i = 1; i < reduce_blocks; i *= 2) { sum += g.shfl_xor(sum, i); } + + sum = g.shfl(sum, threadIdx.x / WARP_SIZE); + } + sum += 1e-6; + + for (int i = 0; i < iterations; i++) { + int data_id = i * (reduceWidth << 2) + (seq_lane); + if (data_id < sequence_length) { + vals[data_id] = data[i].x / sum; + if ((data_id + reduceWidth) < sequence_length) + vals[data_id + reduceWidth] = data[i].y / sum; + if ((data_id + reduceWidth * 2) < sequence_length) + vals[data_id + reduceWidth * 2] = data[i].z / sum; + if ((data_id + reduceWidth * 3) < sequence_length) + vals[data_id + reduceWidth * 3] = data[i].w / sum; + } + } + } +} + +#define LAUNCH_ATTN_SOFTMAX_V2(iterations) \ + attn_softmax_v2<<>>(vals, \ + mask, \ + alibi, \ + layer_scale, \ + triangular, \ + recompute, \ + local_attention, \ + window_size, \ + total_count, \ + heads, \ + sequence_length, \ + num_seq, \ + head_offset, \ + mask_stride, \ + mp_size, \ + reduce_width); + +template +void launch_attn_softmax_v2(T* vals, + T* mask, + T* alibi, + float layer_scale, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + int batch_size, + int heads, + int num_seq, + int sequence_length, + int head_offset, + int mask_stride, + int mp_size, + cudaStream_t stream) +{ + const int total_count = batch_size * heads * num_seq; + + // Scheduling Overview + // 4 element unroll with power of 2 `reduce_width` threads to a ceiling of `attn_threads` + // Each block should be partitioned into as many `reduce_width` blocks + // as can be fit. + constexpr int attn_threads = 256; + constexpr int min_reduce_width = hw_warp_size; + constexpr int internal_unroll = 4; + + // Handle internal unroll then round to next power of 2. Bump up to minimum granularity. + const int thread_steps_rounded = + next_pow2((sequence_length + internal_unroll - 1) / internal_unroll); + const int thread_steps_schedule = + (thread_steps_rounded < min_reduce_width) ? min_reduce_width : thread_steps_rounded; + // Bound reduce width to the number of threads + const int reduce_width = (thread_steps_schedule < attn_threads) ? thread_steps_schedule + : attn_threads; + // Scale for the excess + const int iterations = thread_steps_schedule / reduce_width; + // Should be safe since reduce_width is capped to attn_threads + const int partitions = attn_threads / reduce_width; + + // Launch params + dim3 grid((total_count + partitions - 1) / partitions); + dim3 block(attn_threads); + + if (sequence_length <= 32768) { + if (iterations == 1) { + LAUNCH_ATTN_SOFTMAX_V2(1); + } else if (iterations == 2) { + LAUNCH_ATTN_SOFTMAX_V2(2); + } else if (iterations == 4) { + LAUNCH_ATTN_SOFTMAX_V2(4); + } else if (iterations == 8) { + LAUNCH_ATTN_SOFTMAX_V2(8); + } else if (iterations == 16) { + LAUNCH_ATTN_SOFTMAX_V2(16); + } else if (iterations == 32) { + LAUNCH_ATTN_SOFTMAX_V2(32); + } else if (iterations == 64) { + LAUNCH_ATTN_SOFTMAX_V2(64); + } + } else + throw std::runtime_error("Unsupport Seq_Length!"); +} + +#define INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(T) \ + template void launch_attn_softmax_v2(T* vals, \ + T* mask, \ + T* alibi, \ + float layer_scale, \ + bool triangular, \ + bool recompute, \ + bool local_attention, \ + int window_size, \ + int batch_size, \ + int heads, \ + int num_seq, \ + int sequence_length, \ + int head_offset, \ + int mask_stride, \ + int mp_size, \ + cudaStream_t stream); + +INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(__nv_bfloat16); +#endif +INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(__half); + +#define DEF_ATTN_SOFTMAX_V2_HALF(_iter) \ + template __global__ void attn_softmax_v2<__half, _iter>(__half * vals, \ + __half * mask, \ + __half * alibi, \ + float layer_scale, \ + bool triangular, \ + bool recompute, \ + bool local_attention, \ + int window_size, \ + int total_count, \ + int heads, \ + int sequence_length, \ + int num_seq, \ + int head_offset, \ + int mask_stride, \ + int mp_size, \ + int reduceWidth) + +#define DEF_ATTN_SOFTMAX_V2_BF16(_iter) \ + template __global__ void attn_softmax_v2<__nv_bfloat16, _iter>(__nv_bfloat16 * vals, \ + __nv_bfloat16 * mask, \ + __nv_bfloat16 * alibi, \ + float layer_scale, \ + bool triangular, \ + bool recompute, \ + bool local_attention, \ + int window_size, \ + int total_count, \ + int heads, \ + int sequence_length, \ + int num_seq, \ + int head_offset, \ + int mask_stride, \ + int mp_size, \ + int reduceWidth) + +#define FOREACH_ITERATIONS(cb) \ + cb(1); \ + cb(2); \ + cb(4); \ + cb(8); \ + cb(16); \ + cb(32); \ + cb(64) + +FOREACH_ITERATIONS(DEF_ATTN_SOFTMAX_V2_HALF); +#ifdef BF16_AVAILABLE +FOREACH_ITERATIONS(DEF_ATTN_SOFTMAX_V2_BF16); +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu new file mode 100644 index 0000000000000000000000000000000000000000..e7624363021eef46738cc3a144af987c16aa51d5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu @@ -0,0 +1,733 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#ifndef __HIP_PLATFORM_AMD__ +#include +#endif +#include "conversion_utils.h" +#include "inference_cuda_layers.h" +namespace cg = cooperative_groups; + +// only used to avoid compilation error due to lack of definition. +#ifndef BF16_AVAILABLE +#if defined(__CUDA_BF16_H__) +static_assert(sizeof(__nv_bfloat162) == sizeof(__half2), + "CUDA's __nv_bfloat162 doesn't match __half2 size"); +#else +// Fallback to simple typedef only if CUDA doesn't provide it +using __nv_bfloat162 = __half2; +#endif +#endif + +// Bias add + +__global__ void bias_add_transform_0213(float* output, + float* k_cache, + float* v_cache, + const float* vals, + const float* bias, + int hidden_dim, + int seq_length, + unsigned seq_offset, + int heads, + int head_stride, + int num_kv, + int rotary_dim, + bool rotate_half, + bool rotate_every_two, + int head_ext, + int max_out_tokens, + float rope_theta) +{ + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = blockIdx.z / head_ext; // Hidden count + int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + int d2_out_stride = d2_stride * (cnt == 0 ? seq_length : max_out_tokens); + int d0_out_stride = hidden_dim * (cnt == 0 ? seq_length : max_out_tokens); + + const float4* vals_vec = reinterpret_cast(vals); + float4* output_vec = + reinterpret_cast(cnt == 0 ? output : (cnt == 1 ? k_cache : v_cache)); + + vals_vec += (d0 * (d1_stride + num_kv * 2 * d2_stride) * seq_length); + vals_vec += d1 * (d1_stride + num_kv * 2 * d2_stride); + vals_vec += (cnt == 0 ? 0 : d1_stride) + (cnt == 0 ? 0 : (cnt - 1) * num_kv * d2_stride); + vals_vec += ((cnt == 0 ? d2 : (d2 / head_stride)) * d2_stride); + + output_vec += (d1 * d2_stride); + output_vec += (d0 * d0_out_stride); + output_vec += (d2 * d2_out_stride); + + unsigned seq_id = d1 + seq_offset; + float4 inputs = vals_vec[d3]; + int lane = d3 & 0x1f; + if (cnt < 2 && rotary_dim > 0 && d3 < rotary_dim) { + float4 q = vals_vec[d3]; + float2* q_f = reinterpret_cast(&q); + if (rotate_every_two) { +#pragma unroll + for (int o = 0; o < 2; o++) { + float inv_freq = (float)(((d3 << 1) + o) * 2) / (float)(rotary_dim << 2); + inv_freq = 1.0 / powf(rope_theta, inv_freq) * (float)seq_id; + q_f[o].x = (-1.0 * q_f[o].y * sinf(inv_freq) + q_f[o].x * cosf(inv_freq)); + q_f[o].y = (q_f[o].x * sinf(inv_freq) + q_f[o].y * cosf(inv_freq)); + } + } + output_vec[d3] = q; + } else + output_vec[d3] = inputs; +} + +#define ATTN_H 3 +#define MAX_SEQ_LINE 10 + +template +__global__ void bias_add_transform_0213(T* output, // q + T* k_cache, + T* v_cache, + const T* vals, // qkv + const T* bias, + int hidden_dim, + int seq_length, + unsigned seq_offset, + int all_tokens, + int heads, + int head_stride, + int num_kv, + int rotary_dim, + bool rotate_half, + bool rotate_every_two, + int head_ext, + int max_out_tokens, + float rope_theta) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + unsigned half_dim = (rotary_dim << 3) >> 1; + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = blockIdx.z / head_ext; // Hidden count + int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + int d2_out_stride = d2_stride * (cnt == 0 ? seq_length : max_out_tokens); + int d0_out_stride = hidden_dim * (cnt == 0 ? seq_length : max_out_tokens); + + float4 vals_arr; + float4 output_arr; + + T2* vals_half = reinterpret_cast(&vals_arr); + T2* output_half = reinterpret_cast(&output_arr); + + const float4* vals_vec = reinterpret_cast(vals); + float4* output_vec = + reinterpret_cast(cnt == 0 ? output : (cnt == 1 ? k_cache : v_cache)); + + vals_vec += (d0 * (d1_stride + num_kv * 2 * d2_stride) * seq_length); + vals_vec += (d1 * (d1_stride + num_kv * 2 * d2_stride)); + vals_vec += (cnt == 0 ? 0 : d1_stride) + (cnt == 0 ? 0 : (cnt - 1) * num_kv * d2_stride); + vals_vec += ((cnt == 0 ? d2 : (d2 / head_stride)) * d2_stride); + + output_vec += (d1 * d2_stride); + output_vec += (d0 * d0_out_stride); + output_vec += (d2 * d2_out_stride); + + unsigned seq_id = d1 + seq_offset; + + int lane = d3 & 0x1f; + if (cnt < 2 && rotary_dim > 0 && d3 < rotary_dim) { + float4 q = vals_vec[d3]; + T2* q_h = reinterpret_cast(&q); + if (rotate_every_two) { +#pragma unroll + for (int o = 0; o < 4; o++) { + float inv_freq = (float)(((d3 << 2) + o) * 2) / (float)(rotary_dim << 3); + inv_freq = 1.0 / powf(rope_theta, inv_freq) * (float)seq_id; + float q_data[2]; + q_data[0] = conversion::to(q_h[o].x); + q_data[1] = conversion::to(q_h[o].y); + q_h[o].x = conversion::to(-1.0 * q_data[1] * sinf(inv_freq) + + q_data[0] * cosf(inv_freq)); + q_h[o].y = + conversion::to(q_data[0] * sinf(inv_freq) + q_data[1] * cosf(inv_freq)); + } + } + output_vec[d3] = q; + } else + output_vec[d3] = vals_vec[d3]; +} + +// [B S C*H] - > C * [B A S N] +template <> +void launch_bias_add_transform_0213(float* output, + float* k_cache, + float* v_cache, + const float* vals, + const float* bias, + int batch_size, + int seq_length, + unsigned seq_offset, + int all_tokens, + int hidden_dim, + int heads, + int num_kv, + int rotary_dim, + bool rotate_half, + bool rotate_every_two, + cudaStream_t stream, + int trans_count, + int max_out_tokens, + float rope_theta) +{ + hidden_dim >>= 2; + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + + dim3 block_dim(hidden_dim / heads, (heads / head_ext)); + dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); + + bias_add_transform_0213<<>>(output, + k_cache, + v_cache, + vals, + bias, + hidden_dim, + seq_length, + seq_offset, + heads, + num_kv > 0 ? (heads / num_kv) : 1, + num_kv > 0 ? num_kv : heads, + rotary_dim >> 2, + rotate_half, + rotate_every_two, + head_ext, + max_out_tokens, + rope_theta); +} + +template +void launch_bias_add_transform_0213(T* output, + T* k_cache, + T* v_cache, + const T* vals, + const T* bias, + int batch_size, + int seq_length, + unsigned seq_offset, + int all_tokens, + int hidden_dim, + int heads, + int num_kv, + int rotary_dim, + bool rotate_half, + bool rotate_every_two, + cudaStream_t stream, + int trans_count, + int max_out_tokens, + float rope_theta) +{ + hidden_dim >>= 3; + int head_ext = 1; // (hidden_dim - 1) / MAX_THREADS + 1; + dim3 block_dim(hidden_dim / heads, (heads / head_ext)); + dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); + bias_add_transform_0213<<>>(output, + k_cache, + v_cache, + vals, + bias, + hidden_dim, + seq_length, + seq_offset, + all_tokens, + heads, + num_kv > 0 ? (heads / num_kv) : 1, + num_kv > 0 ? num_kv : heads, + rotary_dim >> 3, + rotate_half, + rotate_every_two, + head_ext, + max_out_tokens, + rope_theta); +} + +#define INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(T) \ + template void launch_bias_add_transform_0213(T*, \ + T*, \ + T*, \ + const T*, \ + const T*, \ + int, \ + int, \ + unsigned, \ + int, \ + int, \ + int, \ + int, \ + int, \ + bool, \ + bool, \ + cudaStream_t, \ + int, \ + int, \ + float) + +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(__nv_bfloat16); +#endif +INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(__half); + +// Bias add + +__global__ void pad_add_transform_0213(float* output, + const float* vals, + int hidden_dim, + int seq_length, + int padded_seq_len, + int heads, + int padded_head_size) +{ +} + +template +__global__ void pad_add_transform_0213(T* output, + const T* vals, + int hidden_dim, + int seq_length, + int padded_seq_len, + int heads, + int padded_head_size) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + float4 ZERO; + const T2 zero_h = conversion::to(0.f); + T2* ZERO_h = reinterpret_cast(&ZERO); +#pragma unroll + for (int i = 0; i < 4; i++) ZERO_h[i] = zero_h; + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y * blockDim.z + threadIdx.z; // Sequence ID (0-127) + int d2 = threadIdx.y; // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + int d2_out_stride = padded_head_size * padded_seq_len; + int d0_out_stride = heads * d2_out_stride; + + const float4* vals_vec = reinterpret_cast(vals); + float4* output_vec = reinterpret_cast(output); + + vals_vec += (d0 * d0_stride); + vals_vec += (d1 * d1_stride); + vals_vec += (d2 * d2_stride); + + output_vec += (d1 * padded_head_size); + output_vec += (d0 * d0_out_stride); + output_vec += (d2 * d2_out_stride); + + if (d3 < d2_stride && d1 < seq_length) + output_vec[d3] = vals_vec[d3]; + else + output_vec[d3] = ZERO; +} + +// [B S C*H] - > C * [B A S N] +template <> +void launch_pad_add_transform_0213(float* output, + const float* vals, + int batch_size, + int hidden_dim, + int seq_length, + int padded_seq_len, + int heads, + int padded_head_size, + cudaStream_t stream) +{ +} + +template +void launch_pad_add_transform_0213(T* output, + const T* vals, + int batch_size, + int hidden_dim, + int seq_length, + int padded_seq_len, + int heads, + int padded_head_size, + cudaStream_t stream) +{ + hidden_dim >>= 3; + dim3 block_dim((padded_head_size >> 3), heads, 2); + dim3 grid_dim(batch_size, padded_seq_len / 2); + pad_add_transform_0213<<>>( + output, vals, hidden_dim, seq_length, padded_seq_len, heads, padded_head_size >> 3); +} + +#define INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(T) \ + template void launch_pad_add_transform_0213( \ + T*, const T*, int, int, int, int, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(__half); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(__nv_bfloat16); +#endif + +// Bias add +template +__global__ void bias_add_transform_0213(T* output, + const T* vals, + const T* bias, + int hidden_dim, + int seq_length, + int heads, + int head_ext); + +template <> +__global__ void bias_add_transform_0213(float* output, + const float* vals, + const float* bias, + int hidden_dim, + int seq_length, + int heads, + int head_ext) +{ + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = blockIdx.z / head_ext; // Hidden count + int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + const float4* vals_vec = reinterpret_cast(vals); + const float4* bias_vec = reinterpret_cast(bias); + float4* output_vec = reinterpret_cast(output); + + float4 inputs = vals_vec[d0 * d0_stride * (gridDim.z / head_ext) + cnt * d1_stride + + d1 * d1_stride * (gridDim.z / head_ext) + d2 * d2_stride + d3]; + float4 biases = bias_vec[cnt * d1_stride + d2 * d2_stride + d3]; + + float4 outputs; + outputs.x = inputs.x + biases.x; + outputs.y = inputs.y + biases.y; + outputs.z = inputs.z + biases.z; + outputs.w = inputs.w + biases.w; + + output_vec[cnt * d0_out_stride * gridDim.x + d0 * d0_out_stride + d1 * d1_out_stride + + d2 * d2_out_stride + d3] = outputs; +} + +template +__global__ void bias_add_transform_0213(T* output, + const T* vals, + const T* bias, + int hidden_dim, + int seq_length, + int heads, + int head_ext) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = blockIdx.z / head_ext; // Hidden count + int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + float4 vals_arr; + float4 bias_arr; + float4 output_arr; + T2* vals_half = reinterpret_cast(&vals_arr); + T2* bias_half = reinterpret_cast(&bias_arr); + T2* output_half = reinterpret_cast(&output_arr); + + const float4* vals_vec = reinterpret_cast(vals); + const float4* bias_vec = reinterpret_cast(bias); + float4* output_vec = reinterpret_cast(output); + + vals_vec += (d0 * d0_stride * (gridDim.z / head_ext)); + vals_vec += (d1 * d1_stride * (gridDim.z / head_ext)); + vals_vec += (cnt * d1_stride); + vals_vec += (d2 * d2_stride); + + bias_vec += (cnt * d1_stride); + bias_vec += (d2 * d2_stride); + + output_vec += (cnt * d0_stride * gridDim.x); + output_vec += (d1 * d2_stride); + output_vec += (d0 * d0_stride); + output_vec += (d2 * d2_out_stride); + + bias_arr = bias_vec[d3]; + vals_arr = vals_vec[d3]; + + output_half[0] = vals_half[0] + bias_half[0]; + output_half[1] = vals_half[1] + bias_half[1]; + output_half[2] = vals_half[2] + bias_half[2]; + output_half[3] = vals_half[3] + bias_half[3]; + output_vec[d3] = output_arr; +} + +template +__global__ void bias_add_transform_0213_v2(T* output, + const T* vals, + const T* bias, + int hidden_dim, + int seq_length, + int heads) +{ + using T2 = + typename std::conditional::value, __half2, __nv_bfloat162>::type; + __shared__ float4 in_data[3072]; + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + int iteration_stride = d1_stride * blockDim.z; // Hidden * 3 / 8 + int batch_stride = d0_stride * blockDim.z; // Hidden * S * 3 / 8 + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = threadIdx.z; // blockIdx.z; // Hidden count + int d2 = threadIdx.y; // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + float4 vals_arr[1]; + float4 bias_arr[1]; + float4 output_arr[1]; + T2* vals_half = reinterpret_cast(vals_arr); + T2* bias_half = reinterpret_cast(bias_arr); + T2* output_half = reinterpret_cast(output_arr); + + const float4* vals_vec = reinterpret_cast(vals); + const float4* bias_vec = reinterpret_cast(bias); + float4* output_vec = reinterpret_cast(output); + + int iter_index = cnt * d1_stride + d2 * d2_stride + d3; + int input_offset = d0 * batch_stride + d1 * (iteration_stride << 1); + bias_arr[0] = bias_vec[iter_index]; + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_id = iter * iteration_stride + iter_index; + vals_arr[0] = vals_vec[input_offset + iter_id]; + + output_half[0] = vals_half[0] + bias_half[0]; + output_half[1] = vals_half[1] + bias_half[1]; + output_half[2] = vals_half[2] + bias_half[2]; + output_half[3] = vals_half[3] + bias_half[3]; + + in_data[iter_id] = output_arr[0]; + } + __syncthreads(); + + iteration_stride = blockDim.z * (blockDim.y >> 1); + int matrix_stride = (d0_out_stride * gridDim.x); + int head_count = (d2 >> 1) + cnt * (blockDim.y >> 1); + + int out_index = d0 * d0_out_stride + d1 * (d1_out_stride << 1) + d3 + (d2 % 2) * d2_stride; + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_row = (iter * iteration_stride) + head_count; + int iter_offset = + (iter_row % blockDim.y) * d2_out_stride + (iter_row / blockDim.y) * matrix_stride; + output_vec[out_index + iter_offset] = + in_data[iter_row * d2_stride + d3 + (d2 % 2) * (d1_stride * blockDim.z)]; + } +} + +template +__global__ void transform4d_0213(T* out, + const T* in, + int heads, + int seq_length, + int hidden_dim, + int head_ext); + +template <> +__global__ void transform4d_0213(float* out, + const float* in, + int heads, + int seq_length, + int hidden_dim, + int head_ext) +{ + int d0_stride = hidden_dim * seq_length; + int d1_stride = d0_stride / heads; + int d2_stride = hidden_dim / heads; + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = hidden_dim; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y / ((seq_length - 1) / blockDim.y + 1); // Head + int d2 = (threadIdx.y + blockDim.y * blockIdx.y) % seq_length; + int cnt = blockIdx.z; + int d3 = threadIdx.x; // Values (groups of 8) + + if (d2 < seq_length) { + const float4* in_vec = reinterpret_cast(in); + float4* out_vec = reinterpret_cast(out); + + float4 vals_vec = in_vec[cnt * d0_stride * gridDim.x + d0 * d0_stride + d1 * d1_stride + + d2 * d2_stride + d3]; + out_vec[d0 * d0_out_stride * gridDim.z + cnt * d2_out_stride + d1 * d1_out_stride + + d2 * d2_out_stride * gridDim.z + d3] = vals_vec; + } +} + +template +__global__ void transform4d_0213(T* out, + const T* in, + int heads, + int seq_length, + int hidden_dim, + int head_ext) +{ + int d0_stride = hidden_dim * (seq_length / head_ext); + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head + int d2 = blockIdx.z / head_ext; // Sequence + int cnt = blockIdx.y; // Hidden count + int d3 = threadIdx.x; // Values (groups of 8) + + const float4* in_vec = reinterpret_cast(in); + float4* out_vec = reinterpret_cast(out); + + in_vec += (cnt * d0_stride * gridDim.x); + in_vec += (d0 * d0_stride); + in_vec += (d2 * d2_stride); + in_vec += (d1 * d2_stride * seq_length); + + out_vec += (cnt * d1_stride); + out_vec += (d1 * d2_stride); + out_vec += (d0 * d0_stride * gridDim.y); + out_vec += (d2 * d1_stride * gridDim.y); + + out_vec[d3] = in_vec[d3]; +} + +template +__global__ void transform4d_0213_v2(T* out, const T* in, int heads, int seq_length, int hidden_dim) +{ + __shared__ float4 in_data[3072]; + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = threadIdx.y; // Head + int d2 = blockIdx.y; // Sequence + int cnt = threadIdx.z; // Hidden count + int d3 = threadIdx.x; // Values (groups of 8) + + const float4* in_vec = reinterpret_cast(in); + float4* out_vec = reinterpret_cast(out); + + int input_offset = d0 * d0_stride + d2 * (d2_stride << 1) + d3 + (d1 % 2) * d2_stride; + int head_count = (d1 >> 1) + cnt * (blockDim.y >> 1); + int iteration_stride = blockDim.z * (blockDim.y >> 1); + int matrix_stride = (d0_stride * gridDim.x); + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_row = iter * iteration_stride + head_count; + int iter_offset = (iter_row % blockDim.y) * d2_stride; + + in_data[d3 + iter_offset + (iter_row / blockDim.y + (d1 % 2) * blockDim.z) * d1_stride] = + in_vec[input_offset + iter_offset * seq_length + + (iter_row / blockDim.y) * matrix_stride]; + } + __syncthreads(); + + iteration_stride = d1_stride * blockDim.z; + int iter_index = cnt * d1_stride + d1 * d2_stride + d3; + int output_offset = d0 * d0_stride * blockDim.z + d2 * (iteration_stride << 1); + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_id = iter * iteration_stride + iter_index; + out_vec[output_offset + iter_id] = in_data[iter_id]; + } +} + +// 3 * [B A S N] - > [B S C*H] +template <> +void launch_transform4d_0213(float* out, + const float* in, + int batch_size, + int heads, + int seq_length, + int hidden_dim, + cudaStream_t stream, + int trans_count) +{ + hidden_dim >>= 2; + dim3 grid_dims(batch_size, heads * ((seq_length - 1) / 8 + 1), trans_count); + dim3 block_dims(hidden_dim / heads, 8); + transform4d_0213 + <<>>(out, in, heads, seq_length, hidden_dim, 1); +} + +template +void launch_transform4d_0213(T* out, + const T* in, + int batch_size, + int heads, + int seq_length, + int hidden_dim, + cudaStream_t stream, + int trans_count) +{ + hidden_dim >>= 3; + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + dim3 grid_dims(batch_size, trans_count, (seq_length * head_ext)); + dim3 block_dims(hidden_dim / heads, (heads / head_ext)); + transform4d_0213<<>>( + out, in, heads, seq_length, hidden_dim, head_ext); +} + +#define INSTANTIATE_2B_LAUNCH_TRANSFORM4D(T) \ + template void launch_transform4d_0213(T*, const T*, int, int, int, int, cudaStream_t, int); + +INSTANTIATE_2B_LAUNCH_TRANSFORM4D(__half) +#ifdef BF16_AVAILABLE +INSTANTIATE_2B_LAUNCH_TRANSFORM4D(__nv_bfloat16) +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_context.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_context.h new file mode 100644 index 0000000000000000000000000000000000000000..378fd4e5e990d5be9b642e2092ab30370506b191 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_context.h @@ -0,0 +1,292 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include +#include "cublas_v2.h" +#include "cuda.h" + +#define MEGABYTE (1024 * 1024) +#define GIGABYTE (1024 * 1024 * 1024) + +// TODO: refactor out +#define WARP_SIZE 32 + +#define CUDA_CHECK(callstr) \ + { \ + cudaError_t error_code = callstr; \ + if (error_code != cudaSuccess) { \ + std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ + assert(0); \ + } \ + } + +#define CUDA_1D_KERNEL_LOOP(i, n) \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) + +#define CUDA_2D_KERNEL_LOOP(i, n, j, m) \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) \ + for (size_t j = blockIdx.y * blockDim.y + threadIdx.y; j < (m); j += blockDim.y * gridDim.y) + +#define DS_CUDA_NUM_THREADS 512 +#define DS_MAXIMUM_NUM_BLOCKS 262144 + +inline int DS_GET_BLOCKS(const int N) +{ + return std::max( + std::min((N + DS_CUDA_NUM_THREADS - 1) / DS_CUDA_NUM_THREADS, DS_MAXIMUM_NUM_BLOCKS), + // Use at least 1 block, since CUDA does not allow empty block + 1); +} + +class InferenceContext { +public: + InferenceContext() + : _workspace(nullptr), + _seed(42), + _curr_offset(0), + _stream(0), + _free_memory_size(0), + _num_tokens(1), + _attention_unfused_workspace_offset(0), + _workSpaceSize(0) + { + _workSpaceSize = 0; + _workspace = 0; + + cublasStatus_t stat = cublasCreate(&_cublasHandle); + if (stat != CUBLAS_STATUS_SUCCESS) { + // It would be nice to use cublasGetStatusName and + // cublasGetStatusString, but they were only added in CUDA 11.4.2. + auto message = std::string("Failed to create cublas handle: cublasStatus_t was ") + + std::to_string(stat); + std::cerr << message << std::endl; + throw std::runtime_error(message); + } +#ifndef __HIP_PLATFORM_AMD__ + cublasSetMathMode(_cublasHandle, CUBLAS_TENSOR_OP_MATH); +#endif + cudaEventCreate(&_comp1_event); + cudaEventCreate(&_comp2_event); + cudaEventCreate(&_comp_event); + cudaEventCreate(&_comm_event); + } + + virtual ~InferenceContext() + { + cublasDestroy(_cublasHandle); + cudaFree(_workspace); + cudaEventDestroy(_comp1_event); + cudaEventDestroy(_comp2_event); + cudaEventDestroy(_comp_event); + cudaEventDestroy(_comm_event); + } + + static InferenceContext& Instance() + { + static InferenceContext _ctx; + return _ctx; + } + + void GenWorkSpace(const unsigned& num_layers, + const unsigned& num_heads, + const size_t& batch_size, + const size_t& prompt_len, + const size_t& hidden_dim, + const unsigned& mp_size, + const bool& external_cache, + const size_t& elem_size, + const unsigned& rank, + unsigned max_out_tokens, + unsigned min_out_tokens) + { + size_t total_size; + if (!_free_memory_size) { cudaMemGetInfo(&_free_memory_size, &total_size); } + + // Flash attention requires padded heads and we'll conservatively allocate + // for that here. Flash attention is only enabled for head size <= 128 right now + const int head_size = hidden_dim / num_heads; + const int padded_head_size = head_size <= 32 ? 32 : (head_size <= 64 ? 64 : 128); + const int effective_head_size = (head_size > 128) ? head_size : padded_head_size; + + size_t activation_size = 10 * (num_heads * effective_head_size) * batch_size; + // Other sequence length dimension is added when the final workSpaceSize is calculated + size_t temp_size = batch_size * (num_heads / mp_size) * max_out_tokens; + size_t cache_size = + num_layers * batch_size * ((num_heads * effective_head_size) / mp_size) * 2; + size_t minimal_requirements = + temp_size + (_free_memory_size > GIGABYTE ? 500 : 100) * MEGABYTE; + if (_free_memory_size < minimal_requirements) { + printf("Requested:\t%lu\nFree:\t%lu\nTotal:\t%lu\n", + minimal_requirements, + _free_memory_size, + total_size); + throw std::runtime_error("Workspace can't be allocated, no enough memory."); + } + + _max_seq_len = ((_free_memory_size - minimal_requirements) / elem_size) / + (activation_size + temp_size + cache_size); + _max_seq_len = std::min((size_t)max_out_tokens, _max_seq_len); + size_t workSpaceSize = ((external_cache ? (activation_size + temp_size) + : (activation_size + temp_size + cache_size))) * + _max_seq_len * elem_size; + temp_size *= _max_seq_len * elem_size; + + if (_max_seq_len < min_out_tokens) { + printf( + "Allocatable workspace available (%ld tokens) is less than minimum requested " + "workspace (%d tokens)\n", + _max_seq_len, + min_out_tokens); + throw std::runtime_error("Workspace can't be allocated, not enough memory"); + } + + if (!_workspace) { + assert(_workspace == nullptr); + cudaMalloc(&_workspace, workSpaceSize); + } else if (_workSpaceSize < workSpaceSize) { + cudaFree(_workspace); + cudaMalloc(&_workspace, workSpaceSize); + } + if (rank == 0 && (!_workspace || _workSpaceSize < workSpaceSize)) + printf( + "------------------------------------------------------\n" + "Free memory : %f (GigaBytes) \n" + "Total memory: %f (GigaBytes) \n" + "Requested memory: %f (GigaBytes) \n" + "Setting maximum total tokens (input + output) to %lu \n" + "WorkSpace: %p \n" + "------------------------------------------------------\n", + (float)_free_memory_size / GIGABYTE, + (float)total_size / GIGABYTE, + (float)workSpaceSize / GIGABYTE, + _max_seq_len, + _workspace); + + if (!_workspace) { + printf("Requested:\t%lu\nFree:\t%lu\nTotal:\t%lu\n", + workSpaceSize, + _free_memory_size, + total_size); + throw std::runtime_error("Workspace is null."); + } + _workSpaceSize = workSpaceSize; + _attention_unfused_workspace_offset = workSpaceSize - temp_size; + } + inline size_t GetMaxTokenLength() const { return _max_seq_len; } + + cudaEvent_t GetCompEvent(int id) { return id == 1 ? _comp1_event : _comp2_event; } + + size_t get_workspace_size() const { return _workSpaceSize; } + void* GetWorkSpace() { return _workspace; } + void* GetAttentionUnfusedWorkspace() + { + return (char*)_workspace + _attention_unfused_workspace_offset; + } + + inline unsigned new_token(unsigned layer_id) + { + if (layer_id == 0) _token_length++; + return _token_length; + } + + inline void reset_tokens(unsigned initial_tokens = 1) + { + _num_tokens = initial_tokens; + } //_token_length = 0; } + + inline unsigned current_tokens() const { return _num_tokens; } + + inline void advance_tokens() { _num_tokens++; } + + cudaStream_t GetCommStream(bool async_op = false) + { + if (!_comm_stream) + _comm_stream = async_op ? at::cuda::getStreamFromPool(true) + : at::cuda::getCurrentCUDAStream(); + return _comm_stream; + } + cudaStream_t GetCurrentStream(bool other_stream = false) + { + // get current pytorch stream. + if (other_stream) { + if (!_stream) _stream = at::cuda::getStreamFromPool(true); + return _stream; + } + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + return stream; + } + + void release_workspace() + { + cudaFree(_workspace); + _workspace = nullptr; + } + bool retake_workspace() + { + if (_workspace != nullptr || _workSpaceSize == 0) return true; + cudaMalloc(&_workspace, _workSpaceSize); + return _workspace != nullptr; + } + cublasHandle_t GetCublasHandle() { return _cublasHandle; } + + std::pair IncrementOffset(uint64_t offset_inc) + { + uint64_t offset = _curr_offset; + _curr_offset += offset_inc; + return std::pair(_seed, offset); + } + + void SetSeed(uint64_t new_seed) { _seed = new_seed; } + + const std::vector>& GetGemmAlgos() const { return _gemm_algos; } + + inline void SynchComp() + { + cudaEventRecord(_comp_event, _comp_stream); + cudaStreamWaitEvent(_comm_stream, _comp_event, 0); + } + inline void SynchComm() + { + cudaEventRecord(_comm_event, _comm_stream); + cudaStreamWaitEvent(_comp_stream, _comm_event, 0); + } + +private: + cublasHandle_t _cublasHandle; + + cudaEvent_t _comp_event; + cudaEvent_t _comm_event; + + void* _workspace; + // offset from _workspace for attention unfused memory + size_t _attention_unfused_workspace_offset; + uint64_t _seed; + uint64_t _curr_offset; + + size_t _workSpaceSize; + size_t _free_memory_size; + + size_t _max_seq_len; + + cudaEvent_t _comp1_event; + cudaEvent_t _comp2_event; + + cudaStream_t _stream; + + unsigned _token_length; + unsigned _num_tokens; + std::vector> _gemm_algos; + + cudaStream_t _comp_stream; + cudaStream_t _comm_stream; + + std::unordered_map _world_sizes; +}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h new file mode 100644 index 0000000000000000000000000000000000000000..40c3e443941d24174ee27c2728cd95907e654074 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h @@ -0,0 +1,504 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#ifdef BF16_AVAILABLE +#include +#endif +#include +#include +#ifndef __HIP_PLATFORM_AMD__ +#include +#endif +#include + +// TODO HIP: Remove backward compatibility for torch<=2.0 in future +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_gemm_ex(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + rocblas_gemm_algo algo, + int b_stride = -1) +#else +int cublas_gemm_ex(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + cublasGemmAlgo_t algo, + int b_stride = -1) +#endif +{ + const int ldb = (b_stride == -1) ? ((transb == CUBLAS_OP_N) ? k : n) : b_stride; +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_status status = rocblas_gemm_ex(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, + rocblas_datatype_f32_r, + (transa == rocblas_operation_none) ? m : k, + (const void*)B, + rocblas_datatype_f32_r, + ldb, + (const void*)beta, + C, + rocblas_datatype_f32_r, + m, + C, + rocblas_datatype_f32_r, + m, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else + cublasStatus_t status = cublasGemmEx(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (transa == CUBLAS_OP_N) ? m : k, + (const void*)B, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + ldb, + (const void*)beta, + C, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + m, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + return 0; +} + +template +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_gemm_ex(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const T* A, + const T* B, + T* C, + rocblas_gemm_algo algo, + int b_stride = -1) +#else +int cublas_gemm_ex(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const T* A, + const T* B, + T* C, + cublasGemmAlgo_t algo, + int b_stride = -1) +#endif +{ + const int ldb = (b_stride == -1) ? ((transb == CUBLAS_OP_N) ? k : n) : b_stride; +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + constexpr auto rocblas_dtype_16 = std::is_same::value ? rocblas_datatype_f16_r + : rocblas_datatype_bf16_r; + rocblas_status status = rocblas_gemm_ex(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, + rocblas_dtype_16, + (transa == rocblas_operation_none) ? m : k, + (const void*)B, + rocblas_dtype_16, + ldb, + (const void*)beta, + (void*)C, + rocblas_dtype_16, + m, + (void*)C, + rocblas_dtype_16, + m, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else +#ifdef __HIP_PLATFORM_AMD__ + constexpr auto cublas_dtype_16 = std::is_same::value ? HIPBLAS_R_16F : HIPBLAS_R_16B; +#else + constexpr auto cublas_dtype_16 = std::is_same::value ? CUDA_R_16F : CUDA_R_16BF; +#endif + cublasStatus_t status = cublasGemmEx(handle, + transa, + transb, + m, + n, + k, + (const void*)alpha, + (const void*)A, + cublas_dtype_16, + (transa == CUBLAS_OP_N) ? m : k, + (const void*)B, + cublas_dtype_16, + ldb, + (const void*)beta, + (void*)C, + cublas_dtype_16, + m, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + return 0; +} + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_strided_batched_gemm(rocblas_handle handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + rocblas_operation op_A, + rocblas_operation op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + rocblas_gemm_algo algo) +#else +int cublas_strided_batched_gemm(cublasHandle_t handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const float* A, + const float* B, + float* C, + cublasOperation_t op_A, + cublasOperation_t op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + cublasGemmAlgo_t algo) +#endif +{ +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + rocblas_status status = + rocblas_gemm_strided_batched_ex(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, + rocblas_datatype_f32_r, + (op_A == rocblas_operation_none) ? m : k, + stride_A, + B, + rocblas_datatype_f32_r, + (op_B == rocblas_operation_none) ? k : n, + stride_B, + beta, + C, + rocblas_datatype_f32_r, + m, + stride_C, + C, + rocblas_datatype_f32_r, + m, + stride_C, + batch, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else + cublasStatus_t status = cublasGemmStridedBatchedEx(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (op_A == CUBLAS_OP_N) ? m : k, + stride_A, + B, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + (op_B == CUBLAS_OP_N) ? k : n, + stride_B, + beta, + C, +#ifdef __HIP_PLATFORM_AMD__ + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + m, + stride_C, + batch, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (batch: %d, m: %d, n: %d, k: %d, error: %d) \n", + batch, + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + return 0; +} + +template +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) +int cublas_strided_batched_gemm(rocblas_handle handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const T* A, + const T* B, + T* C, + rocblas_operation op_A, + rocblas_operation op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + rocblas_gemm_algo algo) +#else +int cublas_strided_batched_gemm(cublasHandle_t handle, + int m, + int n, + int k, + const float* alpha, + const float* beta, + const T* A, + const T* B, + T* C, + cublasOperation_t op_A, + cublasOperation_t op_B, + int stride_A, + int stride_B, + int stride_C, + int batch, + cublasGemmAlgo_t algo) +#endif +{ +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + constexpr auto rocblas_dtype_16 = std::is_same::value ? rocblas_datatype_f16_r + : rocblas_datatype_bf16_r; + rocblas_status status = + rocblas_gemm_strided_batched_ex(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, + rocblas_dtype_16, + (op_A == rocblas_operation_none) ? m : k, + stride_A, + B, + rocblas_dtype_16, + (op_B == rocblas_operation_none) ? k : n, + stride_B, + beta, + C, + rocblas_dtype_16, + m, + stride_C, + C, + rocblas_dtype_16, + m, + stride_C, + batch, + rocblas_datatype_f32_r, + algo, + 0, + 0); +#else +#ifdef __HIP_PLATFORM_AMD__ + constexpr auto cublas_dtype_16 = std::is_same::value ? HIPBLAS_R_16F : HIPBLAS_R_16B; +#else + constexpr auto cublas_dtype_16 = std::is_same::value ? CUDA_R_16F : CUDA_R_16BF; +#endif + cublasStatus_t status = cublasGemmStridedBatchedEx(handle, + op_A, + op_B, + m, + n, + k, + alpha, + A, + cublas_dtype_16, + (op_A == CUBLAS_OP_N) ? m : k, + stride_A, + B, + cublas_dtype_16, + (op_B == CUBLAS_OP_N) ? k : n, + stride_B, + beta, + C, + cublas_dtype_16, + m, + stride_C, + batch, +#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) + HIPBLAS_COMPUTE_32F, +#elif defined(__HIP_PLATFORM_AMD__) + HIPBLAS_R_32F, +#else + CUDA_R_32F, +#endif + algo); +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) + if (status != rocblas_status_success) { +#else + if (status != CUBLAS_STATUS_SUCCESS) { +#endif + fprintf(stderr, + "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", + m, + n, + k, + (int)status); + return EXIT_FAILURE; + } + + return 0; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h new file mode 100644 index 0000000000000000000000000000000000000000..dcc020483687f751f5997863307504449743fa35 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h @@ -0,0 +1,248 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "ds_kernel_utils.h" + +#include +#ifdef BF16_AVAILABLE +#include +#endif +#include +#include +#include +#include +#include + +#define MAX_WARP_NUM 32 +#define WARP_SIZE 32 + +#define MAX_THREADS 1024 +#define SMs 80 + +#define MAX_REGISTERS 256 + +template +void launch_attn_softmax_v2(T* vals, + T* mask, + T* alibi, + float layer_scale, + bool triangular, + bool recompute, + bool local_attention, + int window_size, + int batch_size, + int heads, + int num_seq, + int sequence_length, + int offset, + int mask_stride, + int mp_size, + cudaStream_t stream); + +// Fused bias add with gelu activation +template +void launch_bias_gelu(T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream); + +template +void launch_gated_activation(T* output, + const T* activation, + const T* bias, + int rows, + int output_stride, + int elems_per_row, + bool use_gelu, + cudaStream_t stream); + +// Fused bias add with relu activation +template +void launch_bias_relu(T* input, + const T* bias, + int intermediate_size, + int batch_size, + cudaStream_t stream); + +template +void launch_bias_add(T* input, const T* bias, int hidden_size, int batch_size, cudaStream_t stream); + +template +void launch_bias_residual(T* input, + T* output, + T* attn, + T* bias, + T* attn_bias, + int batch, + int hidden_dim, + int mp_size, + bool preln, + cudaStream_t stream); + +template +void launch_fused_ln(T* output, + const T* vals, + const T* gamma, + const T* beta, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream); + +template +void launch_fused_residual_ln(T* output, + const T* vals, + const T* residual, + const T* bias, + const T* gamma, + const T* beta, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream); + +template +void launch_fused_residual_ln_store_pre_ln_res(T* norm_output, + T* res_output, + const T* vals, + const T* residual, + const T* bias, + const T* gamma, + const T* beta, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream); + +template +void launch_rms_norm(T* norm_output, + T* res_output, + const T* vals, + const T* residual, + const T* gamma, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream); + +template +void launch_dequantize(T* output, + const int8_t* input, + const float* qscale, + unsigned output_size, + unsigned hidden_dim, + unsigned groups, + unsigned merge_count, + cudaStream_t stream); + +template +void launch_dequantize(T* output, + const int8_t* input, + const float* qscale, + unsigned output_size, + unsigned hidden_dim, + unsigned groups, + cudaStream_t stream); +template +void launch_gptj_residual_add(T* input, + T* output, + T* attn, + T* bias, + T* attn_bias, + int batch, + int head_size, + int mp_size, + cudaStream_t stream); + +template +void launch_apply_rotary_pos_emb(T* mixed_query, + T* key_layer, + unsigned head_size, + unsigned seq_len, + unsigned rotary_dim, + unsigned offset, + unsigned num_heads, + unsigned batch, + float rope_theta, + cudaStream_t stream, + int max_out_tokens); + +template +void launch_moe_res_matmul(T* residual, + T* coef, + T* mlp_out, + int seq_len, + int hidden_dim, + cudaStream_t stream); + +// 4D transform [0, 1, 2, 3] -> [0, 2, 1, 3] +template +void launch_transform4d_0213(T* out, + const T* in, + int batch_size, + int heads, + int seq_length, + int hidden_dim, + cudaStream_t stream, + int trans_count); +template +void launch_bias_add_transform_0213(T* outputs, + T* vals, + T* vals1, + const T* vals2, + const T* bias, + int batch_size, + int seq_length, + unsigned seq_offset, + int seq_length1, + int hidden_dim, + int heads, + int num_kv, + int rotary_dim, + bool rotate_half, + bool rotate_every_two, + cudaStream_t stream, + int trans_count, + int max_out_tokens, + float rope_theta); +template +void pad_data(T* padded_output, + T* output, + int bsz, + int head_size, + int padded_head_size, + cudaStream_t stream); + +template +void pad_head_seq(T* padded_output, + T* output, + int bsz, + int seq_len, + int padded_seq_len, + int head_size, + int padded_head_size, + cudaStream_t stream); + +template +void launch_pad_add_transform_0213(T* output, + const T* vals, + int batch_size, + int hidden_dim, + int seq_length, + int padded_seq_len, + int heads, + int padded_head_size, + cudaStream_t stream); + +template +void launch_vector_add(T* out, + const T* a, + const T* b, + float gamma, + int num_elems, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/normalize_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/normalize_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..8b9dbaea6c86d98ed7e431073684b89b8d1499ee --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/normalize_kernels.cu @@ -0,0 +1,2134 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "custom_cuda_layers.h" + +namespace cg = cooperative_groups; + +/* +Fused bias add, residual (elementwise) add, and normalization layer. + +For FP16, this kernel does not promote to FP32 in order to utilize the 2x throughput for +__half2 instructions, and avoid the conversion overhead (1/8 of __hal2 arithmetic). + +For specific launch constraints, see the launch functions. +*/ + +#define NORM_REG (MAX_REGISTERS / 4) + +__global__ void fused_bias_residual_layer_norm(float* vals, + const float* residual, + const float* gamma, + const float* beta, + float epsilon, + bool preLayerNorm, + bool training, + float* vars, + float* means, + int row_stride) +{ + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int gid = id / WARP_SIZE; + + float vals_arr[NORM_REG]; + __shared__ float shr[MAX_WARP_NUM]; + + residual += (row * row_stride); + vals += (row * row_stride); + + float sum = 0.f; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + vals_arr[i] = residual[i * iteration_stride + id]; + sum += vals_arr[i]; + } + if (high_index < row_stride) { + vals_arr[iterations] = residual[high_index]; + sum += vals_arr[iterations]; + iterations++; + } + + for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) shr[gid] = sum; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; + +#if !defined(__STOCHASTIC_MODE__) || __CUDA_ARCH__ < 700 + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + sum += g.shfl_down(sum, i); + } + + sum = g.shfl(sum, 0); + float mean = sum / row_stride; + if (training) + if (threadIdx.x == 0) means[row] = mean; + float variance = 0.f; + for (int i = 0; i < iterations; i++) { + vals_arr[i] -= mean; + variance += vals_arr[i] * vals_arr[i]; + } + + for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } + + if (g.thread_rank() == 0) shr[gid] = variance; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + variance += g.shfl_down(variance, i); + } + variance = g.shfl(variance, 0); + variance /= row_stride; + variance += epsilon; + if (training) + if (threadIdx.x == 0) vars[row] = variance; + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + vals_arr[i] = vals_arr[i] * rsqrtf(variance); + vals_arr[i] = + vals_arr[i] * gamma[i * iteration_stride + id] + beta[i * iteration_stride + id]; + vals[i * iteration_stride + id] = vals_arr[i]; + } + if ((high_index) < row_stride) { + vals_arr[iterations] = vals_arr[iterations] * rsqrtf(variance); + vals_arr[iterations] = vals_arr[iterations] * gamma[high_index] + beta[high_index]; + vals[high_index] = vals_arr[iterations]; + } +} + +__global__ void fused_bias_residual_layer_norm(__half* vals, + const __half* residual, + const __half* gamma, + const __half* beta, + float epsilon, + bool preLayerNorm, + bool training, + __half* vars, + __half* means, + int row_stride) +{ +#ifdef HALF_PRECISION_AVAILABLE + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int gid = id >> WARP_SIZE_BITS; + + float2 vals_f[NORM_REG]; + __shared__ float shr[MAX_WARP_NUM]; + + __half2* vals_cast = reinterpret_cast<__half2*>(vals); + const __half2* residual_cast = reinterpret_cast(residual); + + residual_cast += (row * row_stride); + vals_cast += (row * row_stride); + + float sum = 0.f; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + vals_f[i] = __half22float2(residual_cast[i * iteration_stride + id]); + sum += vals_f[i].x; + sum += vals_f[i].y; + } + if ((high_index) < row_stride) { + vals_f[iterations] = __half22float2(residual_cast[high_index]); + sum += vals_f[iterations].x; + sum += vals_f[iterations].y; + iterations++; + } + + for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) shr[gid] = sum; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + sum += g.shfl_down(sum, i); + } + sum = g.shfl(sum, 0); + float mean = sum / (row_stride * 2); + + float variance = 0.f; + for (int i = 0; i < iterations; i++) { + vals_f[i].x -= mean; + vals_f[i].y -= mean; + variance += vals_f[i].x * vals_f[i].x; + variance += vals_f[i].y * vals_f[i].y; + } + + for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } + + if (g.thread_rank() == 0) shr[gid] = variance; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + variance += g.shfl_down(variance, i); + } + variance = g.shfl(variance, 0); + variance /= (row_stride * 2); + variance += epsilon; + + __half2 variance_h = __float2half2_rn(variance); + const __half2* gamma_cast = reinterpret_cast(gamma); + const __half2* beta_cast = reinterpret_cast(beta); + + if (training && threadIdx.x == 0) { + vars[row] = __float2half(variance); + means[row] = __float2half(mean); + } + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + __half2 vals_arr = __float22half2_rn(vals_f[i]); + vals_arr = vals_arr * h2rsqrt(variance_h); + vals_arr = + vals_arr * gamma_cast[i * iteration_stride + id] + beta_cast[i * iteration_stride + id]; + vals_cast[i * iteration_stride + id] = vals_arr; + } + if ((high_index) < row_stride) { + __half2 vals_arr = __float22half2_rn(vals_f[iterations]); + vals_arr = vals_arr * h2rsqrt(variance_h); + vals_arr = vals_arr * gamma_cast[high_index] + beta_cast[high_index]; + vals_cast[high_index] = vals_arr; + } +#endif +} + +template +void launch_bias_residual_layer_norm(T* vals, + const T* residual, + const T* gamma, + const T* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + T* vars, + T* means); + +template <> +void launch_bias_residual_layer_norm(float* vals, + const float* residual, + const float* gamma, + const float* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + float* vars, + float* means) +{ + int threads = THREADS; + + dim3 grid_dim(batch_size); + + if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 1; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 2; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim(threads); + + fused_bias_residual_layer_norm<<>>( + vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, means, hidden_dim); +} + +template <> +void launch_bias_residual_layer_norm<__half>(__half* vals, + const __half* residual, + const __half* gamma, + const __half* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + __half* vars, + __half* means) +{ + int threads = 128; + + dim3 grid_dim(batch_size); + + if (hidden_dim > 8192 && hidden_dim <= 16384) + threads <<= 1; + else if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 2; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 3; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim(threads); + + fused_bias_residual_layer_norm<<>>( + vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, means, hidden_dim / 2); +} + +__global__ void fused_bias_residual_layer_norm(float* vals, + const float* residual, + const float* gamma, + const float* beta, + float epsilon, + bool preLayerNorm, + bool training, + float* vars, + int row_stride) +{ + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int gid = id / 32; + + float vals_arr[NORM_REG]; + __shared__ float shr[MAX_WARP_NUM]; + + residual += (row * row_stride); + vals += (row * row_stride); + + float sum = 0.f; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + vals_arr[i] = residual[i * iteration_stride + id]; + sum += vals_arr[i]; + } + if ((high_index) < row_stride) { + vals_arr[iterations] = residual[high_index]; + sum += vals_arr[iterations]; + iterations++; + } + + for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) shr[gid] = sum; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; + +#if !defined(__STOCHASTIC_MODE__) || __CUDA_ARCH__ < 700 + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + sum += g.shfl_down(sum, i); + } + + sum = g.shfl(sum, 0); + float mean = sum / row_stride; + float variance = 0.f; + for (int i = 0; i < iterations; i++) { + vals_arr[i] -= mean; + variance += vals_arr[i] * vals_arr[i]; + } + + for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } + + if (g.thread_rank() == 0) shr[gid] = variance; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + variance += g.shfl_down(variance, i); + } + variance = g.shfl(variance, 0); + variance /= row_stride; + variance += epsilon; + if (training) + if (threadIdx.x == 0) vars[row] = variance; + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + vals_arr[i] = vals_arr[i] * rsqrtf(variance); + vals_arr[i] = + vals_arr[i] * gamma[i * iteration_stride + id] + beta[i * iteration_stride + id]; + vals[i * iteration_stride + id] = vals_arr[i]; + } + if ((high_index) < row_stride) { + vals_arr[iterations] = vals_arr[iterations] * rsqrtf(variance); + vals_arr[iterations] = vals_arr[iterations] * gamma[high_index] + beta[high_index]; + vals[high_index] = vals_arr[iterations]; + } +} + +__global__ void fused_bias_residual_layer_norm(__half* vals, + const __half* residual, + const __half* gamma, + const __half* beta, + float epsilon, + bool preLayerNorm, + bool training, + __half* vars, + int row_stride) +{ +#ifdef HALF_PRECISION_AVAILABLE + + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int gid = id >> WARP_SIZE_BITS; + + float2 vals_f[NORM_REG]; + __shared__ float shr[MAX_WARP_NUM]; + + __half2* vals_cast = reinterpret_cast<__half2*>(vals); + const __half2* residual_cast = reinterpret_cast(residual); + + residual_cast += (row * row_stride); + vals_cast += (row * row_stride); + + float sum = 0.f; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + vals_f[i] = __half22float2(residual_cast[i * iteration_stride + id]); + sum += vals_f[i].x; + sum += vals_f[i].y; + } + if ((high_index) < row_stride) { + vals_f[iterations] = __half22float2(residual_cast[high_index]); + sum += vals_f[iterations].x; + sum += vals_f[iterations].y; + iterations++; + } + + for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) shr[gid] = sum; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + sum += g.shfl_down(sum, i); + } + sum = g.shfl(sum, 0); + float mean = sum / (row_stride * 2); + + float variance = 0.f; + for (int i = 0; i < iterations; i++) { + vals_f[i].x -= mean; + vals_f[i].y -= mean; + variance += vals_f[i].x * vals_f[i].x; + variance += vals_f[i].y * vals_f[i].y; + } + + for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } + + if (g.thread_rank() == 0) shr[gid] = variance; + + b.sync(); + + if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { + variance += g.shfl_down(variance, i); + } + variance = g.shfl(variance, 0); + variance /= (row_stride * 2); + variance += epsilon; + + __half2 variance_h = __float2half2_rn(variance); + const __half2* gamma_cast = reinterpret_cast(gamma); + const __half2* beta_cast = reinterpret_cast(beta); + + if (training && threadIdx.x == 0) vars[row] = __float2half(variance); + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + __half2 vals_arr = __float22half2_rn(vals_f[i]); + vals_arr = vals_arr * h2rsqrt(variance_h); + vals_arr = + vals_arr * gamma_cast[i * iteration_stride + id] + beta_cast[i * iteration_stride + id]; + vals_cast[i * iteration_stride + id] = vals_arr; + } + if ((high_index) < row_stride) { + __half2 vals_arr = __float22half2_rn(vals_f[iterations]); + vals_arr = vals_arr * h2rsqrt(variance_h); + vals_arr = vals_arr * gamma_cast[high_index] + beta_cast[high_index]; + vals_cast[high_index] = vals_arr; + } +#endif +} + +template +void launch_bias_residual_layer_norm(T* vals, + const T* residual, + const T* gamma, + const T* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + T* vars); + +/* +To tune this launch the following restrictions must be met: + +For float: +row_stride == hidden_size +threads * iterations == row_stride +threads is in [32, 64, 128, 256, 512, 1024] + +For half: +row_stride == hidden_size / 2 +threads * iterations == row_stride +threads is in [32, 64, 128, 256, 512, 1024] + +*/ + +template <> +void launch_bias_residual_layer_norm(float* vals, + const float* residual, + const float* gamma, + const float* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + float* vars) +{ + int threads = THREADS; + + dim3 grid_dim(batch_size); + + // There are some limitations to call below functions, now just enumerate the situations. + + if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 1; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 2; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim(threads); + + fused_bias_residual_layer_norm<<>>( + vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, hidden_dim); +} + +template <> +void launch_bias_residual_layer_norm<__half>(__half* vals, + const __half* residual, + const __half* gamma, + const __half* beta, + float epsilon, + int batch_size, + int hidden_dim, + cudaStream_t stream, + bool preLayerNorm, + bool training, + __half* vars) +{ + int threads = 128; + + dim3 grid_dim(batch_size); + + // There are some limitations to call below functions, now just enumerate the situations. + + if (hidden_dim > 8192 && hidden_dim <= 16384) + threads <<= 1; + else if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 2; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 3; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim(threads); + fused_bias_residual_layer_norm<<>>( + vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, hidden_dim / 2); +} + +/* Normalize Gamma & Betta gradients + * Compute gradients using either X_hat or + * normalize input (invertible). + * Combine transpose with gradients computation. + */ + +template +__global__ void LayerNormBackward1(const T* __restrict__ out_grad, + const T* __restrict__ vals_hat, + const T* __restrict__ gamma, + const T* __restrict__ betta, + T* __restrict__ gamma_grad, + T* __restrict__ betta_grad, + int rows, + int width, + bool invertible) +{ + __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; + __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int offset = threadIdx.y * width + idx; + int y_stride = width * TILE_DIM; + + float betta_reg = (invertible ? (float)betta[idx] : 0.0f); + float gamma_reg = (float)gamma[idx]; + + // Loop across matrix height + float betta_tmp = 0; + float gamma_tmp = 0; + for (int r = threadIdx.y; r < rows; r += TILE_DIM) { + float grad = (float)out_grad[offset]; + float val = (invertible ? ((float)vals_hat[offset] - betta_reg) / gamma_reg + : (float)vals_hat[offset]); + betta_tmp += grad; + gamma_tmp += (val * grad); + + offset += y_stride; + } + + betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; + gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; + + __syncthreads(); + + // Sum the shared buffer. + float s1 = betta_buffer[threadIdx.y][threadIdx.x]; + float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < TILE_DIM; i <<= 1) { + s1 += g.shfl_down(s1, i); + s2 += g.shfl_down(s2, i); + } + + if (threadIdx.x == 0) { + int pos = blockIdx.x * TILE_DIM + threadIdx.y; + betta_grad[pos] = s1; + gamma_grad[pos] = s2; + } +} + +/* Normalize Gamma & Betta gradients + * Compute gradients using the input to + * the normalize. + * Combine transpose with gradients computation. + */ + +template +__global__ void LayerNormBackward1(const T* __restrict__ out_grad, + const T* __restrict__ X_data, + const T* __restrict__ vars, + const T* __restrict__ means, + T* __restrict__ gamma_grad, + T* __restrict__ betta_grad, + int rows, + int width) +{ + __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; + __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int offset = threadIdx.y * width + idx; + int y_stride = width * TILE_DIM; + + int pos = blockIdx.x * TILE_DIM + threadIdx.y; + // Loop across matrix height + + float betta_tmp = 0; + float gamma_tmp = 0; + for (int r = threadIdx.y; r < rows; r += TILE_DIM) { + float grad = (float)out_grad[offset]; + float val = (float)X_data[offset]; + val = (val - (float)means[r]) * rsqrtf((float)vars[r]); + betta_tmp += grad; + gamma_tmp += (val * grad); + + offset += y_stride; + } + + betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; + gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; + + __syncthreads(); + + // Sum the shared buffer. + float s1 = betta_buffer[threadIdx.y][threadIdx.x]; + float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < TILE_DIM; i <<= 1) { + s1 += g.shfl_down(s1, i); + s2 += g.shfl_down(s2, i); + } + + if (threadIdx.x == 0) { + betta_grad[pos] = s1; + gamma_grad[pos] = s2; + } +} +/* + +/* Backward Normalize (Input-Gradient) + * Using the means and variances from the input + * This type of backward is invertible! + * We do the backward using the X_hat (X - u) / sqrt(variance) or the output of Normalization. + */ + +__global__ void LayerNormBackward2(const float* out_grad, + const float* vals_hat, + const float* gamma, + const float* betta, + const float* vars, + float* inp_grad, + bool invertible, + int row_stride) +{ + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id / WARP_SIZE; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + __shared__ float partialSum[MAX_WARP_NUM]; + + out_grad += (row * row_stride); + vals_hat += (row * row_stride); + inp_grad += (row * row_stride); + + float vals_arr[NORM_REG]; + float vals_hat_arr[NORM_REG]; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + float gamma_reg = gamma[i * iteration_stride + id]; + vals_arr[i] = out_grad[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; + vals_hat_arr[i] = + (invertible ? (vals_hat[i * iteration_stride + id] - betta[i * iteration_stride + id]) / + gamma_reg + : vals_hat[i * iteration_stride + id]); + } + if ((high_index) < row_stride) { + float gamma_reg = gamma[high_index]; + vals_arr[iterations] = out_grad[high_index]; + vals_arr[iterations] *= gamma_reg; + vals_hat_arr[iterations] = + (invertible ? (vals_hat[high_index] - betta[high_index]) / gamma_reg + : vals_hat[high_index]); + iterations++; + } + + float var_reg = vars[row]; + + float sum = 0; + for (int i = 0; i < iterations; i++) { + sum += vals_hat_arr[i] * vals_arr[i] * + sqrtf(var_reg); // dval_hat = gamma * (x - u) * out_grad + vals_arr[i] *= rsqrtf(var_reg); // dvar_inv = gamma * out_grad / sqrt(var) + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= row_stride; + + for (int i = 0; i < iterations; i++) { vals_arr[i] += ((-sum * vals_hat_arr[i]) / var_reg); } + + sum = 0; + for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + sum = g.shfl(sum, 0); + sum /= row_stride; + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) inp_grad[i * iteration_stride + id] = (vals_arr[i] - sum); + if ((high_index) < row_stride) inp_grad[high_index] = (vals_arr[iterations] - sum); +} + +__global__ void LayerNormBackward2(const __half* out_grad, + const __half* vals_hat, + const __half* gamma, + const __half* betta, + const __half* vars, + __half* inp_grad, + bool invertible, + int row_stride) +{ +#ifdef HALF_PRECISION_AVAILABLE + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id / WARP_SIZE; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + __shared__ float partialSum[MAX_WARP_NUM]; + + __half2 vals_arr[NORM_REG]; + float2 vals_arr_f[NORM_REG]; + __half2 vals_hat_arr[NORM_REG]; + + __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); + const __half2* out_grad_h = reinterpret_cast(out_grad); + const __half2* vals_hat_h = reinterpret_cast(vals_hat); + + inp_grad_h += (row * row_stride); + out_grad_h += (row * row_stride); + vals_hat_h += (row * row_stride); + + const __half2* gamma_h = reinterpret_cast(gamma); + const __half2* betta_h = (invertible ? reinterpret_cast(betta) : nullptr); + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + __half2 gamma_reg = gamma_h[i * iteration_stride + id]; + vals_arr[i] = out_grad_h[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; + vals_hat_arr[i] = + (invertible + ? (vals_hat_h[i * iteration_stride + id] - betta_h[i * iteration_stride + id]) / + gamma_reg + : vals_hat_h[i * iteration_stride + id]); + } + if ((high_index) < row_stride) { + __half2 gamma_reg = gamma_h[high_index]; + vals_arr[iterations] = out_grad_h[high_index]; + vals_arr[iterations] *= gamma_reg; + vals_hat_arr[iterations] = + (invertible ? (vals_hat_h[high_index] - betta_h[high_index]) / gamma_reg + : vals_hat_h[high_index]); + iterations++; + } + __half var_h = vars[row]; + __half2 var_reg = __halves2half2(var_h, var_h); + + float sum = 0.f; + for (int i = 0; i < iterations; i++) { + __half2 result_h = (vals_hat_arr[i] * vals_arr[i] * h2sqrt(var_reg)); + float2 result_f = __half22float2(result_h); + sum += result_f.x; + sum += result_f.y; + vals_arr[i] *= h2rsqrt(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + __half2 sum_h = __float2half2_rn(sum); + + for (int i = 0; i < iterations; i++) { + __half2 temp = ((-sum_h * vals_hat_arr[i]) / (var_reg)); + vals_arr_f[i] = __half22float2(vals_arr[i]); + float2 temp_f = __half22float2(temp); + vals_arr_f[i].x += temp_f.x; + vals_arr_f[i].y += temp_f.y; + } + sum = 0.f; + + for (int i = 0; i < iterations; i++) { + sum += (vals_arr_f[i].x); + sum += (vals_arr_f[i].y); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + vals_arr_f[i].x -= sum; + vals_arr_f[i].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[i]); + + inp_grad_h[i * iteration_stride + id] = temp; + } + if ((high_index) < row_stride) { + vals_arr_f[iterations].x -= sum; + vals_arr_f[iterations].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[iterations]); + + inp_grad_h[high_index] = temp; + } +#endif +} + +template <> +void launch_layerNorm_backward(const float* out_grad, + const float* vals_hat, + const float* vars, + const float* gamma, + float* gamma_grad, + float* betta_grad, + float* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2], + bool invertible, + const float* betta) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + LayerNormBackward1<<>>( + out_grad, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); + + dim3 grid_dim2(batch); + + if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 1; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 2; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads); + + LayerNormBackward2<<>>( + out_grad, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim); +} + +template <> +void launch_layerNorm_backward<__half>(const __half* out_grad, + const __half* vals_hat, + const __half* vars, + const __half* gamma, + __half* gamma_grad, + __half* betta_grad, + __half* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2], + bool invertible, + const __half* betta) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + // LayerNormBackward1<__half><<>>( + // out_grad, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); + + dim3 grid_dim2(batch); + + if (hidden_dim > 8192 && hidden_dim <= 16384) + threads <<= 1; + else if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 2; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 3; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads / 2); + + LayerNormBackward2<<>>( + out_grad, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim / 2); +} + +/* Backward Normalize (Input-Gradient) + * Using the means and variances from the input + * This type of backward is not invertible! + * We do the backward using the input (X) + */ + +__global__ void LayerNormBackward2(const float* out_grad, + const float* X_vals, + const float* gamma, + const float* vars, + const float* means, + float* inp_grad, + int row_stride) +{ + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id >> WARP_SIZE_BITS; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + __shared__ float partialSum[MAX_WARP_NUM]; + + out_grad += (row * row_stride); + X_vals += (row * row_stride); + inp_grad += (row * row_stride); + + float vals_arr[NORM_REG]; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + float gamma_reg = gamma[i * iteration_stride + id]; + vals_arr[i] = out_grad[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; + } + if ((high_index) < row_stride) { + float gamma_reg = gamma[high_index]; + vals_arr[iterations] = out_grad[high_index]; + vals_arr[iterations] *= gamma_reg; + iterations++; + } + + float var_reg = vars[row]; + float mean_reg = means[row]; + + float sum = 0; + float xu[NORM_REG]; + for (int i = 0; i < iterations; i++) { + xu[i] = (X_vals[i * iteration_stride + id] - mean_reg); + sum += vals_arr[i] * xu[i]; + vals_arr[i] *= rsqrtf(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= row_stride; + + for (int i = 0; i < iterations; i++) { + vals_arr[i] += (-sum * xu[i] * rsqrtf(var_reg) / (var_reg)); + } + + sum = 0; + for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + sum = g.shfl(sum, 0); + sum /= row_stride; + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) inp_grad[i * iteration_stride + id] = (vals_arr[i] - sum); + if ((high_index) < row_stride) inp_grad[high_index] = (vals_arr[iterations] - sum); +} + +__global__ void LayerNormBackward2(const __half* out_grad, + const __half* X_vals, + const __half* gamma, + const __half* vars, + const __half* means, + __half* inp_grad, + int row_stride) +{ +#ifdef HALF_PRECISION_AVAILABLE + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id >> WARP_SIZE_BITS; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + + __shared__ float partialSum[MAX_WARP_NUM]; + + __half2 vals_arr[NORM_REG]; + float2 vals_arr_f[NORM_REG]; + __half2 xu[NORM_REG]; + + __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); + const __half2* out_grad_h = reinterpret_cast(out_grad); + const __half2* vals_hat_h = reinterpret_cast(X_vals); + + inp_grad_h += (row * row_stride); + out_grad_h += (row * row_stride); + vals_hat_h += (row * row_stride); + + const __half2* gamma_h = reinterpret_cast(gamma); + int high_index = iterations * iteration_stride + id; + + __half mean_h = means[row]; + __half2 mean_reg = __halves2half2(mean_h, mean_h); +#pragma unroll + for (int i = 0; i < iterations; i++) { + __half2 gamma_reg = gamma_h[i * iteration_stride + id]; + vals_arr[i] = out_grad_h[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; // out_grad * gamma + xu[i] = (vals_hat_h[i * iteration_stride + id] - mean_reg); + } + if ((high_index) < row_stride) { + __half2 gamma_reg = gamma_h[high_index]; + vals_arr[iterations] = out_grad_h[high_index]; + vals_arr[iterations] *= gamma_reg; // out_grad * gamma + xu[iterations] = (vals_hat_h[high_index] - mean_reg); + iterations++; + } + __half var_h = vars[row]; + __half2 var_reg = __halves2half2(var_h, var_h); + + float sum = 0.f; + for (int i = 0; i < iterations; i++) { + __half2 result_h = (xu[i] * vals_arr[i]); + float2 result_f = __half22float2(result_h); + sum += result_f.x; + sum += result_f.y; + vals_arr[i] *= h2rsqrt(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + __half2 sum_h = __float2half2_rn(sum); + + for (int i = 0; i < iterations; i++) { + __half2 xu_grad = ((-sum_h * xu[i] * h2rsqrt(var_reg)) / (var_reg)); + vals_arr_f[i] = __half22float2(vals_arr[i]); + float2 xu_grad_f = __half22float2(xu_grad); + vals_arr_f[i].x += xu_grad_f.x; + vals_arr_f[i].y += xu_grad_f.y; + } + + sum = 0.f; + for (int i = 0; i < iterations; i++) { + sum += (vals_arr_f[i].x); + sum += (vals_arr_f[i].y); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + vals_arr_f[i].x -= sum; + vals_arr_f[i].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[i]); + inp_grad_h[i * iteration_stride + id] = temp; + } + if ((high_index) < row_stride) { + vals_arr_f[iterations].x -= sum; + vals_arr_f[iterations].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[iterations]); + inp_grad_h[high_index] = temp; + } +#endif +} + +template <> +void launch_layerNorm_backward(const float* out_grad, + const float* X_data, + const float* vars, + const float* means, + const float* gamma, + float* gamma_grad, + float* betta_grad, + float* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2]) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + LayerNormBackward1<<>>( + out_grad, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); + + dim3 grid_dim2(batch); + + if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 1; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 2; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads); + LayerNormBackward2<<>>( + out_grad, X_data, gamma, vars, means, inp_grad, hidden_dim); +} + +template <> +void launch_layerNorm_backward<__half>(const __half* out_grad, + const __half* X_data, + const __half* vars, + const __half* means, + const __half* gamma, + __half* gamma_grad, + __half* betta_grad, + __half* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2]) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + LayerNormBackward1<__half><<>>( + out_grad, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); + + dim3 grid_dim2(batch); + + if (hidden_dim > 8192 && hidden_dim <= 16384) + threads <<= 1; + else if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 2; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 3; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads / 2); + LayerNormBackward2<<>>( + out_grad, X_data, gamma, vars, means, inp_grad, hidden_dim / 2); +} + +template +__global__ void LayerNormBackward1_fused_add(const T* __restrict__ out_grad1, + const T* __restrict__ out_grad2, + const T* __restrict__ vals_hat, + const T* __restrict__ gamma, + const T* __restrict__ betta, + T* __restrict__ gamma_grad, + T* __restrict__ betta_grad, + int rows, + int width, + bool invertible) +{ + __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; + __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int offset = threadIdx.y * width + idx; + int y_stride = width * TILE_DIM; + + float betta_reg = (invertible ? (float)betta[idx] : 0.0f); + float gamma_reg = (float)gamma[idx]; + + // Loop across matrix height + float betta_tmp = 0; + float gamma_tmp = 0; + for (int r = threadIdx.y; r < rows; r += TILE_DIM) { + float grad = (float)out_grad1[offset] + (float)out_grad2[offset]; + float val = (invertible ? ((float)vals_hat[offset] - betta_reg) / gamma_reg + : (float)vals_hat[offset]); + betta_tmp += grad; + gamma_tmp += (val * grad); + + offset += y_stride; + } + + betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; + gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; + + __syncthreads(); + + // Sum the shared buffer. + float s1 = betta_buffer[threadIdx.y][threadIdx.x]; + float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < TILE_DIM; i <<= 1) { + s1 += g.shfl_down(s1, i); + s2 += g.shfl_down(s2, i); + } + + if (threadIdx.x == 0) { + int pos = blockIdx.x * TILE_DIM + threadIdx.y; + betta_grad[pos] = s1; + gamma_grad[pos] = s2; + } +} + +template +__global__ void LayerNormBackward1_fused_add(const T* __restrict__ out_grad1, + const T* __restrict__ out_grad2, + const T* __restrict__ X_data, + const T* __restrict__ vars, + const T* __restrict__ means, + T* __restrict__ gamma_grad, + T* __restrict__ betta_grad, + int rows, + int width) +{ + __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; + __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int offset = threadIdx.y * width + idx; + int y_stride = width * TILE_DIM; + + int pos = blockIdx.x * TILE_DIM + threadIdx.y; + // Loop across matrix height + + float betta_tmp = 0; + float gamma_tmp = 0; + for (int r = threadIdx.y; r < rows; r += TILE_DIM) { + float grad = (float)out_grad1[offset] + (float)out_grad2[offset]; + float val = (float)X_data[offset]; + val = (val - (float)means[r]) * rsqrtf((float)vars[r]); + betta_tmp += grad; + gamma_tmp += (val * grad); + + offset += y_stride; + } + + betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; + gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; + + __syncthreads(); + + // Sum the shared buffer. + float s1 = betta_buffer[threadIdx.y][threadIdx.x]; + float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < TILE_DIM; i <<= 1) { + s1 += g.shfl_down(s1, i); + s2 += g.shfl_down(s2, i); + } + + if (threadIdx.x == 0) { + betta_grad[pos] = s1; + gamma_grad[pos] = s2; + } +} + +__global__ void LayerNormBackward2_fused_add(const float* out_grad1, + const float* out_grad2, + const float* vals_hat, + const float* gamma, + const float* betta, + const float* vars, + float* inp_grad, + bool invertible, + int row_stride) +{ + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id / WARP_SIZE; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + __shared__ float partialSum[MAX_WARP_NUM]; + + out_grad1 += (row * row_stride); + out_grad2 += (row * row_stride); + vals_hat += (row * row_stride); + inp_grad += (row * row_stride); + + float vals_arr[NORM_REG]; + float vals_hat_arr[NORM_REG]; + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + float gamma_reg = gamma[i * iteration_stride + id]; + vals_arr[i] = out_grad1[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; + vals_hat_arr[i] = + (invertible ? (vals_hat[i * iteration_stride + id] - betta[i * iteration_stride + id]) / + gamma_reg + : vals_hat[i * iteration_stride + id]); + } + if ((high_index) < row_stride) { + float gamma_reg = gamma[high_index]; + vals_arr[iterations] = out_grad1[high_index]; + vals_arr[iterations] *= gamma_reg; + vals_hat_arr[iterations] = + (invertible ? (vals_hat[high_index] - betta[high_index]) / gamma_reg + : vals_hat[high_index]); + iterations++; + } + + float var_reg = vars[row]; + + float sum = 0; + for (int i = 0; i < iterations; i++) { + sum += vals_hat_arr[i] * vals_arr[i] * sqrtf(var_reg); + vals_arr[i] *= rsqrtf(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= row_stride; + + for (int i = 0; i < iterations; i++) { vals_arr[i] += ((-sum * vals_hat_arr[i]) / var_reg); } + + sum = 0; + for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + sum = g.shfl(sum, 0); + sum /= row_stride; + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) + inp_grad[i * iteration_stride + id] = + (vals_arr[i] - sum) + out_grad2[i * iteration_stride + id]; + if ((high_index) < row_stride) + inp_grad[high_index] = (vals_arr[iterations] - sum) + out_grad2[high_index]; +} + +__global__ void LayerNormBackward2_fused_add(const __half* out_grad1, + const __half* out_grad2, + const __half* vals_hat, + const __half* gamma, + const __half* betta, + const __half* vars, + __half* inp_grad, + bool invertible, + int row_stride) +{ +#ifdef HALF_PRECISION_AVAILABLE + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id / WARP_SIZE; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + __shared__ float partialSum[MAX_WARP_NUM]; + + __half2 vals_arr[NORM_REG]; + float2 vals_arr_f[NORM_REG]; + __half2 vals_hat_arr[NORM_REG]; + + // float2 result[iterations]; + + __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); + const __half2* out_grad_h1 = reinterpret_cast(out_grad1); + const __half2* out_grad_h2 = reinterpret_cast(out_grad2); + const __half2* vals_hat_h = reinterpret_cast(vals_hat); + + inp_grad_h += (row * row_stride); + out_grad_h1 += (row * row_stride); + out_grad_h2 += (row * row_stride); + vals_hat_h += (row * row_stride); + + const __half2* gamma_h = reinterpret_cast(gamma); + const __half2* betta_h = (invertible ? reinterpret_cast(betta) : nullptr); + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + __half2 gamma_reg = gamma_h[i * iteration_stride + id]; + vals_arr[i] = out_grad_h1[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; // out_grad * gamma + vals_hat_arr[i] = + (invertible + ? (vals_hat_h[i * iteration_stride + id] - betta_h[i * iteration_stride + id]) / + gamma_reg + : vals_hat_h[i * iteration_stride + id]); + } + if ((high_index) < row_stride) { + __half2 gamma_reg = gamma_h[high_index]; + vals_arr[iterations] = out_grad_h1[high_index]; + vals_arr[iterations] *= gamma_reg; // out_grad * gamma + vals_hat_arr[iterations] = + (invertible ? (vals_hat_h[high_index] - betta_h[high_index]) / gamma_reg + : vals_hat_h[high_index]); + iterations++; + } + __half var_h = vars[row]; + __half2 var_reg = __halves2half2(var_h, var_h); + + float sum = 0.f; + for (int i = 0; i < iterations; i++) { + __half2 result_h = (vals_hat_arr[i] * vals_arr[i] * h2sqrt(var_reg)); + float2 result_f = __half22float2(result_h); + sum += result_f.x; + sum += result_f.y; + vals_arr[i] *= h2rsqrt(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + __half2 sum_h = __float2half2_rn(sum); + + for (int i = 0; i < iterations; i++) { + __half2 temp = ((-sum_h * vals_hat_arr[i]) / (var_reg)); + vals_arr_f[i] = __half22float2(vals_arr[i]); + float2 temp_f = __half22float2(temp); + vals_arr_f[i].x += temp_f.x; + vals_arr_f[i].y += temp_f.y; + } + sum = 0.f; + for (int i = 0; i < iterations; i++) { + sum += (vals_arr_f[i].x); + sum += (vals_arr_f[i].y); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + vals_arr_f[i].x -= sum; + vals_arr_f[i].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[i]); + + inp_grad_h[i * iteration_stride + id] = temp + out_grad_h2[i * iteration_stride + id]; + } + if ((high_index) < row_stride) { + vals_arr_f[iterations].x -= sum; + vals_arr_f[iterations].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[iterations]); + + inp_grad_h[high_index] = temp + out_grad_h2[high_index]; + } +#endif +} + +template <> +void launch_layerNorm_backward_fused_add(const float* out_grad1, + const float* out_grad2, + const float* vals_hat, + const float* vars, + const float* gamma, + float* gamma_grad, + float* betta_grad, + float* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2], + bool invertible, + const float* betta) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + LayerNormBackward1<<>>( + out_grad1, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); + + dim3 grid_dim2(batch); + + if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 1; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 2; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads); + LayerNormBackward2_fused_add<<>>( + out_grad1, out_grad2, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim); +} + +template <> +void launch_layerNorm_backward_fused_add<__half>(const __half* out_grad1, + const __half* out_grad2, + const __half* vals_hat, + const __half* vars, + const __half* gamma, + __half* gamma_grad, + __half* betta_grad, + __half* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2], + bool invertible, + const __half* betta) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + LayerNormBackward1<__half><<>>( + out_grad1, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); + + dim3 grid_dim2(batch); + + if (hidden_dim > 8192 && hidden_dim <= 16384) + threads <<= 1; + else if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 2; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 3; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads / 2); + LayerNormBackward2_fused_add<<>>( + out_grad1, out_grad2, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim / 2); +} + +/* Backward Normalize (Input-Gradient) + * Using the means and variances from the input + * This type of backward is not invertible! + * We do the backward using the input (X) + */ + +__global__ void LayerNormBackward2_fused_add(const float* out_grad1, + const float* out_grad2, + const float* X_vals, + const float* gamma, + const float* vars, + const float* means, + float* inp_grad, + int row_stride) +{ + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id / WARP_SIZE; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + __shared__ float partialSum[MAX_WARP_NUM]; + + float vals_arr[NORM_REG]; + float vals_hat_arr[NORM_REG]; + + out_grad1 += (row * row_stride); + out_grad2 += (row * row_stride); + X_vals += (row * row_stride); + inp_grad += (row * row_stride); + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + float gamma_reg = gamma[i * iteration_stride + id]; + vals_arr[i] = out_grad1[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; + vals_hat_arr[i] = X_vals[i * iteration_stride + id]; + } + if ((high_index) < row_stride) { + float gamma_reg = gamma[high_index]; + vals_arr[iterations] = out_grad1[high_index]; + vals_arr[iterations] *= gamma_reg; + vals_hat_arr[iterations] = X_vals[high_index]; + iterations++; + } + + float var_reg = vars[row]; + float mean_reg = means[row]; + + float sum = 0; + float xu[NORM_REG]; + for (int i = 0; i < iterations; i++) { + xu[i] = (vals_hat_arr[i] - mean_reg); + sum += vals_arr[i] * xu[i]; + vals_arr[i] *= rsqrtf(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= row_stride; + + for (int i = 0; i < iterations; i++) { + vals_arr[i] += (-sum * xu[i] * rsqrtf(var_reg) / (var_reg)); + } + + sum = 0; + for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + sum = g.shfl(sum, 0); + sum /= row_stride; + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) + inp_grad[i * iteration_stride + id] = + (vals_arr[i] - sum) + out_grad2[i * iteration_stride + id]; + if ((high_index) < row_stride) + inp_grad[high_index] = (vals_arr[iterations] - sum) + out_grad2[high_index]; +} + +__global__ void LayerNormBackward2_fused_add(const __half* out_grad1, + const __half* out_grad2, + const __half* X_vals, + const __half* gamma, + const __half* vars, + const __half* means, + __half* inp_grad, + int row_stride) +{ +#ifdef HALF_PRECISION_AVAILABLE + int iteration_stride = blockDim.x; + int iterations = row_stride / iteration_stride; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + int wid = id / WARP_SIZE; + int warp_num = iteration_stride >> WARP_SIZE_BITS; + + __shared__ float partialSum[MAX_WARP_NUM]; + + __half2 vals_arr[NORM_REG]; + float2 vals_arr_f[NORM_REG]; + __half2 vals_hat_arr[NORM_REG]; + + __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); + const __half2* out_grad_h1 = reinterpret_cast(out_grad1); + const __half2* out_grad_h2 = reinterpret_cast(out_grad2); + const __half2* vals_hat_h = reinterpret_cast(X_vals); + + out_grad_h1 += (row * row_stride); + out_grad_h2 += (row * row_stride); + inp_grad_h += (row * row_stride); + vals_hat_h += (row * row_stride); + + const __half2* gamma_h = reinterpret_cast(gamma); + int high_index = iterations * iteration_stride + id; +#pragma unroll + for (int i = 0; i < iterations; i++) { + __half2 gamma_reg = gamma_h[i * iteration_stride + id]; + vals_arr[i] = out_grad_h1[i * iteration_stride + id]; + vals_arr[i] *= gamma_reg; // out_grad * gamma + vals_hat_arr[i] = vals_hat_h[i * iteration_stride + id]; + } + if ((high_index) < row_stride) { + __half2 gamma_reg = gamma_h[high_index]; + vals_arr[iterations] = out_grad_h1[high_index]; + vals_arr[iterations] *= gamma_reg; // out_grad * gamma + vals_hat_arr[iterations] = vals_hat_h[high_index]; + iterations++; + } + + __half mean_h = means[row]; + __half var_h = vars[row]; + __half2 var_reg = __halves2half2(var_h, var_h); + __half2 mean_reg = __halves2half2(mean_h, mean_h); + __half2 xu[NORM_REG]; + + float sum = 0.f; + for (int i = 0; i < iterations; i++) { + xu[i] = (vals_hat_arr[i] - mean_reg); + __half2 result_h = (xu[i] * vals_arr[i]); + float2 result_f = __half22float2(result_h); + sum += result_f.x; + sum += result_f.y; + vals_arr[i] *= h2rsqrt(var_reg); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + __half2 sum_h = __float2half2_rn(sum); + + for (int i = 0; i < iterations; i++) { + __half2 xu_grad = ((-sum_h * xu[i] * h2rsqrt(var_reg)) / (var_reg)); + vals_arr_f[i] = __half22float2(vals_arr[i]); + float2 xu_grad_f = __half22float2(xu_grad); + vals_arr_f[i].x += xu_grad_f.x; + vals_arr_f[i].y += xu_grad_f.y; + } + + sum = 0.f; + for (int i = 0; i < iterations; i++) { + sum += (vals_arr_f[i].x); + sum += (vals_arr_f[i].y); + } + + for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } + + if (g.thread_rank() == 0) partialSum[wid] = sum; + + __syncthreads(); + + if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; + +#ifndef __STOCHASTIC_MODE__ + __syncthreads(); +#endif + + for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); + + sum = g.shfl(sum, 0); + sum /= (2 * row_stride); + + iterations = row_stride / iteration_stride; + for (int i = 0; i < iterations; i++) { + vals_arr_f[i].x -= sum; + vals_arr_f[i].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[i]); + inp_grad_h[i * iteration_stride + id] = temp + out_grad_h2[i * iteration_stride + id]; + } + if ((high_index) < row_stride) { + vals_arr_f[iterations].x -= sum; + vals_arr_f[iterations].y -= sum; + __half2 temp = __float22half2_rn(vals_arr_f[iterations]); + inp_grad_h[high_index] = temp + out_grad_h2[high_index]; + } +#endif +} + +template <> +void launch_layerNorm_backward_fused_add(const float* out_grad1, + const float* out_grad2, + const float* X_data, + const float* vars, + const float* means, + const float* gamma, + float* gamma_grad, + float* betta_grad, + float* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2]) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + LayerNormBackward1<<>>( + out_grad1, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); + + dim3 grid_dim2(batch); + + if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 1; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 2; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads); + LayerNormBackward2_fused_add<<>>( + out_grad1, out_grad2, X_data, gamma, vars, means, inp_grad, hidden_dim); +} + +template <> +void launch_layerNorm_backward_fused_add<__half>(const __half* out_grad1, + const __half* out_grad2, + const __half* X_data, + const __half* vars, + const __half* means, + const __half* gamma, + __half* gamma_grad, + __half* betta_grad, + __half* inp_grad, + int batch, + int hidden_dim, + cudaStream_t stream[2]) +{ + int threads = THREADS; + + dim3 grid_dim(hidden_dim / TILE_DIM); + dim3 block_dim(TILE_DIM, TILE_DIM); + + LayerNormBackward1<__half><<>>( + out_grad1, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); + + dim3 grid_dim2(batch); + + if (hidden_dim > 8192 && hidden_dim <= 16384) + threads <<= 1; + else if (hidden_dim > 16384 && hidden_dim <= 32768) + threads <<= 2; + else if (hidden_dim > 32768 && hidden_dim <= 65536) + threads <<= 3; + else if (hidden_dim > 65536) + throw std::runtime_error("Unsupport hidden_dim."); + + dim3 block_dim2(threads / 2); + LayerNormBackward2_fused_add<<>>( + out_grad1, out_grad2, X_data, gamma, vars, means, inp_grad, hidden_dim / 2); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/softmax_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/softmax_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..d9efed286462f1be971922c2345efb8d2049e227 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/softmax_kernels.cu @@ -0,0 +1,701 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "custom_cuda_layers.h" +#include "general_kernels.h" + +namespace cg = cooperative_groups; + +dim3 get_attn_softmax_grid(int batch_size, int heads, int sequence_length, int threads) +{ + int seq_length4 = sequence_length / 4; + int block_compute_size = + (seq_length4 < threads ? (int)pow(2.0, floor(log2((float)(threads / seq_length4)))) : 1); + // Note that the Y and Z dimensions are limited to 65535, while X is basically unlimited: + // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications + // The batch size is typically relatively small, while the sequence length could potentially be + // arbitrarily large. We therefore place the batch size second to avoid hitting the Y limit. + unsigned x = heads * sequence_length / block_compute_size; + unsigned y = batch_size; + return {x, y}; +} + +// Fused attention + softmax +template +__global__ void attn_softmax(float* vals, + const float* attn_mask, + int heads, + int seq_length, + int iterations) +{ + __shared__ float partialSum[MAX_WARP_NUM]; + + int warp_num = blockDim.x >> WARP_SIZE_BITS; + + int iteration_stride = blockDim.x; + int block_width = blockStride * seq_length; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int batch = blockIdx.y; + int row = blockIdx.x; + int max_threads_in_sequence = std::max(seq_length, tbSeq); + int seq_lane = threadIdx.x % max_threads_in_sequence; + + int data_offset = batch * (gridDim.x * block_width) + row * block_width + + (threadIdx.x / max_threads_in_sequence) * seq_length; + int mask_offset = batch * seq_length; + + int wid = threadIdx.x >> WARP_SIZE_BITS; + int lane = threadIdx.x & 0x1f; + + float4* val_cast = reinterpret_cast(vals); + const float4* attn_mask_cast = reinterpret_cast(attn_mask); + + float4 data[MAX_THREAD_ITERATIONS]; + + float max_val = minus_infinity; + + for (int i = 0; i < iterations; i++) { + int data_id = i * iteration_stride + seq_lane; + if (data_id < seq_length) { + float4 mask = attn_mask_cast[mask_offset + data_id]; + data[i] = val_cast[data_offset + data_id]; + + data[i].x += mask.x; + data[i].y += mask.y; + data[i].z += mask.z; + data[i].w += mask.w; + + max_val = (data[i].x > max_val ? data[i].x : max_val); + max_val = (data[i].y > max_val ? data[i].y : max_val); + max_val = (data[i].z > max_val ? data[i].z : max_val); + max_val = (data[i].w > max_val ? data[i].w : max_val); + } else { + data[i].x = minus_infinity; + data[i].y = minus_infinity; + data[i].z = minus_infinity; + data[i].w = minus_infinity; + } + } + + for (int i = 1; i < tbSize; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + if (seq_length > tbSize) { + if (lane == 0) partialSum[wid] = max_val; + b.sync(); + + if (lane < warp_num) max_val = partialSum[lane]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + int iters = warp_num; + if (seq_length < iteration_stride) + iters = warp_num / (iteration_stride / max_threads_in_sequence); + + for (int i = 1; i < iters; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + max_val = g.shfl(max_val, threadIdx.x / tbSize); + } + + float sum = 0; + for (int i = 0; i < iterations; i++) { + data[i].x = __expf(data[i].x - max_val); + data[i].y = __expf(data[i].y - max_val); + data[i].z = __expf(data[i].z - max_val); + data[i].w = __expf(data[i].w - max_val); + + sum += (data[i].x + data[i].y + data[i].z + data[i].w); + } + + for (int i = 1; i < tbSize; i *= 2) { sum += g.shfl_xor(sum, i); } + + if (seq_length > tbSize) { + if (lane == 0) partialSum[wid] = sum; + b.sync(); + + if (lane < warp_num) sum = partialSum[lane]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + int iters = warp_num; + if (seq_length < iteration_stride) + iters = warp_num / (iteration_stride / max_threads_in_sequence); + + for (int i = 1; i < iters; i *= 2) { sum += g.shfl_xor(sum, i); } + + sum = g.shfl(sum, threadIdx.x / tbSize); + } + + sum += 1e-6; + + for (int i = 0; i < iterations; i++) { + data[i].x /= sum; + data[i].y /= sum; + data[i].z /= sum; + data[i].w /= sum; + + int data_id = i * iteration_stride + seq_lane; + if (data_id < seq_length) val_cast[data_offset + data_id] = data[i]; + } +} + +template +__global__ void attn_softmax(__half* vals, + const __half* attn_mask, + int heads, + int seq_length, + int iterations) +{ +#ifdef HALF_PRECISION_AVAILABLE + __shared__ float partialSum[MAX_WARP_NUM]; + + int warp_num = blockDim.x >> WARP_SIZE_BITS; + + int iteration_stride = blockDim.x; + int block_width = blockStride * seq_length; + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int batch = blockIdx.y; + int row = blockIdx.x; + int max_threads_in_sequence = std::max(seq_length, tbSeq); + int seq_lane = threadIdx.x % max_threads_in_sequence; + + int data_offset = batch * (gridDim.x * block_width) + row * block_width + + (threadIdx.x / max_threads_in_sequence) * seq_length; + int mask_offset = batch * seq_length; + + int wid = threadIdx.x >> WARP_SIZE_BITS; + int lane = threadIdx.x & 0x1f; + + float2* val_cast = reinterpret_cast(vals); + const float2* attn_mask_cast = reinterpret_cast(attn_mask); + + val_cast += data_offset; + attn_mask_cast += mask_offset; + + float2 low_data[MAX_THREAD_ITERATIONS]; + float2 high_data[MAX_THREAD_ITERATIONS]; + + float max_val = minus_infinity; + + for (int i = 0; i < iterations; i++) { + int data_id = i * iteration_stride + seq_lane; + if (data_id < seq_length) { + float2 data = val_cast[data_id]; + float2 mask = attn_mask_cast[data_id]; + + __half2* data_arr = reinterpret_cast<__half2*>(&data); + __half2* mask_arr = reinterpret_cast<__half2*>(&mask); + + low_data[i] = __half22float2(data_arr[0]); + high_data[i] = __half22float2(data_arr[1]); + float2 low_mask = __half22float2(mask_arr[0]); + float2 high_mask = __half22float2(mask_arr[1]); + + low_data[i].x += low_mask.x; + low_data[i].y += low_mask.y; + high_data[i].x += high_mask.x; + high_data[i].y += high_mask.y; + + max_val = (low_data[i].x > max_val ? low_data[i].x : max_val); + max_val = (low_data[i].y > max_val ? low_data[i].y : max_val); + max_val = (high_data[i].x > max_val ? high_data[i].x : max_val); + max_val = (high_data[i].y > max_val ? high_data[i].y : max_val); + } + } + + for (int i = 1; i < tbSize; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + if (seq_length > tbSize) { + if (lane == 0) partialSum[wid] = max_val; + b.sync(); + + if (lane < warp_num) max_val = partialSum[lane]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + int iters = warp_num; + if (seq_length < iteration_stride) + iters = warp_num / (iteration_stride / max_threads_in_sequence); + + for (int i = 1; i < iters; i *= 2) { + auto temp = g.shfl_xor(max_val, i); + max_val = (temp > max_val ? temp : max_val); + } + + max_val = g.shfl(max_val, threadIdx.x / tbSize); + } + + float sum = 0; + for (int i = 0; i < iterations; i++) { + int data_id = i * iteration_stride + seq_lane; + if (data_id < seq_length) { + low_data[i].x = __expf(low_data[i].x - max_val); + low_data[i].y = __expf(low_data[i].y - max_val); + high_data[i].x = __expf(high_data[i].x - max_val); + high_data[i].y = __expf(high_data[i].y - max_val); + + sum += (low_data[i].x + low_data[i].y + high_data[i].x + high_data[i].y); + } + } + + for (int i = 1; i < tbSize; i *= 2) { sum += g.shfl_xor(sum, i); } + + if (seq_length > tbSize) { + if (lane == 0) partialSum[wid] = sum; + b.sync(); + + if (lane < warp_num) sum = partialSum[lane]; + +#ifndef __STOCHASTIC_MODE__ + b.sync(); +#endif + + int iters = warp_num; + if (seq_length < iteration_stride) + iters = warp_num / (iteration_stride / max_threads_in_sequence); + + for (int i = 1; i < iters; i *= 2) { sum += g.shfl_xor(sum, i); } + + sum = g.shfl(sum, threadIdx.x / tbSize); + } + + sum += 1e-6; + + for (int i = 0; i < iterations; i++) { + int data_id = i * iteration_stride + seq_lane; + if (data_id < seq_length) { + float2 result_f; + __half2* result_h = reinterpret_cast<__half2*>(&result_f); + + low_data[i].x /= sum; + low_data[i].y /= sum; + high_data[i].x /= sum; + high_data[i].y /= sum; + + result_h[0] = __float22half2_rn(low_data[i]); + result_h[1] = __float22half2_rn(high_data[i]); + + val_cast[data_id] = result_f; + } + } + +#endif +} + +template +void launch_attn_softmax(T*, const T*, int, int, int, cudaStream_t); + +template <> +void launch_attn_softmax(float* vals, + const float* attn_mask, + int batch_size, + int heads, + int sequence_length, + cudaStream_t stream) +{ + const int threads = 128; + int seq_length4 = sequence_length / 4; + + dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); + + int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; + + dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / + subblock_max_workload * threads) + : threads); + int iterations = + (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads + : MAX_THREAD_ITERATIONS); + + if (sequence_length <= 8) + attn_softmax<2, (threads / 2), 2> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 16) + attn_softmax<4, (threads / 4), 4> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 32) + attn_softmax<8, (threads / 8), 8> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 64) + attn_softmax<16, (threads / 16), 16> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 128) + attn_softmax<32, (threads / 32), 32> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 256) + attn_softmax<32, (threads / 64), 64> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else { + const int threads = 256; + dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); + + int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; + + dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / + subblock_max_workload * threads) + : threads); + iterations = + (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads + : MAX_THREAD_ITERATIONS); + if (sequence_length <= 512) + attn_softmax<32, (threads / 128), 128><<>>( + vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length < (MAX_THREADS * MAX_THREAD_ITERATIONS * 4)) + attn_softmax<32, 1, 128><<>>( + vals, attn_mask, heads, seq_length4, iterations); + else + throw std::runtime_error( + "Unsupport Seq_Length! Check the restriction of the max_threads and " + "max_thread_iterations!"); + } +} + +template <> +void launch_attn_softmax<__half>(__half* vals, + const __half* attn_mask, + int batch_size, + int heads, + int sequence_length, + cudaStream_t stream) +{ + const int threads = 128; + int seq_length4 = sequence_length / 4; + + dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); + + int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; + + dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / + subblock_max_workload * threads) + : threads); + + int iterations = + (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads + : MAX_THREAD_ITERATIONS); + + if (sequence_length <= 8) + attn_softmax<2, (threads / 2), 2> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 16) + attn_softmax<4, (threads / 4), 4> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 32) + attn_softmax<8, (threads / 8), 8> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 64) + attn_softmax<16, (threads / 16), 16> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 128) + attn_softmax<32, (threads / 32), 32> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length <= 256) + attn_softmax<32, (threads / 64), 64> + <<>>(vals, attn_mask, heads, seq_length4, iterations); + else { + const int threads = 256; + dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); + + int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; + + dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / + subblock_max_workload * threads) + : threads); + iterations = + (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads + : MAX_THREAD_ITERATIONS); + if (sequence_length <= 512) + attn_softmax<32, (threads / 128), 128><<>>( + vals, attn_mask, heads, seq_length4, iterations); + else if (sequence_length < (MAX_THREADS * MAX_THREAD_ITERATIONS * 4)) + attn_softmax<32, 1, 128><<>>( + vals, attn_mask, heads, seq_length4, iterations); + else + throw std::runtime_error( + "Unsupport Seq_Length! Check the restriction of the max_threads and " + "max_thread_iterations!"); + } +} + +template +__global__ void softmax_backward_kernel(T* out_grad, const T* soft_inp, int seq_length) +{ + __shared__ float partialSum[MAX_WARP_NUM]; + + int warp_num = blockDim.x >> WARP_SIZE_BITS; // warp-count = num_threads / WARP_SIZE (32) + + int iteration_stride = blockDim.x; + int block_width = blockStride * seq_length; + + int iterations = (seq_length < (MAX_THREAD_ITERATIONS * iteration_stride) + ? (seq_length + iteration_stride - 1) / iteration_stride + : MAX_THREAD_ITERATIONS); + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + int row = blockIdx.x; + int id = threadIdx.x; + + int wid = id >> WARP_SIZE_BITS; + int lane = id & 0x1f; + + T val_reg[MAX_THREAD_ITERATIONS]; + T soft_reg[MAX_THREAD_ITERATIONS]; + float grad_reg = 0.0f; + +#pragma unroll + for (int i = 0; i < iterations; i++) { + int data_id = i * iteration_stride + id; + if (data_id < block_width) { + val_reg[i] = out_grad[row * block_width + data_id]; + soft_reg[i] = soft_inp[row * block_width + data_id]; + + grad_reg += ((float)val_reg[i] * + (float)soft_reg[i]); // if done in half, the multiplication, we may lose + // 2% of accuracy in computation!! + } + } + for (int i = 1; i < tbSize; i *= 2) grad_reg += g.shfl_xor(grad_reg, i); + + if (seq_length > tbSize) { + if (lane == 0) partialSum[wid] = grad_reg; + b.sync(); + + if (lane < warp_num) grad_reg = partialSum[lane]; + + int iters = warp_num; + if (seq_length < iteration_stride) iters = warp_num / (iteration_stride / seq_length); + + for (int i = 1; i < iters; i *= 2) grad_reg += g.shfl_xor(grad_reg, i); + + grad_reg = g.shfl(grad_reg, id / tbSize); + } + + for (int i = 0; i < iterations; i++) { + int data_id = i * iteration_stride + id; + if (data_id < block_width) { + float temp = (float)soft_reg[i] * ((float)val_reg[i] - grad_reg); + out_grad[row * block_width + data_id] = (T)temp; + } + } +} + +template +__global__ void softmax_backward_kernel_v2(T* grad /* input & output*/, + const T* output, + int softmax_length) +{ + int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; + int offset = batch_idx * softmax_length + threadIdx.x; + + grad += offset; + output += offset; + + T grad_reg[ITERATIONS]; + T output_reg[ITERATIONS]; + float sum = 0.0; + +#pragma unroll + for (int i = 0; i < ITERATIONS; ++i) { + int curr_idx = threadIdx.x + i * WARP_SIZE; + if (curr_idx < softmax_length) { + grad_reg[i] = grad[i * WARP_SIZE]; + output_reg[i] = output[i * WARP_SIZE]; + sum += (float)grad_reg[i] * (float)output_reg[i]; + } + } + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + + for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_xor(sum, i); + +#pragma unroll + for (int i = 0; i < ITERATIONS; ++i) { + int curr_idx = threadIdx.x + i * WARP_SIZE; + if (curr_idx < softmax_length) + grad[i * WARP_SIZE] = (float)output_reg[i] * ((float)grad_reg[i] - sum); + } +} + +__global__ void softmax_backward_kernel_arbitrary_length(__half* grad /* input & output*/, + const __half* output, + int softmax_length) +{ + int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; + int offset = batch_idx * softmax_length + threadIdx.x; + + const float4* output_cast = reinterpret_cast(output); + float4* grad_cast = reinterpret_cast(grad); + + grad_cast += offset; + output_cast += offset; + + float sum = 0.0; + int curr_idx = threadIdx.x; + while (curr_idx < softmax_length) { + float4 out_reg = output_cast[curr_idx]; + float4 grad_reg = grad_cast[curr_idx]; + __half2* out_h = reinterpret_cast<__half2*>(&out_reg); + __half2* grad_h = reinterpret_cast<__half2*>(&grad_reg); +#pragma unroll + for (int m = 0; m < 4; m++) grad_h[m] *= out_h[m]; + sum += ((float)grad_h[0].x + (float)grad_h[0].y + (float)grad_h[1].x + (float)grad_h[1].y) + + ((float)grad_h[2].x + (float)grad_h[2].y + (float)grad_h[3].x + (float)grad_h[3].y); + curr_idx += WARP_SIZE; + } + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_xor(sum, i); + + curr_idx = threadIdx.x; + while (curr_idx < softmax_length) { + float4 out_reg = output_cast[curr_idx]; + float4 grad_reg = grad_cast[curr_idx]; + __half* grad_h = reinterpret_cast<__half*>(&grad_reg); + __half* out_h = reinterpret_cast<__half*>(&out_reg); + +#pragma unroll + for (int m = 0; m < 8; m++) grad_h[m] = (float)out_h[m] * ((float)grad_h[m] - sum); + + grad_cast[curr_idx] = grad_reg; + curr_idx += WARP_SIZE; + } +} + +__global__ void softmax_backward_kernel_arbitrary_length(float* grad /* input & output*/, + const float* output, + int softmax_length) +{ + int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; + int offset = batch_idx * softmax_length + threadIdx.x; + + const float4* output_cast = reinterpret_cast(output); + float4* grad_cast = reinterpret_cast(grad); + + grad_cast += offset; + output_cast += offset; + + float sum = 0.0; + int curr_idx = threadIdx.x; + while (curr_idx < softmax_length) { + float4 out_reg = output_cast[curr_idx]; + float4 grad_reg = grad_cast[curr_idx]; + + grad_reg.x *= out_reg.x; + grad_reg.y *= out_reg.y; + grad_reg.z *= out_reg.z; + grad_reg.w *= out_reg.w; + sum += (grad_reg.x + grad_reg.y + grad_reg.z + grad_reg.w); + + curr_idx += WARP_SIZE; + } + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile g = cg::tiled_partition(b); + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_xor(sum, i); + + curr_idx = threadIdx.x; + while (curr_idx < softmax_length) { + float4 out_reg = output_cast[curr_idx]; + float4 grad_reg = grad_cast[curr_idx]; + grad_reg.x = out_reg.x * (grad_reg.x - sum); + grad_reg.y = out_reg.y * (grad_reg.y - sum); + grad_reg.z = out_reg.z * (grad_reg.z - sum); + grad_reg.w = out_reg.w * (grad_reg.w - sum); + + grad_cast[curr_idx] = grad_reg; + curr_idx += WARP_SIZE; + } +} + +template +void launch_attn_softmax_backward_v2(T* out_grad, + const T* soft_inp, + int batch_size, + int heads, + int seq_length, + cudaStream_t stream) +{ + const int warps_per_block = 4; + dim3 grid_dim(batch_size * heads * seq_length / warps_per_block); + dim3 block_dim(WARP_SIZE, warps_per_block); + + if (seq_length <= 32) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 64) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 128) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 256) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 384) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 512) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 768) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 1024) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 2048) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 4096) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else if (seq_length <= 8192) + softmax_backward_kernel_v2 + <<>>(out_grad, soft_inp, seq_length); + else + softmax_backward_kernel_arbitrary_length<<>>( + out_grad, soft_inp, seq_length / (4 << ((sizeof(T) & 2) >> 1))); +} + +template void launch_attn_softmax_backward_v2<__half>(__half* out_grad, + const __half* soft_inp, + int batch_size, + int heads, + int seq_length, + cudaStream_t stream); +template void launch_attn_softmax_backward_v2(float* out_grad, + const float* soft_inp, + int batch_size, + int heads, + int seq_length, + cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/transform_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/transform_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..dbecc31536c6c9355c4a3e243f2a0219bc476e6e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/transform_kernels.cu @@ -0,0 +1,580 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "custom_cuda_layers.h" + +#define rows_trans 16 +#define cols_trans 16 + +template +__global__ void Transpose_Kernel(const T* inp, T* out, int row_width, int col_width) +{ + __shared__ T data_block[rows_trans * (cols_trans + 1)]; + + int r = threadIdx.x / cols_trans; + int c = threadIdx.x % cols_trans; + + int m = row_width / cols_trans; + + int i = blockIdx.x / m * rows_trans + r; + int j = blockIdx.x % m * cols_trans + c; + + int row_stride = rows_trans / ((rows_trans * cols_trans + THREADS - 1) / THREADS); + + for (int k = 0; k < rows_trans; k += row_stride) + data_block[(k + r) * cols_trans + c] = inp[(i + k) * row_width + j]; + + __syncthreads(); + + i = blockIdx.x % m * rows_trans + r; + j = blockIdx.x / m * cols_trans + c; + + for (int k = 0; k < rows_trans; k += row_stride) + out[(i + k) * col_width + j] = data_block[c * cols_trans + r + k]; +} + +template <> +void Transpose<__half>(const __half* inp_mat, + __half* out_mat, + int rows, + int cols, + cudaStream_t stream) +{ + int threads = THREADS; + + Transpose_Kernel<__half><<<(rows * cols + threads - 1) / threads, threads, 0, stream>>>( + inp_mat, out_mat, cols, rows); +} + +template <> +void Transpose(const float* inp_mat, float* out_mat, int rows, int cols, cudaStream_t stream) +{ + int threads = THREADS; + + Transpose_Kernel<<<(rows * cols + threads - 1) / threads, threads, 0, stream>>>( + inp_mat, out_mat, cols, rows); +} + +template +__global__ void transform_0213(T* output, + const T* vals, + int hidden_dim, + int seq_length, + int heads, + int head_ext); + +template <> +__global__ void transform_0213(float* output, + const float* vals, + int hidden_dim, + int seq_length, + int heads, + int head_ext) +{ + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y / head_ext; // Sequence ID (0-127) + int d2 = threadIdx.y + (blockIdx.y % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + const float4* vals_vec = reinterpret_cast(vals); + float4* output_vec = reinterpret_cast(output); + + float4 inputs = vals_vec[d0 * d0_stride + d1 * d1_stride + d2 * d2_stride + d3]; + output_vec[d0 * d0_out_stride + d1 * d1_out_stride + d2 * d2_out_stride + d3] = inputs; +} + +template <> +__global__ void transform_0213<__half>(__half* output, + const __half* vals, + int hidden_dim, + int seq_length, + int heads, + int head_ext) +{ +#ifdef HALF_PRECISION_AVAILABLE + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y / head_ext; // Sequence ID (0-127) + int d2 = threadIdx.y + (blockIdx.y % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + float4 vals_arr[1]; + + const float4* vals_vec = reinterpret_cast(vals); + float4* output_vec = reinterpret_cast(output); + + vals_arr[0] = vals_vec[d0 * d0_stride + d1 * d1_stride + d2 * d2_stride + d3]; + output_vec[d0 * d0_out_stride + d1 * d1_out_stride + d2 * d2_out_stride + d3] = vals_arr[0]; +#endif +} + +template <> +void launch_transform_0213(float* output, + const float* vals, + int batch_size, + int seq_length, + int hidden_dim, + int heads, + cudaStream_t stream) +{ + hidden_dim >>= 2; + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + dim3 block_dim(hidden_dim / heads, (heads / head_ext)); + dim3 grid_dim(batch_size, (seq_length * head_ext)); + + transform_0213 + <<>>(output, vals, hidden_dim, seq_length, heads, head_ext); +} + +template <> +void launch_transform_0213<__half>(__half* output, + const __half* vals, + int batch_size, + int seq_length, + int hidden_dim, + int heads, + cudaStream_t stream) +{ + hidden_dim >>= 3; + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + dim3 block_dim(hidden_dim / heads, (heads / head_ext)); + dim3 grid_dim(batch_size, (seq_length * head_ext)); + transform_0213<__half> + <<>>(output, vals, hidden_dim, seq_length, heads, head_ext); +} + +// Bias add +template +__global__ void bias_add_transform_0213(T* output, + const T* vals, + const T* bias, + int hidden_dim, + int seq_length, + int heads, + int head_ext); + +template <> +__global__ void bias_add_transform_0213(float* output, + const float* vals, + const float* bias, + int hidden_dim, + int seq_length, + int heads, + int head_ext) +{ + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = blockIdx.z / head_ext; // Hidden count + int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + const float4* vals_vec = reinterpret_cast(vals); + const float4* bias_vec = reinterpret_cast(bias); + float4* output_vec = reinterpret_cast(output); + + float4 inputs = vals_vec[d0 * d0_stride * (gridDim.z / head_ext) + cnt * d1_stride + + d1 * d1_stride * (gridDim.z / head_ext) + d2 * d2_stride + d3]; + float4 biases = bias_vec[cnt * d1_stride + d2 * d2_stride + d3]; + + float4 outputs; + outputs.x = inputs.x + biases.x; + outputs.y = inputs.y + biases.y; + outputs.z = inputs.z + biases.z; + outputs.w = inputs.w + biases.w; + + output_vec[cnt * d0_out_stride * gridDim.x + d0 * d0_out_stride + d1 * d1_out_stride + + d2 * d2_out_stride + d3] = outputs; +} + +#define ATTN_H 3 +#define MAX_SEQ_LINE 10 + +template <> +__global__ void bias_add_transform_0213<__half>(__half* output, + const __half* vals, + const __half* bias, + int hidden_dim, + int seq_length, + int heads, + int head_ext) +{ +#ifdef HALF_PRECISION_AVAILABLE + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = blockIdx.z / head_ext; // Hidden count + int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + float4 vals_arr; + float4 bias_arr; + float4 output_arr; + __half2* vals_half = reinterpret_cast<__half2*>(&vals_arr); + __half2* bias_half = reinterpret_cast<__half2*>(&bias_arr); + __half2* output_half = reinterpret_cast<__half2*>(&output_arr); + + const float4* vals_vec = reinterpret_cast(vals); + const float4* bias_vec = reinterpret_cast(bias); + float4* output_vec = reinterpret_cast(output); + + vals_vec += (d0 * d0_stride * (gridDim.z / head_ext)); + vals_vec += (d1 * d1_stride * (gridDim.z / head_ext)); + vals_vec += (cnt * d1_stride); + vals_vec += (d2 * d2_stride); + + bias_vec += (cnt * d1_stride); + bias_vec += (d2 * d2_stride); + + output_vec += (cnt * d0_stride * gridDim.x); + output_vec += (d1 * d2_stride); + output_vec += (d0 * d0_stride); + output_vec += (d2 * d2_out_stride); + + bias_arr = bias_vec[d3]; + vals_arr = vals_vec[d3]; + +#if defined(__ACC_HALF__) + output_half[0] = vals_half[0] + bias_half[0]; + output_half[1] = vals_half[1] + bias_half[1]; + output_half[2] = vals_half[2] + bias_half[2]; + output_half[3] = vals_half[3] + bias_half[3]; +#else + float2 bias_arr_f[4]; + float2 vals_arr_f[4]; +#pragma unroll + for (int l = 0; l < 4; l++) { + bias_arr_f[l] = __half22float2(bias_half[l]); + vals_arr_f[l] = __half22float2(vals_half[l]); + vals_arr_f[l].x += bias_arr_f[l].x; + vals_arr_f[l].y += bias_arr_f[l].y; + output_half[l] = __float22half2_rn(vals_arr_f[l]); + } +#endif + output_vec[d3] = output_arr; + +#endif +} + +__global__ void bias_add_transform_0213_v2(__half* output, + const __half* vals, + const __half* bias, + int hidden_dim, + int seq_length, + int heads) +{ +#ifdef HALF_PRECISION_AVAILABLE + __shared__ float4 in_data[3072]; + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + int iteration_stride = d1_stride * blockDim.z; // Hidden * 3 / 8 + int batch_stride = d0_stride * blockDim.z; // Hidden * S * 3 / 8 + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = d2_stride * seq_length; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y; // Sequence ID (0-127) + int cnt = threadIdx.z; // blockIdx.z; // Hidden count + int d2 = threadIdx.y; // Head (0-11) + int d3 = threadIdx.x; // Values (groups of 4) + + float4 vals_arr[1]; + float4 bias_arr[1]; + float4 output_arr[1]; + __half2* vals_half = reinterpret_cast<__half2*>(vals_arr); + __half2* bias_half = reinterpret_cast<__half2*>(bias_arr); + __half2* output_half = reinterpret_cast<__half2*>(output_arr); + + const float4* vals_vec = reinterpret_cast(vals); + const float4* bias_vec = reinterpret_cast(bias); + float4* output_vec = reinterpret_cast(output); + + int iter_index = cnt * d1_stride + d2 * d2_stride + d3; + int input_offset = d0 * batch_stride + d1 * (iteration_stride << 1); + bias_arr[0] = bias_vec[iter_index]; + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_id = iter * iteration_stride + iter_index; + vals_arr[0] = vals_vec[input_offset + iter_id]; + + output_half[0] = vals_half[0] + bias_half[0]; + output_half[1] = vals_half[1] + bias_half[1]; + output_half[2] = vals_half[2] + bias_half[2]; + output_half[3] = vals_half[3] + bias_half[3]; + + in_data[iter_id] = output_arr[0]; + } + __syncthreads(); + + iteration_stride = blockDim.z * (blockDim.y >> 1); + int matrix_stride = (d0_out_stride * gridDim.x); + int head_count = (d2 >> 1) + cnt * (blockDim.y >> 1); + + int out_index = d0 * d0_out_stride + d1 * (d1_out_stride << 1) + d3 + (d2 % 2) * d2_stride; + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_row = (iter * iteration_stride) + head_count; + int iter_offset = + (iter_row % blockDim.y) * d2_out_stride + (iter_row / blockDim.y) * matrix_stride; + output_vec[out_index + iter_offset] = + in_data[iter_row * d2_stride + d3 + (d2 % 2) * (d1_stride * blockDim.z)]; + } +#endif +} + +// [B S C*H] - > C * [B A S N] +template <> +void launch_bias_add_transform_0213(float* output, + const float* vals, + const float* bias, + int batch_size, + int seq_length, + int hidden_dim, + int heads, + cudaStream_t stream, + int trans_count) +{ + hidden_dim >>= 2; + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + + dim3 block_dim(hidden_dim / heads, (heads / head_ext)); + dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); + + bias_add_transform_0213<<>>( + output, vals, bias, hidden_dim, seq_length, heads, head_ext); +} + +template <> +void launch_bias_add_transform_0213<__half>(__half* output, + const __half* vals, + const __half* bias, + int batch_size, + int seq_length, + int hidden_dim, + int heads, + cudaStream_t stream, + int trans_count) +{ + hidden_dim >>= 3; + if (hidden_dim > 128 || hidden_dim < 16) { + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + dim3 block_dim(hidden_dim / heads, (heads / head_ext)); + dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); + bias_add_transform_0213<__half><<>>( + output, vals, bias, hidden_dim, seq_length, heads, head_ext); + } else { + dim3 block_dim(hidden_dim / heads, heads, trans_count); + dim3 grid_dim(batch_size, seq_length / 2); + bias_add_transform_0213_v2<<>>( + output, vals, bias, hidden_dim, seq_length, heads); + } +} + +template +__global__ void transform4d_0213(T* out, + const T* in, + int heads, + int seq_length, + int hidden_dim, + int head_ext); + +template <> +__global__ void transform4d_0213(float* out, + const float* in, + int heads, + int seq_length, + int hidden_dim, + int head_ext) +{ + int d0_stride = hidden_dim * seq_length; + int d1_stride = d0_stride / heads; + int d2_stride = hidden_dim / heads; + + int d0_out_stride = d0_stride; + int d1_out_stride = d2_stride; + int d2_out_stride = hidden_dim; + + int d0 = blockIdx.x; // Batch + int d1 = blockIdx.y / ((seq_length - 1) / blockDim.y + 1); // Head + int d2 = (threadIdx.y + blockDim.y * blockIdx.y) % seq_length; + int cnt = blockIdx.z; + int d3 = threadIdx.x; // Values (groups of 8) + + if (d2 < seq_length) { + const float4* in_vec = reinterpret_cast(in); + float4* out_vec = reinterpret_cast(out); + + float4 vals_vec = in_vec[cnt * d0_stride * gridDim.x + d0 * d0_stride + d1 * d1_stride + + d2 * d2_stride + d3]; + out_vec[d0 * d0_out_stride * gridDim.z + cnt * d2_out_stride + d1 * d1_out_stride + + d2 * d2_out_stride * gridDim.z + d3] = vals_vec; + } +} + +template <> +__global__ void transform4d_0213<__half>(__half* out, + const __half* in, + int heads, + int seq_length, + int hidden_dim, + int head_ext) +{ +#ifdef HALF_PRECISION_AVAILABLE + + int d0_stride = hidden_dim * (seq_length / head_ext); + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head + int d2 = blockIdx.z / head_ext; // Sequence + int cnt = blockIdx.y; // Hidden count + int d3 = threadIdx.x; // Values (groups of 8) + + const float4* in_vec = reinterpret_cast(in); + float4* out_vec = reinterpret_cast(out); + + in_vec += (cnt * d0_stride * gridDim.x); + in_vec += (d0 * d0_stride); + in_vec += (d2 * d2_stride); + in_vec += (d1 * d2_stride * seq_length); + + out_vec += (cnt * d1_stride); + out_vec += (d1 * d2_stride); + out_vec += (d0 * d0_stride * gridDim.y); + out_vec += (d2 * d1_stride * gridDim.y); + + out_vec[d3] = in_vec[d3]; + +#endif +} + +__global__ void transform4d_0213_v2(__half* out, + const __half* in, + int heads, + int seq_length, + int hidden_dim) +{ +#ifdef HALF_PRECISION_AVAILABLE + __shared__ float4 in_data[3072]; + + int d0_stride = hidden_dim * seq_length; + int d1_stride = hidden_dim; + int d2_stride = hidden_dim / heads; + + int d0 = blockIdx.x; // Batch + int d1 = threadIdx.y; // Head + int d2 = blockIdx.y; // Sequence + int cnt = threadIdx.z; // Hidden count + int d3 = threadIdx.x; // Values (groups of 8) + + const float4* in_vec = reinterpret_cast(in); + float4* out_vec = reinterpret_cast(out); + + int input_offset = d0 * d0_stride + d2 * (d2_stride << 1) + d3 + (d1 % 2) * d2_stride; + int head_count = (d1 >> 1) + cnt * (blockDim.y >> 1); + int iteration_stride = blockDim.z * (blockDim.y >> 1); + int matrix_stride = (d0_stride * gridDim.x); + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_row = iter * iteration_stride + head_count; + int iter_offset = (iter_row % blockDim.y) * d2_stride; + + in_data[d3 + iter_offset + (iter_row / blockDim.y + (d1 % 2) * blockDim.z) * d1_stride] = + in_vec[input_offset + iter_offset * seq_length + + (iter_row / blockDim.y) * matrix_stride]; + } + __syncthreads(); + + iteration_stride = d1_stride * blockDim.z; + int iter_index = cnt * d1_stride + d1 * d2_stride + d3; + int output_offset = d0 * d0_stride * blockDim.z + d2 * (iteration_stride << 1); + +#pragma unroll + for (int iter = 0; iter < 2; iter++) { + int iter_id = iter * iteration_stride + iter_index; + out_vec[output_offset + iter_id] = in_data[iter_id]; + } +#endif +} + +// 3 * [B A S N] - > [B S C*H] +template <> +void launch_transform4d_0213(float* out, + const float* in, + int batch_size, + int heads, + int seq_length, + int hidden_dim, + cudaStream_t stream, + int trans_count) +{ + hidden_dim >>= 2; + dim3 grid_dims(batch_size, heads * ((seq_length - 1) / 8 + 1), trans_count); + dim3 block_dims(hidden_dim / heads, 8); + transform4d_0213 + <<>>(out, in, heads, seq_length, hidden_dim, 1); +} + +template <> +void launch_transform4d_0213<__half>(__half* out, + const __half* in, + int batch_size, + int heads, + int seq_length, + int hidden_dim, + cudaStream_t stream, + int trans_count) +{ + hidden_dim >>= 3; + if (hidden_dim > 128 || hidden_dim < 16) { + int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; + dim3 grid_dims(batch_size, trans_count, (seq_length * head_ext)); + dim3 block_dims(hidden_dim / heads, (heads / head_ext)); + transform4d_0213<__half><<>>( + out, in, heads, seq_length, hidden_dim, head_ext); + } else { + dim3 grid_dims(batch_size, seq_length / 2); + dim3 block_dims(hidden_dim / heads, heads, trans_count); + transform4d_0213_v2<<>>( + out, in, heads, seq_length, hidden_dim); + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/py_ds_utils.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/py_ds_utils.cpp new file mode 100644 index 0000000000000000000000000000000000000000..df5c9f361c616f07fd13640804200dca25d71191 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/py_ds_utils.cpp @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* + Collection of system utilities. +*/ + +#include +#include "tensor_cast.h" +using namespace pybind11::literals; + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("cast_to_byte_tensor", + py::overload_cast(&cast_to_byte_tensor), + "Cast a 1-dimensional tensor of any type to byte tensor.", + "src_tensor"_a); + + m.def("cast_to_byte_tensor", + py::overload_cast&>(&cast_to_byte_tensor), + "Cast a multi-dimensional tensor of any type to byte tensor.", + "src_tensor"_a); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.cpp new file mode 100644 index 0000000000000000000000000000000000000000..8352bc72bb0286e30268e77f243787c2923a0d40 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.cpp @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "tensor_cast.h" + +at::Tensor cast_to_byte_tensor(at::Tensor& src_tensor) +{ + if (src_tensor.nbytes() <= 1) return src_tensor; + + auto options = torch::TensorOptions() + .dtype(torch::kUInt8) + .layout(src_tensor.layout()) + .device(src_tensor.device()); + return at::from_blob( + src_tensor.data_ptr(), static_cast(src_tensor.nbytes()), options); +} + +std::vector cast_to_byte_tensor(std::vector& tensor_list) +{ + std::vector byte_tensors; + for (auto src_tensor : tensor_list) { byte_tensors.push_back(cast_to_byte_tensor(src_tensor)); } + + return byte_tensors; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.h new file mode 100644 index 0000000000000000000000000000000000000000..86155567122d592c58979606a21aa9c3f00d5511 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* + Utilities for type casting torch tensors without data movement. +*/ + +#include +#include + +using namespace std; +at::Tensor cast_to_byte_tensor(at::Tensor& src_tensor); + +std::vector cast_to_byte_tensor(std::vector& tensor_list); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp new file mode 100644 index 0000000000000000000000000000000000000000..dc727f8fa2168ef7897f685de6ac5d80aaad8309 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_adagrad.h" +#include +#include +#include +#include +#include +#include + +static std::unordered_map> s_optimizers; + +// C++ interface + +void Adagrad_Optimizer::Step_1(float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<1>( + &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); +#endif + if (_param_size > rounded_size) { + float step_size = -1 * _alpha; + ds_half_precision_t* grads_cast_h; + ds_half_precision_t* params_cast_h; + if (half_precision) { + grads_cast_h = reinterpret_cast(grads); + params_cast_h = reinterpret_cast(_params); + } + for (size_t t = rounded_size; t < _param_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > _param_size) copy_size = _param_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t k = t; k < offset; k++) { + float grad = half_precision ? (float)grads_cast_h[k] : grads[k]; + float param = half_precision ? (float)params_cast_h[k] : _params[k]; + float momentum = grads[k]; + float variance = _exp_avg_sq[k]; + if (_weight_decay > 0) { grad = param * _weight_decay + grad; } + + variance += grad * grad; + + grad = sqrt(variance); + grad += _eps; + grad = momentum / grad; + param = grad * step_size + param; + if (half_precision) + params_cast_h[k] = (ds_half_precision_t)param; + else + _params[k] = param; + // STORE UPDATE TERM TO GRAD'S MEMORY + grads[k] = grad * step_size; + _exp_avg_sq[k] = variance; + } + } + } +} + +void Adagrad_Optimizer::Step_4(float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<4>( + &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); +#endif + if (_param_size > rounded_size) + Step_1((_params + rounded_size), + (grads + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size), + (dev_params != nullptr ? (dev_params + rounded_size) : dev_params), + half_precision); +} + +int create_adagrad_optimizer(int optimizer_id, + float alpha = 1e-2, + float eps = 1e-8, + float weight_decay = 0, + bool should_log = false) +{ + auto opt = std::make_shared(alpha, eps, weight_decay); + + s_optimizers[optimizer_id] = opt; + + if (should_log) { + std::string avx_type = ""; +#if defined(__AVX512__) + avx_type = "AVX512"; +#else +#if defined(__AVX256__) + avx_type = "AVX2"; +#else + avx_type = "scalar"; +#endif +#endif + + printf("Adagrad Optimizer #%d is created with %s arithmetic capability.\n", + optimizer_id, + avx_type.c_str()); + printf("Config: alpha=%f, weight_decay=%f\n", alpha, weight_decay); + } + + return 0; +} + +void Adagrad_Optimizer::Step_8(float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<8>( + &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); +#endif + if (_param_size > rounded_size) + Step_4((_params + rounded_size), + (grads + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size), + (dev_params != nullptr ? (dev_params + rounded_size) : dev_params), + half_precision); +} + +int ds_adagrad_step(int optimizer_id, + size_t step, + float lr, + float epsilon, + float weight_decay, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg_sq) +{ + auto params_c = params.contiguous(); + auto grads_c = grads.contiguous(); + auto exp_avg_sq_c = exp_avg_sq.contiguous(); + + float* params_ptr = (float*)params_c.data_ptr(); + float* grads_ptr = (float*)grads_c.data_ptr(); + float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr(); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + opt->IncrementStep(step); + opt->update_state(lr, epsilon, weight_decay); + opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.numel()); + + return 0; +} + +int ds_adagrad_step_plus_copy(int optimizer_id, + size_t step, + float lr, + float epsilon, + float weight_decay, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg_sq, + torch::Tensor& gpu_params) +{ + assert(false); + return 0; +} + +int destroy_adagrad_optimizer(int optimizer_id) +{ + s_optimizers.erase(optimizer_id); + + return 0; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("adagrad_update", &ds_adagrad_step, "DeepSpeed CPU Adagrad update (C++)"); + m.def("adagrad_update_copy", + &ds_adagrad_step_plus_copy, + "DeepSpeed CPU Adagrad update and param copy (C++)"); + m.def("create_adagrad", &create_adagrad_optimizer, "DeepSpeed CPU Adagrad (C++)"); + m.def("destroy_adagrad", &destroy_adagrad_optimizer, "DeepSpeed CPU Adagrad destroy (C++)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp new file mode 100644 index 0000000000000000000000000000000000000000..13b390248608b046dab443f85346b5446a47d722 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_adam", + &multi_tensor_adam_cuda, + "Compute and apply gradient update to parameters for Adam optimizer"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp new file mode 100644 index 0000000000000000000000000000000000000000..0720a020247a3ebf80fc231e11333ea56cb09924 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp @@ -0,0 +1,159 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include + +#include + +#include +#include "multi_tensor_apply.dp.hpp" +#include "type_shim.h" + +#define BLOCK_SIZE 512 +#define ILP 4 + +typedef enum : int { + ADAM_MODE_0 = 0, // L2 regularization mode + ADAM_MODE_1 = 1 // Decoupled weight decay mode(AdamW) +} adamMode_t; + +using MATH_T = float; + +template +struct AdamFunctor { + __inline__ __attribute__((always_inline)) void operator()(int chunk_size, + volatile int* noop_gmem, + TensorListMetadata<4>& tl, + const float beta1, + const float beta2, + const float beta1_correction, + const float beta2_correction, + const float epsilon, + const float lr, + adamMode_t mode, + const float decay) + { + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int tensor_loc = tl.block_to_tensor[item_ct1.get_group(2)]; + + int chunk_idx = tl.block_to_chunk[item_ct1.get_group(2)]; + int n = tl.sizes[tensor_loc]; + + T* g = (T*)tl.addresses[0][tensor_loc]; + g += chunk_idx * chunk_size; + + T* p = (T*)tl.addresses[1][tensor_loc]; + p += chunk_idx * chunk_size; + + T* m = (T*)tl.addresses[2][tensor_loc]; + m += chunk_idx * chunk_size; + + T* v = (T*)tl.addresses[3][tensor_loc]; + v += chunk_idx * chunk_size; + + n -= chunk_idx * chunk_size; + + // see note in multi_tensor_scale_kernel.cu + for (int i_start = 0; i_start < n && i_start < chunk_size; + i_start += item_ct1.get_local_range(2) * ILP) { + MATH_T r_g[ILP]; + MATH_T r_p[ILP]; + MATH_T r_m[ILP]; + MATH_T r_v[ILP]; +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + item_ct1.get_local_id(2) + ii * item_ct1.get_local_range(2); + if (i < n && i < chunk_size) { + r_g[ii] = g[i]; + r_p[ii] = p[i]; + r_m[ii] = m[i]; + r_v[ii] = v[i]; + } else { + r_g[ii] = MATH_T(0); + r_p[ii] = MATH_T(0); + r_m[ii] = MATH_T(0); + r_v[ii] = MATH_T(0); + } + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + if (mode == ADAM_MODE_0) { // L2 + r_g[ii] = r_g[ii] + (decay * r_p[ii]); + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sycl::sqrt(next_v_unbiased) + epsilon; + MATH_T update = next_m_unbiased / denom; + r_p[ii] = r_p[ii] - (lr * update); + } else { // weight decay + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sycl::sqrt(next_v_unbiased) + epsilon; + MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); + r_p[ii] = r_p[ii] - (lr * update); + } + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + item_ct1.get_local_id(2) + ii * item_ct1.get_local_range(2); + if (i < n && i < chunk_size) { + p[i] = r_p[ii]; + m[i] = r_m[ii]; + v[i] = r_v[ii]; + } + } + } + } +}; + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay) +{ + using namespace at; + + // Handle bias correction mode + float bias_correction1 = 1.0f, bias_correction2 = 1.0f; + if (bias_correction == 1) { + bias_correction1 = 1 - std::pow(beta1, step); + bias_correction2 = 1 - std::pow(beta2, step); + } + + // Assume single type across p,g,m1,m2 now + DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), + 0, + "adam", + multi_tensor_apply<4>(BLOCK_SIZE, + chunk_size, + noop_flag, + tensor_lists, + AdamFunctor(), + beta1, + beta2, + bias_correction1, + bias_correction2, + epsilon, + lr, + (adamMode_t)mode, + weight_decay);) +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp new file mode 100644 index 0000000000000000000000000000000000000000..14a130e2a23e906ba9b305ee3968aa4e27c17cf8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp @@ -0,0 +1,226 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +#include +#include "compat.h" + +#include +#include +#include + +namespace at { +namespace cuda { +sycl::queue* getCurrentCUDAStream() +{ + c10::xpu::XPUStream stream = c10::xpu::getCurrentXPUStream(); + auto& queue = stream.queue(); + return &queue; +} + +sycl::queue* getStreamFromPool(bool) +{ + // not implemented + return nullptr; +} +} // namespace cuda +} // namespace at +// #include + +// This header is the one-stop shop for all your multi-tensor apply needs. + +// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) +constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; +constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; + +template +struct TensorListMetadata { + void* addresses[n][depth_to_max_tensors[n - 1]]; + int sizes[depth_to_max_tensors[n - 1]]; + unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. + int start_tensor_this_launch; +}; + +template +class multi_tensor_apply_kernel { +public: + multi_tensor_apply_kernel(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) + : chunk_size(chunk_size), noop_flag(noop_flag), tl(tl), callable(callable), args(args...) + { + } + + // This should be identical to original __global__ function + static void inline __global__function(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) + { + callable(chunk_size, noop_flag, tl, args...); + } + + // If global function template contains parameter pack, + // we only deal with parameter pack at the end of template parameter list + template + static void inline __tuple_expand_driver(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + Tuple args, + std::index_sequence) + { + __global__function(chunk_size, noop_flag, tl, callable, std::get(args)...); + } + + // + // Because __global__ function can't really use any reference types, we can sure that args + // are all good behaviors + // + void operator()(sycl::nd_item<3>) const + { + __tuple_expand_driver(chunk_size, + noop_flag, + tl, + callable, + args, + std::make_index_sequence()); + } + +private: + int chunk_size; + volatile int* noop_flag; + T tl; + U callable; + std::tuple args; +}; + +// to make sure multi_tensor_apply_kernel can be used in sycl::buffer +namespace sycl { +template +struct is_device_copyable> : std::true_type {}; +} // namespace sycl + +template +void multi_tensor_apply(int block_size, + int chunk_size, + const at::Tensor& noop_flag, + const std::vector>& tensor_lists, + T callable, + ArgTypes... args) +{ + TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); + int len0 = tensor_lists[0].size(); + TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); + auto ref_device = tensor_lists[0][0].device(); + TORCH_CHECK(ref_device.type() == at::kXPU, "expected input to be on cuda"); + for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices + { + TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); + for (int t = 0; t < tensor_lists[l].size(); t++) { + // TODO: Print which tensor fails. + bool contiguous_memory = tensor_lists[l][t].is_contiguous(); +#ifdef VERSION_GE_1_5 + contiguous_memory = (contiguous_memory || + tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); +#endif + TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); + TORCH_CHECK(tensor_lists[l][t].device() == ref_device, + "A tensor was not on the same device as the first tensor"); + TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); + } + } + + int ntensors = tensor_lists[0].size(); + + TensorListMetadata tl; + + /* const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); */ + auto stream = at::cuda::getCurrentCUDAStream(); + + tl.start_tensor_this_launch = 0; + int loc_block_info = 0; + int loc_tensor_info = 0; + for (int t = 0; t < ntensors; t++) { + tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); + for (int d = 0; d < depth; d++) + tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); + loc_tensor_info++; + + int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; + + for (int chunk = 0; chunk < chunks_this_tensor; chunk++) { + // std::cout << chunks_this_tensor << std::endl; + tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; + tl.block_to_chunk[loc_block_info] = chunk; + loc_block_info++; + + bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && + chunk == chunks_this_tensor - 1); + bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); + bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); + if (tensors_full || blocks_full || last_chunk) { + // using accscalar_t = acc_type; + /* multi_tensor_apply_kernel, T, ArgTypes...> + * fn(chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); */ + if constexpr (sizeof(multi_tensor_apply_kernel( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...)) < + 2048) { + ((sycl::queue*)(stream)) + ->parallel_for( + sycl::nd_range<3>(sycl::range<3>(1, 1, loc_block_info) * + sycl::range<3>(1, 1, block_size), + sycl::range<3>(1, 1, block_size)), + multi_tensor_apply_kernel( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...)); + } else { + auto capture = multi_tensor_apply_kernel( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); + sycl::buffer params(const_cast(&capture), + sycl::range<1>(1)); + stream->submit([&](sycl::handler& cgh) { + auto device_params = + params.template get_access(cgh); + cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, loc_block_info) * + sycl::range<3>(1, 1, block_size), + sycl::range<3>(1, 1, block_size)), + [=](sycl::nd_item<3> item) { device_params[0](item); }); + }); + } + 0; + + // Reset. The control flow possibilities here make my brain hurt. + loc_block_info = 0; + if (chunk == chunks_this_tensor - 1) { + // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + loc_tensor_info = 0; + tl.start_tensor_this_launch = t + 1; + } else { + // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; + for (int d = 0; d < depth; d++) + tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; + loc_tensor_info = 1; + tl.start_tensor_this_launch = t; + } + } + } + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp new file mode 100644 index 0000000000000000000000000000000000000000..cfd004ef13577c02f958fa18d116131398443562 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +inline void has_capability_or_fail(const sycl::device& dev, + const std::initializer_list& props) +{ + for (const auto& it : props) { + if (dev.has(it)) continue; + switch (it) { + case sycl::aspect::fp64: + throw std::runtime_error("'double' is not supported in '" + + dev.get_info() + "' device"); + break; + case sycl::aspect::fp16: + throw std::runtime_error("'half' is not supported in '" + + dev.get_info() + "' device"); + break; + default: +#define __SYCL_ASPECT(ASPECT, ID) \ + case sycl::aspect::ASPECT: return #ASPECT; +#define __SYCL_ASPECT_DEPRECATED(ASPECT, ID, MESSAGE) __SYCL_ASPECT(ASPECT, ID) +#define __SYCL_ASPECT_DEPRECATED_ALIAS(ASPECT, ID, MESSAGE) + auto getAspectNameStr = [](sycl::aspect AspectNum) -> std::string { + switch (AspectNum) { +#include +#include + default: return "unknown aspect"; + } + }; +#undef __SYCL_ASPECT_DEPRECATED_ALIAS +#undef __SYCL_ASPECT_DEPRECATED +#undef __SYCL_ASPECT + throw std::runtime_error("'" + getAspectNameStr(it) + "' is not supported in '" + + dev.get_info() + "' device"); + } + break; + } +} + +void param_update_kernel(const float* input, sycl::half* output, int size) +{ + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); + + if (id < size) { output[id] = (sycl::half)input[id]; } +} + +void launch_param_update(const float* input, sycl::half* output, int size, sycl::queue* stream) +{ + int threads = 1024; + + sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1); + sycl::range<3> block_dim(1, 1, threads); + + { + has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); + stream->parallel_for( + sycl::nd_range<3>(grid_dim * block_dim, block_dim), + [=](sycl::nd_item<3> item_ct1) { param_update_kernel(input, output, size); }); + } +} + +void param_update_kernel_half(const float* input, sycl::half* output, int size) +{ + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); + sycl::half2* output_cast = reinterpret_cast(output); + if (id < size) { + float input_f = input[id]; + sycl::half2* input_h = reinterpret_cast(&input_f); + output_cast[id] = *input_h; + } +} + +void launch_param_update_half(const float* input, sycl::half* output, int size, sycl::queue* stream) +{ + int threads = 1024; + size /= 2; + sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1); + sycl::range<3> block_dim(1, 1, threads); + + { + has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); + stream->parallel_for( + sycl::nd_range<3>(grid_dim * block_dim, block_dim), + [=](sycl::nd_item<3> item_ct1) { param_update_kernel_half(input, output, size); }); + } +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h new file mode 100644 index 0000000000000000000000000000000000000000..6d54446d472e9dfa6141436f55e0e623be48acb4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#ifndef TORCH_CHECK +#define TORCH_CHECK AT_CHECK +#endif + +#ifdef VERSION_GE_1_3 +#define DATA_PTR data_ptr +#else +#define DATA_PTR data +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h new file mode 100644 index 0000000000000000000000000000000000000000..660f860917f6c8b345825e6110a7508ba19a584f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include "simd.h" + +typedef unsigned short ds_half_precision_t; + +#define STEP(SPAN) \ + void Step_##SPAN(float* _params, \ + float* grads, \ + float* _exp_avg_sq, \ + size_t _param_size, \ + ds_half_precision_t* dev_param = nullptr, \ + bool half_precision = false); + +class Adagrad_Optimizer { +public: + Adagrad_Optimizer(float alpha = 1e-2, float eps = 1e-8, float weight_decay = 0) + : _alpha(alpha), _eps(eps), _weight_decay(weight_decay) + { + } + ~Adagrad_Optimizer() {} +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg_sq, + size_t param_size, + ds_half_precision_t* dev_param = nullptr, + bool half_precision = false); +#endif + STEP(1) + STEP(4) + STEP(8) + inline void IncrementStep(size_t step) + { + _step++; + if (_step != step) { _step = step; } + } + inline void update_state(float lr, float epsilon, float weight_decay) + { + _alpha = lr; + _eps = epsilon; + _weight_decay = weight_decay; + } + +private: + float _alpha; + float _eps; + float _weight_decay; + + float _betta1_t; + float _betta2_t; + size_t _step; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t new_rounded_size = 0; + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + AVX_Data weight_decay4; + if (_weight_decay > 0) weight_decay4.data = SIMD_SET(_weight_decay); + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i, half_precision); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, grads + i, false); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i, false); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i, half_precision); + + if (_weight_decay > 0) { simd_fma(grad_4, param_4, weight_decay4, grad_4); } + + simd_fma(variance_4, grad_4, grad_4, variance_4); + simd_sqrt(grad_4, variance_4); + simd_add(grad_4, grad_4, eps_4); + simd_div(grad_4, momentum_4, grad_4); + simd_fma(param_4, grad_4, step_size_4, param_4); + + simd_store(_params + i, param_4, half_precision); + simd_store(_exp_avg_sq + i, variance_4, false); + } + } + *rounded_size = new_rounded_size; +} +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h new file mode 100644 index 0000000000000000000000000000000000000000..7bc0364c569d00d01b9f0b526d21d3c57fb478dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h @@ -0,0 +1,237 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include +#include "simd.h" + +#include +typedef unsigned short ds_half_precision_t; + +#define STEP(SPAN) \ + void Step_##SPAN(float* _params, \ + float* grads, \ + float* _exp_avg, \ + float* _exp_avg_sq, \ + size_t _param_size, \ + ds_half_precision_t* dev_param = nullptr, \ + bool half_precision = false); + +class Adam_Optimizer { +public: + Adam_Optimizer(float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float eps = 1e-8, + float weight_decay = 0, + bool adamw_mode = true) + : _alpha(alpha), + _betta1(betta1), + _betta2(betta2), + _eps(eps), + _weight_decay(weight_decay), + _betta1_t(1.0), + _betta2_t(1.0), + _step(0), + _adamw_mode(adamw_mode) + { + } + ~Adam_Optimizer() {} + +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t param_size, + ds_half_precision_t* dev_param = nullptr, + bool half_precision = false); +#endif + STEP(1) + STEP(4) + STEP(8) + inline void IncrementStep(size_t step, float beta1, float beta2) + { + if (beta1 != _betta1 || beta2 != _betta2) { + _step = step; + _betta1 = beta1; + _betta2 = beta2; + _betta1_t = std::pow(_betta1, step); + _betta2_t = std::pow(_betta2, step); + } else { + _step++; + if (_step != step) { + _betta1_t = std::pow(_betta1, step); + _betta2_t = std::pow(_betta2, step); + _step = step; + } else { + _betta1_t *= _betta1; + _betta2_t *= _betta2; + } + } + } + inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) + { + _alpha = lr; + _eps = epsilon; + _weight_decay = weight_decay; + + _bias_correction1 = 1.0f; + _bias_correction2 = 1.0f; + if (bias_correction == 1) { + _bias_correction1 = 1 - _betta1_t; + _bias_correction2 = 1 / sqrt(1 - _betta2_t); + } + } + +private: + float _alpha; + float _betta1; + float _betta2; + float _eps; + float _weight_decay; + + float _betta1_t; + float _betta2_t; + size_t _step; + + float _bias_correction1; + float _bias_correction2; + + bool _adamw_mode; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Adam_Optimizer::Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t new_rounded_size = 0; + int rshft = half_precision ? 1 : 0; + + AVX_Data betta1_4; + betta1_4.data = SIMD_SET(_betta1); + AVX_Data betta2_4; + betta2_4.data = SIMD_SET(_betta2); + + float betta1_minus1 = 1 - _betta1; + float betta2_minus1 = 1 - _betta2; + AVX_Data betta1_minus1_4; + betta1_minus1_4.data = SIMD_SET(betta1_minus1); + AVX_Data betta2_minus1_4; + betta2_minus1_4.data = SIMD_SET(betta2_minus1); + + AVX_Data bias2_sqrt; + bias2_sqrt.data = SIMD_SET(_bias_correction2); + + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha / _bias_correction1; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + float w_decay = -1 * _alpha * _weight_decay; + AVX_Data weight_decay4; + if (_weight_decay > 0) + weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + (i >> rshft), half_precision); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, _exp_avg + i, false); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i, false); + + AVX_Data param_4[span]; + simd_load(param_4, _params + (i >> rshft), half_precision); + + if (_weight_decay > 0 && !_adamw_mode) { + simd_fma(grad_4, param_4, weight_decay4, grad_4); + } + + simd_mul(momentum_4, momentum_4, betta1_4); + simd_fma(momentum_4, grad_4, betta1_minus1_4, momentum_4); + simd_mul(variance_4, variance_4, betta2_4); + simd_mul(grad_4, grad_4, grad_4); + simd_fma(variance_4, grad_4, betta2_minus1_4, variance_4); + simd_sqrt(grad_4, variance_4); + simd_fma(grad_4, grad_4, bias2_sqrt, eps_4); + simd_div(grad_4, momentum_4, grad_4); + + if (_weight_decay > 0 && _adamw_mode) { + simd_fma(param_4, param_4, weight_decay4, param_4); + } + + simd_fma(param_4, grad_4, step_size_4, param_4); + + simd_store(_params + (i >> rshft), param_4, half_precision); + simd_store(_exp_avg + i, momentum_4, false); + simd_store(_exp_avg_sq + i, variance_4, false); + } + } + *rounded_size = new_rounded_size; +} +#endif + +int create_adam_optimizer(int optimizer_id, + float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float eps = 1e-8, + float weight_decay = 0, + bool adamw_mode = true, + bool should_log = false); + +int ds_adam_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + bool bias_correction, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq); + +int ds_adam_step_plus_copy(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + bool bias_correction, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq, + torch::Tensor& gpu_params); + +int destroy_adam_optimizer(int optimizer_id); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h new file mode 100644 index 0000000000000000000000000000000000000000..097e2d8585ccbdb83f1c249c1b6d89272a134999 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h @@ -0,0 +1,198 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#if (__x86_64__ || __i386__) +#include +#include +#endif + +#define TILE (128 * 1024 * 1024) +#if defined(__AVX512__) or defined(__AVX256__) + +#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) + +#if defined(__AVX512__) +#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d) +#define SIMD_LOAD(x) _mm512_loadu_ps(x) +#define SIMD_SET(x) _mm512_set1_ps(x) +#define SIMD_ADD(x, y) _mm512_add_ps(x, y) +#define SIMD_MUL(x, y) _mm512_mul_ps(x, y) +#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) +#define SIMD_SQRT(x) _mm512_sqrt_ps(x) +#define SIMD_DIV(x, y) _mm512_div_ps(x, y) +#define SIMD_AND(x, y) _mm512_and_ps(x, y) +#define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) +#define SIMD_OR(x, y) _mm512_or_ps(x, y) +#define SIMD_XOR(x, y) _mm512_xor_ps(x, y) +#define SIMD_WIDTH 16 + +#define SIMD_LOAD2(x, h) \ + ((h) ? _mm512_cvtph_ps(_mm256_castps_si256(_mm256_loadu_ps(x))) : _mm512_loadu_ps(x)) +#define SIMD_STORE2(x, d, h) \ + ((h) ? _mm256_store_ps(x, _mm256_castsi256_ps(_mm512_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) \ + : _mm512_storeu_ps(x, d)) + +#define INTV __m256i +#elif defined(__AVX256__) +#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d) +#define SIMD_LOAD(x) _mm256_loadu_ps(x) +#define SIMD_SET(x) _mm256_set1_ps(x) +#define SIMD_ADD(x, y) _mm256_add_ps(x, y) +#define SIMD_MUL(x, y) _mm256_mul_ps(x, y) +#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) +#define SIMD_SQRT(x) _mm256_sqrt_ps(x) +#define SIMD_DIV(x, y) _mm256_div_ps(x, y) +#define SIMD_AND(x, y) _mm256_and_ps(x, y) +#define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) +#define SIMD_OR(x, y) _mm256_or_ps(x, y) +#define SIMD_XOR(x, y) _mm256_xor_ps(x, y) +#define SIMD_WIDTH 8 + +#define SIMD_LOAD2(x, h) \ + ((h) ? _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*)x)) : _mm256_loadu_ps(x)) +#define SIMD_STORE2(x, d, h) \ + ((h) ? _mm_store_ps(x, _mm_castsi128_ps(_mm256_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) \ + : _mm256_storeu_ps(x, d)) + +#define INTV __m128i +#endif + +union AVX_Data { +#if defined(__AVX512__) + __m512 data; +#elif defined(__AVX256__) + __m256 data; +#endif + // float data_f[16]; +}; + +template +inline void simd_store(float* dst, AVX_Data* src, bool half_precision) +{ + size_t width = (half_precision ? SIMD_WIDTH / 2 : SIMD_WIDTH); +#pragma unroll + for (size_t i = 0; i < span; ++i) { SIMD_STORE2(dst + width * i, src[i].data, half_precision); } +} +template +inline void simd_load(AVX_Data* dst, float* src, bool half_precision) +{ + size_t width = (half_precision ? 1 : SIMD_WIDTH); +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD2(src + width * i, half_precision); } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data* src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a[i].data); + } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a.data); + } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Data* src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r[i].data, src_a[i].data); + } +} +template +inline void simd_sqrt(AVX_Data* dst, AVX_Data* src) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_SQRT(src[i].data); } +} +template +inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r[i].data); + } +} +template +inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r[i].data); } +} + +#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h new file mode 100644 index 0000000000000000000000000000000000000000..1897afd1fea248f497acf3fb6ba049baa1ca56c3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h @@ -0,0 +1,155 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* Taken from NVIDIA/apex commit 855808f3fc268e9715d613f3c2e56469d8c986d8 */ +#include +/* #include */ +#include + +// Forward/backward compatibility hack around +// https://github.com/pytorch/pytorch/commit/3aeb78079bcd68282fe9117088e138b77318e288 +// pending more future-proof guidance from upstream. +// struct TypeShim +// { +// const at::Type& payload; +// TypeShim(const at::Type& type) : payload(type) {} +// // Enable trivial conversion to a const at::Type& for pre-3aeb78 +// operator const at::Type&(){ return payload; }; +// // Enable dispatch switch statements to take *this directly for post-3aeb78 +// //operator at::ScalarType(){ return payload.; }; +// }; + +#define DISPATCH_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t_##LEVEL = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t_##LEVEL = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +#define DISPATCH_DOUBLE_AND_FLOAT(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +template +__inline__ __attribute__((always_inline)) T +reduce_block_into_lanes(T* x, + T val, + int lanes = 1, + bool share_result = false) // lanes is intended to be <= 32. +{ + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int tid = item_ct1.get_local_id(2) + item_ct1.get_local_id(1) * item_ct1.get_local_range(2); + int blockSize = item_ct1.get_local_range(2) * + item_ct1.get_local_range(1); // blockSize is intended to be a multiple of 32. + + if (blockSize >= 64) { + x[tid] = val; + /* + DPCT1118:1: SYCL group functions and algorithms must be encountered in converged control + flow. You may need to adjust the code. + */ + /* + DPCT1065:6: Consider replacing sycl::nd_item::barrier() with + sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if + there is no access to global memory. + */ + item_ct1.barrier(); + } + +#pragma unroll + for (int i = (blockSize >> 1); i >= 64; i >>= 1) { + if (tid < i) x[tid] = x[tid] + x[tid + i]; + /* + DPCT1118:2: SYCL group functions and algorithms must be encountered in converged control + flow. You may need to adjust the code. + */ + /* + DPCT1065:7: Consider replacing sycl::nd_item::barrier() with + sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if + there is no access to global memory. + */ + item_ct1.barrier(); + } + + T final; + + if (tid < 32) { + if (blockSize >= 64) + final = x[tid] + x[tid + 32]; + else + final = val; + // __SYNCWARP(); + +#pragma unroll + for (int i = 16; i >= lanes; i >>= 1) + final = final + __shfl_down_sync(0xffffffff, final, i); + } + + if (share_result) { + if (tid < lanes) x[tid] = final; // EpilogueOp + // Make sure the smem result is visible to all warps. + /* + DPCT1118:3: SYCL group functions and algorithms must be encountered in converged control + flow. You may need to adjust the code. + */ + /* + DPCT1065:8: Consider replacing sycl::nd_item::barrier() with + sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if + there is no access to global memory. + */ + item_ct1.barrier(); + } + + return final; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp new file mode 100644 index 0000000000000000000000000000000000000000..fb09bb16efdda5b578c239c35be7ee493dc87c75 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include + +using namespace sycl; +using namespace xpu; + +void packbitskernel(const float* input, uint8_t* output, const int input_size, id<1> item_ct1) +{ + // get the sign bit of each float and pack them into byte + int i = item_ct1; + for (int j = 0; j < 8; ++j) { + int k = i * 8 + j; + int bit = k < input_size && (!sycl::signbit(input[k])); + output[i] |= bit << (7 - j); + } +} + +void unpackbitskernel(const uint8_t* input, float* output, id<1> item_ct1) +{ + // use the bit value to set float, bit 0 -> float -1, bit 1 -> float 1 + int i = item_ct1; + output[i] = (float((input[i / 8] >> (7 - i % 8)) & 1) - 0.5) * 2; +} + +sycl::queue get_current_queue(at::Device device) +{ + c10::impl::VirtualGuardImpl impl(device.type()); + c10::Stream _stream = impl.getStreamFromGlobalPool(device, /*isHighPriority=*/false); + sycl::queue queue = xpu::get_queue_from_stream(_stream); + return queue; +} + +/* +pack float tensor into uint8 tensor. Every eight float elements get packed into one uint8 +if float x >= 0, will be packed as a '1' bit, or will be packed as '0' +Arguments: + tensor: A bool tensor that get packed. + input_size: numel of input tensor + rank: device id in order to get corresponding stream +*/ +at::Tensor packbits(at::Tensor tensor, int input_size, int rank) +{ + at::Device device = "xpu:" + std::to_string(rank); + sycl::queue q = get_current_queue(device); + + int packed_size = (input_size + 7) / 8; + auto unit8_options = at::TensorOptions().dtype(at::kByte).device(at::kXPU); + at::Tensor packed = torch::zeros({packed_size}, unit8_options); + + float* input = (float*)tensor.data_ptr(); + uint8_t* output = (uint8_t*)packed.data_ptr(); + + auto event = q.submit([&](sycl::handler& cgh) { + cgh.parallel_for<>(range(packed_size), [=](id<1> item_ct1) { + packbitskernel(input, output, input_size, item_ct1); + }); + }); + + return packed; +} + +/* +unpack uint8 tensor into float tensor. Every uint8 element get unpacked into eight float +a '1' bit will be converted to a float(1), a '0' bit will be converted to a float(-1). +Arguments: + tensor: A uint8 tensor that get unpacked. + input_size: numel of input tensor + rank: device id in order to get corresponding stream +*/ +at::Tensor unpackbits(at::Tensor tensor, int input_size, int rank) +{ + at::Device device = "xpu:" + std::to_string(rank); + sycl::queue q = get_current_queue(device); + + auto float_options = at::TensorOptions().dtype(at::kFloat).device(at::kXPU); + at::Tensor unpacked = torch::empty({input_size * 8}, float_options); + + uint8_t* input = (uint8_t*)tensor.data_ptr(); + float* output = (float*)unpacked.data_ptr(); + + auto event = q.submit([&](sycl::handler& cgh) { + cgh.parallel_for<>(range(input_size * 8), + [=](id<1> item_ct1) { unpackbitskernel(input, output, item_ct1); }); + }); + + return unpacked; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("packbits", &packbits, "DeepSpeed XPU packbits (C++)"); + m.def("unpackbits", &unpackbits, "DeepSpeed XPU unpackbits (C++)"); +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c5fd280fc32661bae081ad8824c23efa36c9c97 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .evoformer_attn import DS4Sci_EvoformerAttention, EvoformerFusedAttention diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a12e2f2a54200db66cc3c0ff9fe637c23c37b6ab Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..572e6cc40a018cedc713cd4f67326c33eab8e2c3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/evoformer_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/evoformer_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..da5843d6de31ed94c27c466bbcf5ae5966bb620f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/evoformer_attn.py @@ -0,0 +1,106 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import numpy as np +from deepspeed.ops.op_builder import EvoformerAttnBuilder +from deepspeed.accelerator import get_accelerator + +kernel_ = None + + +def _attention(Q, K, V, bias1, bias2): + assert Q.shape[-3] > 16, "seq_len must be greater than 16" + O = torch.empty_like(Q, dtype=Q.dtype) + assert get_accelerator().on_accelerator(Q), "Q must be on cuda" + assert get_accelerator().on_accelerator(K), "K must be on cuda" + assert get_accelerator().on_accelerator(V), "V must be on cuda" + assert get_accelerator().on_accelerator(bias1), "bias1 must be on cuda" + assert get_accelerator().on_accelerator(bias2), "bias2 must be on cuda" + global kernel_ + if kernel_ is None: + kernel_ = EvoformerAttnBuilder().load() + nheads = Q.shape[-2] + nq = (Q.shape[-3] + 31) // 32 * 32 + nb = np.prod(Q.shape[:-3]) + lse = torch.empty((nb, nheads, nq), dtype=torch.float32, device=Q.device) + kernel_.attention(Q, K, V, bias1, bias2, O, lse) + return O, lse + + +def attention_bwd(dO, Q, K, V, O, lse, bias1, bias2, bias1_grad, bias2_grad): + assert max(Q.shape[-1], V.shape[-1]) <= 64, "Hidden size is too large. Need to change kMax to a larger value" + dQ = torch.empty_like(Q, dtype=Q.dtype) + dK = torch.empty_like(K, dtype=K.dtype) + dV = torch.empty_like(V, dtype=V.dtype) + assert get_accelerator().on_accelerator(dO), "dO must be on cuda" + assert get_accelerator().on_accelerator(Q), "Q must be on cuda" + assert get_accelerator().on_accelerator(K), "K must be on cuda" + assert get_accelerator().on_accelerator(V), "V must be on cuda" + assert get_accelerator().on_accelerator(O), "O must be on cuda" + global kernel_ + if kernel_ is None: + kernel_ = EvoformerAttnBuilder().load() + delta = torch.empty_like(lse) + if bias1_grad: + dB1 = torch.zeros_like(bias1, dtype=torch.float32) + else: + dB1 = torch.tensor([], dtype=torch.float32, device=bias1.device) + if bias2_grad: + dB2 = torch.zeros_like(bias2, dtype=torch.float32) + else: + dB2 = torch.tensor([], dtype=torch.float32, device=bias2.device) + kernel_.attention_bwd(dO, Q, K, V, O, lse, delta, bias1, bias2, dQ, dK, dV, dB1, dB2) + return dQ, dK, dV, dB1.to(dO.dtype), dB2.to(dO.dtype) + + +class EvoformerFusedAttention(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, bias1=None, bias2=None): + """ + q, k, v: are in shape [*, L, H, D] + """ + bias1_ = bias1.contiguous() if bias1 is not None else torch.tensor([], dtype=q.dtype, device=q.device) + bias2_ = bias2.contiguous() if bias2 is not None else torch.tensor([], dtype=q.dtype, device=q.device) + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + o, lse = _attention(q, k, v, bias1_, bias2_) + ctx.save_for_backward(q, k, v, o, lse, bias1_, bias2_) + return o + + @staticmethod + def backward(ctx, grad_output): + q, k, v, o, lse, bias1, bias2 = ctx.saved_tensors + is_b1_grad = bias1.numel() != 0 and ctx.needs_input_grad[3] + is_b2_grad = bias2.numel() != 0 and ctx.needs_input_grad[4] + dQ, dK, dV, dB1, dB2 = attention_bwd(grad_output, q, k, v, o, lse, bias1, bias2, is_b1_grad, is_b2_grad) + if not is_b1_grad: + dB1 = None + if not is_b2_grad: + dB2 = None + return dQ, dK, dV, dB1, dB2 + + +def DS4Sci_EvoformerAttention(Q, K, V, biases): + assert len(biases) <= 2 + + if (len(biases) == 0): + biases.append(None) + + if (len(biases) == 1): + biases.append(None) + + bias_1_shape = lambda x: (x.shape[0], x.shape[1], 1, 1, x.shape[2]) + bias_2_shape = lambda x: (x.shape[0], 1, x.shape[3], x.shape[2], x.shape[2]) + + if biases[0] is not None: + assert biases[0].shape == bias_1_shape(Q), "bias1 shape is incorrect" + + if biases[1] is not None: + assert biases[1].shape == bias_2_shape(Q), "bias2 shape is incorrect" + + return EvoformerFusedAttention.apply(Q, K, V, biases[0], biases[1]) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9cf23373c265a1c0521a9eaa3b713cdd7aca57c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .quantize import FP_Quantize, Quantizer +from .fp8_gemm import matmul_fp8 diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bdb0d9c972e4d6353b211a31e59fc9105786d6d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24194fccc374fa7f9cd20898e6df9753d0ebcc2e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b973add6fe3e02b5e5138e3588ffef65323ee64a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2761ae6b6d35c85c520ad268a07204d1bb8a372 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..db4fa5ae2c9271b480b6d209da28f0f4fdbbe85d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######## Fused MoE kernel ######### +# These kernels are implemented for +# fusing GeMM with dequantization of +# fp8 weight data when using bit-16 +# activation. +################################### + +import torch + + +def matmul_fp8(inp, weight, scale, quantization_group_size, quantizer): + from deepspeed import get_accelerator + + if not get_accelerator().is_triton_supported(): + return matmul_fp8_fallback(inp, weight, scale, quantization_group_size, quantizer) + else: + # Import dynamically to prevent failures on systems without triton. + from .fp8_gemm_triton import matmul_fp8_triton + return matmul_fp8_triton(inp, weight, scale, quantization_group_size) + + +def matmul_fp8_fallback(inp, weight, scale, quantization_group_size, quantizer): + return torch.matmul(inp, quantizer.dequantize(weight, scale=scale)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py new file mode 100644 index 0000000000000000000000000000000000000000..086525cc64425558399886cc534c57a039ee0a41 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py @@ -0,0 +1,165 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######## Fused MoE kernel ######### +# These kernels are implemented for +# fusing GeMM with dequantization of +# fp8 weight data when using bit-16 +# activation. +################################### + +import torch +import triton +import triton.language as tl + + +@triton.jit +def matmul_kernel_fp8_bf16(inp_ptr, weight_ptr, out_ptr, scale_ptr, M, N, K, stride_am, stride_ak, stride_bk, + stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + quantization_group_size: tl.constexpr): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + + inp_data = inp_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + weight_data = weight_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + weight_ptrs_offset = offs_k[:, None] * (stride_bk // quantization_group_size) + ( + (pid_n * BLOCK_SIZE_N) // quantization_group_size) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + inp = tl.load(inp_data, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + weight = tl.load(weight_data, mask=offs_k[:, None] < K, other=0.0) + scale = tl.load(scale_ptr + weight_ptrs_offset + ((k * BLOCK_SIZE_K * stride_bk) // quantization_group_size)) + # Dequantize weight (fp8 -> bf16) + w = (weight & 0x80).to(tl.uint16) << 8 + w = w | ((weight & 0x7f).to(tl.uint16) << 4) + w = (w + 0x3C00).to(tl.uint16) + w = (w.to(tl.bfloat16, bitcast=True).to(tl.float32) * scale).to(tl.bfloat16) + + inp_data += BLOCK_SIZE_K * stride_ak + weight_data += BLOCK_SIZE_K * stride_bk + + accumulator += tl.dot(inp, w) + + out = accumulator.to(tl.bfloat16) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + out_data = out_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + tl.store(out_data, out, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N)) + + +@triton.jit +def matmul_kernel_fp8_fp16(inp_ptr, weight_ptr, out_ptr, scale_ptr, M, N, K, stride_am, stride_ak, stride_bk, + stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + quantization_group_size: tl.constexpr): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + + inp_data = inp_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + weight_data = weight_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + weight_ptrs_offset = offs_k[:, None] * (stride_bk // quantization_group_size) + ( + (pid_n * BLOCK_SIZE_N) // quantization_group_size) + + weight = tl.load(weight_data, mask=offs_k[:, None] < K, other=0.0) + scale = tl.load(scale_ptr + weight_ptrs_offset) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + inp = tl.load(inp_data, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + # Dequantize weight (fp8 -> fp16) + w = (((weight & 0x80) << 8) | ((weight & 0x7f) << 7)).to(tl.uint16) + w = (w + 0x2000).to(tl.uint16) + w = (w.to(tl.float16, bitcast=True) * scale).to(tl.float16) + + inp_data += BLOCK_SIZE_K * stride_ak + weight_data += BLOCK_SIZE_K * stride_bk + + weight = tl.load(weight_data, mask=offs_k[:, None] < K - (k + 1) * BLOCK_SIZE_K, other=0.0) + scale = tl.load(scale_ptr + (weight_ptrs_offset + + (((k + 1) * BLOCK_SIZE_K * stride_bk) // quantization_group_size))) + + accumulator += tl.dot(inp, w) + + out = accumulator.to(tl.float16) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + out_data = out_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + tl.store(out_data, out, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N)) + + +def matmul_fp8_triton(inp, weight, scale, quantization_group_size): + + assert inp.shape[1] == weight.shape[0], \ + f"Incompatible dimensions (input: {inp.shape}, weight: {weight.shape})" + + M, K = inp.shape + K, N = weight.shape + + out = torch.empty((M, N), device=inp.device, dtype=inp.dtype) + + # GEMM tuning parameters! + # TODO: Add a more configurable tuning for selecting the best GeMM + BLOCK_SIZE_M = 16 if M <= 16 else 32 if M <= 32 else 64 if M <= 64 else 128 + BLOCK_SIZE_N = 64 + BLOCK_SIZE_K = max(64, quantization_group_size) + GROUP_SIZE_M = 8 + num_stages = 4 + num_warps = 4 + if M >= 256: + BLOCK_SIZE_M = 256 + BLOCK_SIZE_N = 128 + BLOCK_SIZE_K = max(128, quantization_group_size) + num_stages = 3 + num_warps = 8 + + grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), ) + kernel = matmul_kernel_fp8_bf16 if inp.dtype == torch.bfloat16 else matmul_kernel_fp8_fp16 + kernel[grid](inp, + weight, + out, + scale, + M, + N, + K, + inp.stride(0), + inp.stride(1), + weight.stride(0), + weight.stride(1), + out.stride(0), + out.stride(1), + quantization_group_size=quantization_group_size, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=GROUP_SIZE_M, + num_stages=num_stages, + num_warps=num_warps) + return out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/quantize.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..47b3b08c7e03d94a98cc6c79af071beb82e40860 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/quantize.py @@ -0,0 +1,167 @@ +# 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): + + def __init__(self, quantization_config) -> None: + global fp_quant_module + super().__init__(group_size=quantization_config.group_size) + if fp_quant_module is None: + fp_quant_module = FPQuantizerBuilder().load() + self.cuda_impl = getattr(fp_quant_module, "CUDA_IMPL", True) + self.q_config = quantization_config + + self.orig_dtype = None + self.num_groups = None + self.input_q = None + self.scale = None + + def quantize(self, + input, + q_bits=8, + q_mantisa_bits=3, + stochastic_mode=False, + return_meta_tensor=False) -> torch.Tensor: + assert input.dtype == torch.bfloat16, "only support bf16 for now" + if return_meta_tensor: + assert q_bits == 8, "meta tensor is only supported with q_bit=8" + + self.orig_dtype = input.dtype + self.orig_shape = input.shape + + if q_bits == 8: + pass + elif q_bits == 12: + q_mantisa_bits = 4 + elif q_bits == 6: + q_mantisa_bits = 2 + elif q_bits == 4: + q_mantisa_bits = 1 + else: + assert (0), \ + f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!" + self.num_groups = input.numel() // self.group_size + self.input_q = torch.ones(self.num_groups, + int(self.group_size * q_bits) // 8 + 4, + dtype=torch.uint8, + device=input.device) + out = fp_quant_module.quantize(self.input_q, input, self.group_size, stochastic_mode, q_bits, q_mantisa_bits) + if return_meta_tensor: + data, self.scale = out.split(self.group_size, dim=-1) + data = data.contiguous().reshape(input.shape) + self.scale = self.scale.contiguous() + del self.input_q + del out + gc.collect() + get_accelerator().empty_cache() + return data, self.scale + + return out + + def to(self, *args, **kwargs): + # Intermediate tensors may need to be moved to different devices + if hasattr(self, 'input_q'): + self.input_q = self.input_q.to(*args, **kwargs) + if hasattr(self, 'scale'): + self.scale = self.scale.to(*args, **kwargs) + + def get_scales(self): + return fp_quant_module.get_scales(self.scale, self.num_groups) + + def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: + assert (self.orig_dtype is not None), \ + "[De-quantization Error]: you need to call quantize before dequantizing!" + fp_out = torch.empty(self.orig_shape, dtype=self.orig_dtype, + device=input_q.device) if fp_out is None else fp_out + if q_bits == 8: + pass + elif q_bits == 12: + q_mantisa_bits = 4 + elif q_bits == 6: + q_mantisa_bits = 2 + elif q_bits == 4: + q_mantisa_bits = 1 + else: + assert (0), \ + f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" + + if scale is not None: + assert input_q.numel() == fp_out.numel(), \ + f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' + input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() + fp_quant_module.dequantize(fp_out, input_q, self.group_size, q_mantisa_bits, q_bits - q_mantisa_bits - 1) + return fp_out + + def selective_dequantize(self, + input_q, + indexes, + fp_out=None, + q_bits=8, + q_mantisa_bits=3, + scale=None) -> torch.Tensor: + assert (not hasattr(self, 'orig_shape') or len(self.orig_shape) == 3), \ + "Selective-Dequantization works on 3d tensor only! Please reshape the tensor before calling dequantize function." + assert (self.orig_dtype is not None), \ + "[De-quantization Error]: you need to call quantize before dequantizing!" + fp_out = torch.empty( + (indexes.shape[0], + *self.orig_shape[1:]), dtype=self.orig_dtype, device=input_q.device) if fp_out is None else fp_out + if q_bits == 8: + pass + elif q_bits == 12: + q_mantisa_bits = 4 + elif q_bits == 6: + q_mantisa_bits = 2 + elif q_bits == 4: + q_mantisa_bits = 1 + else: + assert (0), \ + f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" + + if scale is not None: + assert input_q.numel() == fp_out.numel(), \ + f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' + input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() + + fp_quant_module.selective_dequantize(fp_out, input_q, indexes, self.group_size, q_mantisa_bits, + q_bits - q_mantisa_bits - 1) + return fp_out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0762c810766ebc7bc832e737d7e773216475b9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..op_builder import GDSBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e3498d2be482d8e2b822b5ca1445fcfd9e2b4273 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..81bc07e827c88e3f7c134eab78c761813fdf7826 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .fused_lamb import FusedLamb diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1bef2dbda816956cc8737c73b9ddf832c5391992 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..819f035891f9b80e747479172f986bb159593ad1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/fused_lamb.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/fused_lamb.py new file mode 100644 index 0000000000000000000000000000000000000000..6ccd9d4c6b066601bb2f6e5d8d8def6be09fd22b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/fused_lamb.py @@ -0,0 +1,174 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from NVIDIA/apex/optimizer/fused_adam and implements the LAMB optimizer +""" +import types +import torch +from deepspeed.ops.op_builder import FusedLambBuilder + + +class FusedLamb(torch.optim.Optimizer): + """Implements the LAMB algorithm. Currently GPU-only. + + LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. + https://arxiv.org/abs/1904.00962 + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + bias_correction (bool, optional): bias correction (default: True) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + max_grad_norm (float, optional): value used to clip global grad norm + (default: 0.0) + max_coeff(float, optional): maximum value of the lamb coefficient (default: 10.0) + min_coeff(float, optional): minimum value of the lamb coefficient (default: 0.01) + amsgrad (boolean, optional): NOT SUPPORTED in FusedLamb! + """ + + def __init__(self, + params, + lr=1e-3, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + max_coeff=10.0, + min_coeff=0.01, + amsgrad=False): + self.fused_lamb_cuda = FusedLambBuilder().load() + + if amsgrad: + raise RuntimeError('FusedLamb does not support the AMSGrad variant.') + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm, + max_coeff=max_coeff, + min_coeff=min_coeff) + super(FusedLamb, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.lamb_coeffs = [] + + def step(self, closure=None, grads=None, output_params=None, scale=1., grad_norms=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + output params (list of tensors, optional): A reduced precision copy + of the updated weights written out in addition to the regular + updated weights. Have to be of same type as gradients. (default: None) + scale (float, optional): factor to divide gradient tensor values + by before applying to weights. (default: 1) + """ + loss = None + if closure is not None: + loss = closure() + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + if output_params is None: + output_params_group = [None] * len(self.param_groups) + elif isinstance(output_params, types.GeneratorType): + output_params_group = [output_params] + elif type(output_params[0]) != list: + output_params_group = [output_params] + else: + output_params_group = output_params + + if grad_norms is None: + grad_norms = [None] * len(self.param_groups) + + #remove the previous coeffs + del self.lamb_coeffs[:] + + for group, grads_this_group, output_params_this_group, grad_norm_group in zip( + self.param_groups, grads_group, output_params_group, grad_norms): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + if output_params_this_group is None: + output_params_this_group = [None] * len(group['params']) + + if grad_norm_group is None: + grad_norm_group = [None] * len(group['params']) + elif not isinstance(grad_norm_group, list): + grad_norm_group = [grad_norm_group] + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad, output_param, grad_norm in zip(group['params'], grads_this_group, output_params_this_group, + grad_norm_group): + + # compute combined scale factor for this group + combined_scale = scale + if group['max_grad_norm'] > 0: + # norm is in fact norm*scale + clip = ((grad_norm / scale) + 1e-6) / group['max_grad_norm'] + if clip > 1: + combined_scale = clip * scale + + #note: p.grad should not ever be set for correct operation of mixed precision optimizer that sometimes sends None gradients + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('FusedLamb does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + max_coeff = group['max_coeff'] + min_coeff = group['min_coeff'] + + state['step'] += 1 + + out_p = torch.tensor([], dtype=torch.float) if output_param is None else output_param + lamb_coeff = self.fused_lamb_cuda.lamb(p.data, out_p, exp_avg, exp_avg_sq, grad, group['lr'], beta1, + beta2, max_coeff, min_coeff, group['eps'], combined_scale, + state['step'], self.eps_mode, bias_correction, + group['weight_decay']) + self.lamb_coeffs.append(lamb_coeff) + return loss + + def get_lamb_coeffs(self): + lamb_coeffs = [lamb_coeff.item() for lamb_coeff in self.lamb_coeffs] + return lamb_coeffs diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2f90e5ec2e808057a26cc36f2d7caedecb44c86e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .cpu_lion import DeepSpeedCPULion +from .fused_lion import FusedLion diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..915ddd5b12c8e9624e97bae763c4d3c9277d7cd4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57689253a6ccf0e599cacf22a1ee7603efb48fd0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d484992084b17f9527f6d12f870b9b00b8aedcf Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4f0420c5f3b895d5cd24f591fa52ce6fc56aadb Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/cpu_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/cpu_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..03342a3fcd34fdb852434a97e03fa7f99426c7d0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/cpu_lion.py @@ -0,0 +1,127 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from cpuinfo import get_cpu_info +from deepspeed.utils import logger +from deepspeed.utils.logging import should_log_le +from deepspeed.ops.op_builder import CPULionBuilder + + +class DeepSpeedCPULion(torch.optim.Optimizer): + optimizer_id = 0 + + def __init__(self, model_params, lr=1e-3, betas=(0.9, 0.999), weight_decay=0, fp32_optimizer_states=True): + """Fast vectorized implementation of Lion optimizer on CPU: + + See Symbolic Discovery of Optimization Algorithms (https://doi.org/10.48550/arXiv.2302.06675). + + .. note:: + We recommend using our `config + `_ + to allow :meth:`deepspeed.initialize` to build this optimizer + for you. + + + Arguments: + model_params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + full_precision_optimizer_states: creates momentum and variance in full precision regardless of + the precision of the parameters (default: True) + """ + + default_args = dict(lr=lr, betas=betas, weight_decay=weight_decay) + super(DeepSpeedCPULion, self).__init__(model_params, default_args) + + cpu_info = get_cpu_info() + self.cpu_vendor = cpu_info["vendor_id_raw"].lower() if "vendor_id_raw" in cpu_info else "unknown" + if "amd" in self.cpu_vendor: + for group_id, group in enumerate(self.param_groups): + for param_id, p in enumerate(group['params']): + if p.dtype == torch.half: + logger.warning("FP16 params for CPULion may not work on AMD CPUs") + break + else: + continue + break + + self.opt_id = DeepSpeedCPULion.optimizer_id + DeepSpeedCPULion.optimizer_id = DeepSpeedCPULion.optimizer_id + 1 + self.fp32_optimizer_states = fp32_optimizer_states + self.ds_opt_lion = CPULionBuilder().load() + + self.ds_opt_lion.create_lion(self.opt_id, lr, betas[0], betas[1], weight_decay, should_log_le("info")) + + def __del__(self): + # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize + # is used multiple times in the same process (notebook or pytest worker) + self.ds_opt_lion.destroy_lion(self.opt_id) + + def __setstate__(self, state): + super(DeepSpeedCPULion, self).__setstate__(state) + for group in self.param_groups: + group.setdefault('amsgrad', False) + + @torch.no_grad() + def step(self, closure=None): + """Update the model parameters. + + .. note:: + This method will be called internally by ZeRO-Offload. DeepSpeed + users should still use ``engine.step()`` as shown in the + `Getting Started + `_ guide. + + Args: + closure (callable, optional): closure to compute the loss. + Defaults to ``None``. + + Returns: + loss: if ``closure`` is provided. Otherwise ``None``. + """ + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # intended device for step + device = torch.device('cpu') + + for group_id, group in enumerate(self.param_groups): + for param_id, p in enumerate(group['params']): + + if p.grad is None: + continue + + assert p.device == device, f"CPULion param is on {p.device} and must be 'cpu', make " \ + "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config." + + state = self.state[p] + # State initialization + if len(state) == 0: + #print(f'group {group_id} param {param_id} = {p.numel()}') + state['step'] = 0 + + #use full precision by default unless self.fp32_optimizer_states is off + state_dtype = torch.float if self.fp32_optimizer_states else p.dtype + + # gradient momentums + state['exp_avg'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) + #memory_format=torch.preserve_format) + # gradient variances + state['exp_avg_sq'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) + #memory_format=torch.preserve_format) + + state['step'] += 1 + beta1, beta2 = group['betas'] + + self.ds_opt_lion.lion_update(self.opt_id, state['step'], group['lr'], beta1, beta2, + group['weight_decay'], p.data, p.grad.data, state['exp_avg']) + return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/fused_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/fused_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..7332a7f96361a1a05d770d945a90efe6b24ef217 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/fused_lion.py @@ -0,0 +1,131 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +This file is modified from fused_adam.py +""" + +import torch +from .multi_tensor_apply import MultiTensorApply + +multi_tensor_applier = MultiTensorApply(2048 * 32) +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import FusedLionBuilder + + +class FusedLion(torch.optim.Optimizer): + """Implements Lion algorithm. + + Currently GPU-only. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + set_grad_none (bool, optional): whether set grad to None when zero_grad() + method is called. (default: True) + + .. _Symbolic Discovery of Optimization Algorithms: + https://doi.org/10.48550/arXiv.2302.06675 + """ + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), weight_decay=0., set_grad_none=True): + + defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) + super(FusedLion, self).__init__(params, defaults) + self.set_grad_none = set_grad_none + + fused_lion_cuda = FusedLionBuilder().load() + # Skip buffer + self._dummy_overflow_buf = get_accelerator().IntTensor([0]) + self.multi_tensor_lion = fused_lion_cuda.multi_tensor_lion + + def zero_grad(self): + if self.set_grad_none: + for group in self.param_groups: + for p in group['params']: + p.grad = None + else: + super(FusedLion, self).zero_grad() + + def step(self, closure=None, grads=None, output_params=None, scale=None, grad_norms=None, grad_scaler=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + + The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes. + """ + if any(p is not None for p in [grads, output_params, scale, grad_norms]): + raise RuntimeError('FusedLion has been updated.') + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + if len(group['params']) == 0: + continue + beta1, beta2 = group['betas'] + + # assume same step across group now to simplify things + # per parameter step can be easily support by making it tensor, or pass list into kernel + if 'step' not in group: + group['step'] = 0 + + # create lists for multi-tensor apply + g_16, p_16, m_16 = [], [], [] + g_bf, p_bf, m_bf = [], [], [] + g_32, p_32, m_32 = [], [], [] + + for p in group['params']: + if p.grad is None: + continue + if p.grad.data.is_sparse: + raise NotImplementedError('FusedLion does not support sparse gradients') + + state = self.state[p] + # State initialization + if len(state) == 0: + # DeepSpeed ZeRO 3 processes each subgroup a time, so we need to keep tracking step count for each tensor separately. + # While this is not an issue for ZeRO 1 & 2, since they apply a single optimization step to the whole param group at the same time. + # In order to keep backward compatibility for the existing checkpoints, we use group['state'] to initialize state['step'] if it exists. + state['step'] = group.get('step', 0) + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + + if p.dtype == torch.float16: + g_16.append(p.grad.data) + p_16.append(p.data) + m_16.append(state['exp_avg']) + elif p.dtype == torch.bfloat16: + g_bf.append(p.grad) + p_bf.append(p) + m_bf.append(state['exp_avg']) + elif p.dtype == torch.float32: + g_32.append(p.grad.data) + p_32.append(p.data) + m_32.append(state['exp_avg']) + else: + raise RuntimeError('FusedLion only support fp16, bf16 and fp32.') + + if len(g_16) > 0: + state['step'] += 1 + multi_tensor_applier(self.multi_tensor_lion, self._dummy_overflow_buf, [g_16, p_16, m_16], group['lr'], + beta1, beta2, state['step'], group['weight_decay']) + + if len(g_bf) > 0: + state['step'] += 1 + multi_tensor_applier(self.multi_tensor_lion, self._dummy_overflow_buf, [g_bf, p_bf, m_bf], group['lr'], + beta1, beta2, state['step'], group['weight_decay']) + + if len(g_32) > 0: + state['step'] += 1 + multi_tensor_applier(self.multi_tensor_lion, self._dummy_overflow_buf, [g_32, p_32, m_32], group['lr'], + beta1, beta2, state['step'], group['weight_decay']) + + return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/multi_tensor_apply.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/multi_tensor_apply.py new file mode 100644 index 0000000000000000000000000000000000000000..0ba228505cef747eea4fec62f3e68707fa4daa0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/multi_tensor_apply.py @@ -0,0 +1,17 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from NVIDIA/apex, commit a109f85 +""" + + +class MultiTensorApply(object): + + def __init__(self, chunk_size): + self.chunk_size = chunk_size + + def __call__(self, op, noop_flag_buffer, tensor_lists, *args): + return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..afe48159933c75045e9445a577d545b7ae260073 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__init__.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import os +import pkgutil +import importlib + +from .builder import get_default_compute_capabilities, OpBuilder + +# Do not remove, required for abstract accelerator to detect if we have a deepspeed or 3p op_builder +__deepspeed__ = True + +# List of all available op builders from deepspeed op_builder +try: + import deepspeed.ops.op_builder # noqa: F401 # type: ignore + op_builder_dir = "deepspeed.ops.op_builder" +except ImportError: + op_builder_dir = "op_builder" + +__op_builders__ = [] + +this_module = sys.modules[__name__] + + +def builder_closure(member_name): + if op_builder_dir == "op_builder": + # during installation time cannot get builder due to torch not installed, + # return closure instead + def _builder(): + from deepspeed.accelerator import get_accelerator + builder = get_accelerator().create_op_builder(member_name) + return builder + + return _builder + else: + # during runtime, return op builder class directly + from deepspeed.accelerator import get_accelerator + builder = get_accelerator().get_op_builder(member_name) + return builder + + +# reflect builder names and add builder closure, such as 'TransformerBuilder()' creates op builder wrt current accelerator +for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(this_module.__file__)]): + if module_name != 'all_ops' and module_name != 'builder': + module = importlib.import_module(f".{module_name}", package=op_builder_dir) + for member_name in module.__dir__(): + if member_name.endswith('Builder') and member_name != "OpBuilder" and member_name != "CUDAOpBuilder": + # assign builder name to variable with same name + # the following is equivalent to i.e. TransformerBuilder = "TransformerBuilder" + this_module.__dict__[member_name] = builder_closure(member_name) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31b16fcdde12b70ca256c9e1f56b0eed8de2c20e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e2ea343b24aaf34644b72419379c9b8414a3a18 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4b378bdce86424372bfdca95a044880f47fbe38 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c42dd6b1a31c5935651f779f3bf1793480084a2b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..179c01c29b836af966e6660a6d204e1460d324e7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41ba016eccad83a5f43c5ab0799aead24863562a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed07f53a5f3ac67bfb052a7bb098075a02f75a82 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..02e6d63f4c80c8b90e3977391e693a2dfb3047df Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c937d3760c41a85a6ac6623241a16214e743ba7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a1e7cbd9b2b9ae57698bb559f3a20e350ff15ed Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d38644f7b9abb5a2c4aa0140d6d552ab06d4c4e0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4060487e61f2364d39d7d1c55a43fa6b4e55b7e1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..163133ab38feca34de1f7f2bc9c401dc44e743c9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14b3c1a68427f1ff28ad6164beec96e1ecffb178 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e154874ef49c7189e04b8b29725276042ec8400e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79056246654e7e796876e47883c15ba1319052fe Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b76fe296e68a7e7725be39e9f3dac9f44f0e9301 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37e310c96336647289aa3f26a01cc628814ebdb4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb642b4d541fabef240bb807a82d67ad0ff3299e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d193d3ffda639f93a5fc2e9564c52733d66c42f4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fb0a063861b02d89496018168d1c38f3492158a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb695b224f3113b9a6ffc1745084e1cb095b98bf Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59b8e29b17d49867fe2a0db0269e36779f029ec4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b3a30fd82c041fa79668543032be1256f31ecfa Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..819193b74ad8b078f2ac03abe77bddcfe9e802ac Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f667907ddc531df5e04ebeb874f09211d5952c99 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/all_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/all_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ff11ca180072f7dc72918b417b5e6e44eb53e3c4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/all_ops.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import pkgutil +import importlib +try: + # during installation time accelerator is visible, otherwise return deepspeed.accelerator + from accelerator import get_accelerator +except ImportError: + from deepspeed.accelerator import get_accelerator + +# List of all available ops + +# reflect all builder names into __op_builders__ +op_builder_dir = get_accelerator().op_builder_dir() +op_builder_module = importlib.import_module(op_builder_dir) +__op_builders__ = [] + +for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]): + # avoid self references + if module_name != 'all_ops' and module_name != 'builder': + module = importlib.import_module("{}.{}".format(op_builder_dir, module_name)) + for member_name in module.__dir__(): + if member_name.endswith('Builder'): + # append builder to __op_builders__ list + builder = get_accelerator().create_op_builder(member_name) + __op_builders__.append(builder) + +ALL_OPS = {op.name: op for op in __op_builders__ if op is not None} +accelerator_name = get_accelerator()._name diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..f59cc6810c6f04a9c9b323093b9bfa202294f9ee --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/async_io.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import shutil +import subprocess + +from .builder import TorchCPUOpBuilder + + +class AsyncIOBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def lib_sources(self): + src_list = [ + 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', + 'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp', + 'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', + 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' + ] + return src_list + + def sources(self): + return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp'] + + def include_paths(self): + import torch + if self.build_for_cpu: + CUDA_INCLUDE = [] + elif not self.is_rocm_pytorch(): + CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] + else: + CUDA_INCLUDE = [ + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"), + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "rocrand"), + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "hiprand"), + ] + return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE + + def cxx_args(self): + # -O0 for improved debugging, since performance is bound by I/O + args = super().cxx_args() + import torch + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): + args.remove('-std=c++17') + args.append('-std=c++14') + args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder'] + return args + + def extra_ldflags(self): + if self.build_for_cpu: + return ['-fopenmp'] + + import torch.utils.cpp_extension + CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME + if CUDA_HOME is None: + ldflags = ['-laio'] # the ROCM case + else: + CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") + ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart'] + return ldflags + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_submit', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..aef7d50358ddfed809c1a7744db4a37f16037e8c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py @@ -0,0 +1,861 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import re +import sys +import time +import importlib +from pathlib import Path +import subprocess +import shlex +import shutil +import tempfile +import distutils.ccompiler +import distutils.log +import distutils.sysconfig +from distutils.errors import CompileError, LinkError +from abc import ABC, abstractmethod +from typing import List + +YELLOW = '\033[93m' +END = '\033[0m' +WARNING = f"{YELLOW} [WARNING] {END}" + +DEFAULT_TORCH_EXTENSION_PATH = "/tmp/torch_extensions" +DEFAULT_COMPUTE_CAPABILITIES = "6.0;6.1;7.0" + +try: + import torch +except ImportError: + print(f"{WARNING} unable to import torch, please install it if you want to pre-compile any deepspeed ops.") +else: + TORCH_MAJOR = int(torch.__version__.split('.')[0]) + TORCH_MINOR = int(torch.__version__.split('.')[1]) + + +class MissingCUDAException(Exception): + pass + + +class CUDAMismatchException(Exception): + pass + + +def installed_cuda_version(name=""): + import torch.utils.cpp_extension + cuda_home = torch.utils.cpp_extension.CUDA_HOME + if cuda_home is None: + raise MissingCUDAException("CUDA_HOME does not exist, unable to compile CUDA op(s)") + # Ensure there is not a cuda version mismatch between torch and nvcc compiler + output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True) + output_split = output.split() + release_idx = output_split.index("release") + release = output_split[release_idx + 1].replace(',', '').split(".") + # Ignore patch versions, only look at major + minor + cuda_major, cuda_minor = release[:2] + return int(cuda_major), int(cuda_minor) + + +def get_default_compute_capabilities(): + compute_caps = DEFAULT_COMPUTE_CAPABILITIES + # Update compute capability according to: https://en.wikipedia.org/wiki/CUDA#GPUs_supported + import torch.utils.cpp_extension + if torch.utils.cpp_extension.CUDA_HOME is not None: + if installed_cuda_version()[0] == 11: + if installed_cuda_version()[1] >= 0: + compute_caps += ";8.0" + if installed_cuda_version()[1] >= 1: + compute_caps += ";8.6" + if installed_cuda_version()[1] >= 8: + compute_caps += ";9.0" + elif installed_cuda_version()[0] == 12: + compute_caps += ";8.0;8.6;9.0" + if installed_cuda_version()[1] >= 8: + compute_caps += ";10.0;12.0" + return compute_caps + + +# list compatible minor CUDA versions - so that for example pytorch built with cuda-11.0 can be used +# to build deepspeed and system-wide installed cuda 11.2 +cuda_minor_mismatch_ok = { + 10: ["10.0", "10.1", "10.2"], + 11: ["11.0", "11.1", "11.2", "11.3", "11.4", "11.5", "11.6", "11.7", "11.8"], + 12: ["12.0", "12.1", "12.2", "12.3", "12.4", "12.5", "12.6", "12.8", "12.9"], # There is no CUDATk 12.7 +} + + +def assert_no_cuda_mismatch(name=""): + cuda_major, cuda_minor = installed_cuda_version(name) + sys_cuda_version = f'{cuda_major}.{cuda_minor}' + torch_cuda_version = ".".join(torch.version.cuda.split('.')[:2]) + # This is a show-stopping error, should probably not proceed past this + if sys_cuda_version != torch_cuda_version: + if (cuda_major in cuda_minor_mismatch_ok and sys_cuda_version in cuda_minor_mismatch_ok[cuda_major] + and torch_cuda_version in cuda_minor_mismatch_ok[cuda_major]): + print(f"Installed CUDA version {sys_cuda_version} does not match the " + f"version torch was compiled with {torch.version.cuda} " + "but since the APIs are compatible, accepting this combination") + return True + elif os.getenv("DS_SKIP_CUDA_CHECK", "0") == "1": + print( + f"{WARNING} DeepSpeed Op Builder: Installed CUDA version {sys_cuda_version} does not match the " + f"version torch was compiled with {torch.version.cuda}." + "Detected `DS_SKIP_CUDA_CHECK=1`: Allowing this combination of CUDA, but it may result in unexpected behavior." + ) + return True + raise CUDAMismatchException( + f">- DeepSpeed Op Builder: Installed CUDA version {sys_cuda_version} does not match the " + f"version torch was compiled with {torch.version.cuda}, unable to compile " + "cuda/cpp extensions without a matching cuda version.") + return True + + +class OpBuilder(ABC): + _rocm_version = None + _rocm_gpu_arch = None + _rocm_wavefront_size = None + _is_rocm_pytorch = None + _is_sycl_enabled = None + _loaded_ops = {} + + def __init__(self, name): + self.name = name + self.jit_mode = False + self.build_for_cpu = False + self.enable_bf16 = False + self.error_log = None + + @abstractmethod + def absolute_name(self): + ''' + Returns absolute build path for cases where the op is pre-installed, e.g., deepspeed.ops.adam.cpu_adam + will be installed as something like: deepspeed/ops/adam/cpu_adam.so + ''' + pass + + @abstractmethod + def sources(self): + ''' + Returns list of source files for your op, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed) + ''' + pass + + def hipify_extension(self): + pass + + def sycl_extension(self): + pass + + @staticmethod + def validate_torch_version(torch_info): + install_torch_version = torch_info['version'] + current_torch_version = ".".join(torch.__version__.split('.')[:2]) + if install_torch_version != current_torch_version: + raise RuntimeError("PyTorch version mismatch! DeepSpeed ops were compiled and installed " + "with a different version than what is being used at runtime. " + f"Please re-install DeepSpeed or switch torch versions. " + f"Install torch version={install_torch_version}, " + f"Runtime torch version={current_torch_version}") + + @staticmethod + def validate_torch_op_version(torch_info): + if not OpBuilder.is_rocm_pytorch(): + current_cuda_version = ".".join(torch.version.cuda.split('.')[:2]) + install_cuda_version = torch_info['cuda_version'] + if install_cuda_version != current_cuda_version: + raise RuntimeError("CUDA version mismatch! DeepSpeed ops were compiled and installed " + "with a different version than what is being used at runtime. " + f"Please re-install DeepSpeed or switch torch versions. " + f"Install CUDA version={install_cuda_version}, " + f"Runtime CUDA version={current_cuda_version}") + else: + current_hip_version = ".".join(torch.version.hip.split('.')[:2]) + install_hip_version = torch_info['hip_version'] + if install_hip_version != current_hip_version: + raise RuntimeError("HIP version mismatch! DeepSpeed ops were compiled and installed " + "with a different version than what is being used at runtime. " + f"Please re-install DeepSpeed or switch torch versions. " + f"Install HIP version={install_hip_version}, " + f"Runtime HIP version={current_hip_version}") + + @staticmethod + def is_rocm_pytorch(): + if OpBuilder._is_rocm_pytorch is not None: + return OpBuilder._is_rocm_pytorch + + _is_rocm_pytorch = False + try: + import torch + except ImportError: + pass + else: + if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5): + _is_rocm_pytorch = hasattr(torch.version, 'hip') and torch.version.hip is not None + if _is_rocm_pytorch: + from torch.utils.cpp_extension import ROCM_HOME + _is_rocm_pytorch = ROCM_HOME is not None + OpBuilder._is_rocm_pytorch = _is_rocm_pytorch + return OpBuilder._is_rocm_pytorch + + @staticmethod + def is_sycl_enabled(): + if OpBuilder._is_sycl_enabled is not None: + return OpBuilder._is_sycl_enabled + + _is_sycl_enabled = False + try: + result = subprocess.run(["c2s", "--version"], capture_output=True) + except: + pass + else: + _is_sycl_enabled = True + + OpBuilder._is_sycl_enabled = _is_sycl_enabled + return OpBuilder._is_sycl_enabled + + @staticmethod + def installed_rocm_version(): + if OpBuilder._rocm_version: + return OpBuilder._rocm_version + + ROCM_MAJOR = '0' + ROCM_MINOR = '0' + ROCM_VERSION_DEV_RAW = "" + if OpBuilder.is_rocm_pytorch(): + from torch.utils.cpp_extension import ROCM_HOME + rocm_ver_file = Path(ROCM_HOME).joinpath(".info/version") + if rocm_ver_file.is_file(): + with open(rocm_ver_file, 'r') as file: + ROCM_VERSION_DEV_RAW = file.read() + elif "rocm" in torch.__version__: + ROCM_VERSION_DEV_RAW = torch.__version__.split("rocm")[1] + if ROCM_VERSION_DEV_RAW != "": + ROCM_MAJOR = ROCM_VERSION_DEV_RAW.split('.')[0] + ROCM_MINOR = ROCM_VERSION_DEV_RAW.split('.')[1] + else: + # Look in /usr/include/rocm-version.h + rocm_ver_file = Path("/usr/include/rocm_version.h") + if rocm_ver_file.is_file(): + with open(rocm_ver_file, 'r') as file: + for ln in file.readlines(): + if "#define ROCM_VERSION_MAJOR" in ln: + ROCM_MAJOR = re.findall(r'\S+', ln)[2] + elif "#define ROCM_VERSION_MINOR" in ln: + ROCM_MINOR = re.findall(r'\S+', ln)[2] + if ROCM_MAJOR == '0': + assert False, "Could not detect ROCm version" + + OpBuilder._rocm_version = (int(ROCM_MAJOR), int(ROCM_MINOR)) + return OpBuilder._rocm_version + + @staticmethod + def get_rocm_gpu_arch(): + if OpBuilder._rocm_gpu_arch: + return OpBuilder._rocm_gpu_arch + rocm_info = Path("/opt/rocm/bin/rocminfo") + if (not rocm_info.is_file()): + rocm_info = Path("rocminfo") + rocm_gpu_arch_cmd = str(rocm_info) + " | grep -o -m 1 'gfx.*'" + try: + result = subprocess.check_output(rocm_gpu_arch_cmd, shell=True) + rocm_gpu_arch = result.decode('utf-8').strip() + except subprocess.CalledProcessError: + rocm_gpu_arch = "" + OpBuilder._rocm_gpu_arch = rocm_gpu_arch + return OpBuilder._rocm_gpu_arch + + @staticmethod + def get_rocm_wavefront_size(): + if OpBuilder._rocm_wavefront_size: + return OpBuilder._rocm_wavefront_size + + rocm_info = Path("/opt/rocm/bin/rocminfo") + if (not rocm_info.is_file()): + rocm_info = Path("rocminfo") + rocm_wavefront_size_cmd = str( + rocm_info) + " | grep -Eo -m1 'Wavefront Size:[[:space:]]+[0-9]+' | grep -Eo '[0-9]+'" + try: + result = subprocess.check_output(rocm_wavefront_size_cmd, shell=True) + rocm_wavefront_size = result.decode('utf-8').strip() + except subprocess.CalledProcessError: + rocm_wavefront_size = "32" + OpBuilder._rocm_wavefront_size = rocm_wavefront_size + return OpBuilder._rocm_wavefront_size + + def include_paths(self): + ''' + Returns list of include paths, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed) + ''' + return [] + + def nvcc_args(self): + ''' + Returns optional list of compiler flags to forward to nvcc when building CUDA sources + ''' + return [] + + def cxx_args(self): + ''' + Returns optional list of compiler flags to forward to the build + ''' + return [] + + def is_compatible(self, verbose=False): + ''' + Check if all non-python dependencies are satisfied to build this op + ''' + return True + + def extra_ldflags(self): + return [] + + def has_function(self, funcname, libraries, library_dirs=None, verbose=False): + ''' + Test for existence of a function within a tuple of libraries. + + This is used as a smoke test to check whether a certain library is available. + As a test, this creates a simple C program that calls the specified function, + and then distutils is used to compile that program and link it with the specified libraries. + Returns True if both the compile and link are successful, False otherwise. + ''' + tempdir = None # we create a temporary directory to hold various files + filestderr = None # handle to open file to which we redirect stderr + oldstderr = None # file descriptor for stderr + try: + # Echo compile and link commands that are used. + if verbose: + distutils.log.set_verbosity(1) + + # Create a compiler object. + compiler = distutils.ccompiler.new_compiler(verbose=verbose) + + # Configure compiler and linker to build according to Python install. + distutils.sysconfig.customize_compiler(compiler) + + # Create a temporary directory to hold test files. + tempdir = tempfile.mkdtemp() + + # Define a simple C program that calls the function in question + prog = "void %s(void); int main(int argc, char** argv) { %s(); return 0; }" % (funcname, funcname) + + # Write the test program to a file. + filename = os.path.join(tempdir, 'test.c') + with open(filename, 'w') as f: + f.write(prog) + + # Redirect stderr file descriptor to a file to silence compile/link warnings. + if not verbose: + filestderr = open(os.path.join(tempdir, 'stderr.txt'), 'w') + oldstderr = os.dup(sys.stderr.fileno()) + os.dup2(filestderr.fileno(), sys.stderr.fileno()) + + # Workaround for behavior in distutils.ccompiler.CCompiler.object_filenames() + # Otherwise, a local directory will be used instead of tempdir + drive, driveless_filename = os.path.splitdrive(filename) + root_dir = driveless_filename[0] if os.path.isabs(driveless_filename) else '' + output_dir = os.path.join(drive, root_dir) + + # Attempt to compile the C program into an object file. + cflags = shlex.split(os.environ.get('CFLAGS', "")) + objs = compiler.compile([filename], output_dir=output_dir, extra_preargs=self.strip_empty_entries(cflags)) + + # Attempt to link the object file into an executable. + # Be sure to tack on any libraries that have been specified. + ldflags = shlex.split(os.environ.get('LDFLAGS', "")) + compiler.link_executable(objs, + os.path.join(tempdir, 'a.out'), + extra_preargs=self.strip_empty_entries(ldflags), + libraries=libraries, + library_dirs=library_dirs) + + # Compile and link succeeded + return True + + except CompileError: + return False + + except LinkError: + return False + + except: + return False + + finally: + # Restore stderr file descriptor and close the stderr redirect file. + if oldstderr is not None: + os.dup2(oldstderr, sys.stderr.fileno()) + if filestderr is not None: + filestderr.close() + + # Delete the temporary directory holding the test program and stderr files. + if tempdir is not None: + shutil.rmtree(tempdir) + + def strip_empty_entries(self, args): + ''' + Drop any empty strings from the list of compile and link flags + ''' + return [x for x in args if len(x) > 0] + + def cpu_arch(self): + try: + from cpuinfo import get_cpu_info + except ImportError as e: + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return "-march=native" + + try: + cpu_info = get_cpu_info() + except Exception as e: + self.warning(f"{self.name} attempted to use py-cpuinfo but failed (exception type: {type(e)}, {e}), " + "falling back to lscpu to get this information.") + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return "-march=native" + + if cpu_info['arch'].startswith('PPC_'): + # gcc does not provide -march on PowerPC, use -mcpu instead + return '-mcpu=native' + return '-march=native' + + def get_cuda_compile_flag(self): + try: + if not self.is_rocm_pytorch(): + assert_no_cuda_mismatch(self.name) + return "-D__ENABLE_CUDA__" + except MissingCUDAException: + print(f"{WARNING} {self.name} cuda is missing or is incompatible with installed torch, " + "only cpu ops can be compiled!") + return '-D__DISABLE_CUDA__' + return '-D__DISABLE_CUDA__' + + def _backup_cpuinfo(self): + # Construct cpu_info dict from lscpu that is similar to what py-cpuinfo provides + if not self.command_exists('lscpu'): + self.warning(f"{self.name} attempted to query 'lscpu' after failing to use py-cpuinfo " + "to detect the CPU architecture. 'lscpu' does not appear to exist on " + "your system, will fall back to use -march=native and non-vectorized execution.") + return None + result = subprocess.check_output(['lscpu']) + result = result.decode('utf-8').strip().lower() + + cpu_info = {} + cpu_info['arch'] = None + cpu_info['flags'] = "" + if 'genuineintel' in result or 'authenticamd' in result: + cpu_info['arch'] = 'X86_64' + if 'avx512' in result: + cpu_info['flags'] += 'avx512,' + elif 'avx512f' in result: + cpu_info['flags'] += 'avx512f,' + if 'avx2' in result: + cpu_info['flags'] += 'avx2' + elif 'ppc64le' in result: + cpu_info['arch'] = "PPC_" + + return cpu_info + + def simd_width(self): + try: + from cpuinfo import get_cpu_info + except ImportError as e: + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return '-D__SCALAR__' + + try: + cpu_info = get_cpu_info() + except Exception as e: + self.warning(f"{self.name} attempted to use py-cpuinfo but failed (exception type: {type(e)}, {e}), " + "falling back to lscpu to get this information.") + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return '-D__SCALAR__' + + if cpu_info['arch'] == 'X86_64': + if 'avx512' in cpu_info['flags'] or 'avx512f' in cpu_info['flags']: + return '-D__AVX512__' + elif 'avx2' in cpu_info['flags']: + return '-D__AVX256__' + return '-D__SCALAR__' + + def command_exists(self, cmd): + if '|' in cmd: + cmds = cmd.split("|") + else: + cmds = [cmd] + valid = False + for cmd in cmds: + safe_cmd = ["bash", "-c", f"type {cmd}"] + result = subprocess.Popen(safe_cmd, stdout=subprocess.PIPE) + valid = valid or result.wait() == 0 + + if not valid and len(cmds) > 1: + print(f"{WARNING} {self.name} requires one of the following commands '{cmds}', but it does not exist!") + elif not valid and len(cmds) == 1: + print(f"{WARNING} {self.name} requires the '{cmd}' command, but it does not exist!") + return valid + + def warning(self, msg): + self.error_log = f"{msg}" + print(f"{WARNING} {msg}") + + def deepspeed_src_path(self, code_path): + if os.path.isabs(code_path): + return code_path + else: + return os.path.join(Path(__file__).parent.parent.absolute(), code_path) + + def builder(self): + from torch.utils.cpp_extension import CppExtension + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + return CppExtension(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + extra_compile_args={'cxx': self.strip_empty_entries(self.cxx_args())}, + extra_link_args=self.strip_empty_entries(self.extra_ldflags())) + + def load(self, verbose=True): + if self.name in __class__._loaded_ops: + return __class__._loaded_ops[self.name] + + from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name + from deepspeed.accelerator import get_accelerator + if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: + # Ensure the op we're about to load was compiled with the same + # torch/cuda versions we are currently using at runtime. + self.validate_torch_version(torch_info) + if torch.cuda.is_available() and isinstance(self, CUDAOpBuilder): + self.validate_torch_op_version(torch_info) + + op_module = importlib.import_module(self.absolute_name()) + __class__._loaded_ops[self.name] = op_module + return op_module + else: + return self.jit_load(verbose) + + def jit_load(self, verbose=True): + if not self.is_compatible(verbose): + raise RuntimeError( + f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" + ) + try: + import ninja # noqa: F401 # type: ignore + except ImportError: + raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") + + if isinstance(self, CUDAOpBuilder) and not self.is_rocm_pytorch(): + self.build_for_cpu = not torch.cuda.is_available() + + self.jit_mode = True + from torch.utils.cpp_extension import load + + start_build = time.time() + sources = [os.path.abspath(self.deepspeed_src_path(path)) for path in self.sources()] + extra_include_paths = [os.path.abspath(self.deepspeed_src_path(path)) for path in self.include_paths()] + + # Torch will try and apply whatever CCs are in the arch list at compile time, + # we have already set the intended targets ourselves we know that will be + # needed at runtime. This prevents CC collisions such as multiple __half + # implementations. Stash arch list to reset after build. + torch_arch_list = None + if "TORCH_CUDA_ARCH_LIST" in os.environ: + torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") + os.environ["TORCH_CUDA_ARCH_LIST"] = "" + + nvcc_args = self.strip_empty_entries(self.nvcc_args()) + cxx_args = self.strip_empty_entries(self.cxx_args()) + + cxx_args.append("-UC10_USE_GLOG") + nvcc_args.append("-UC10_USE_GLOG") + if isinstance(self, CUDAOpBuilder): + if not self.build_for_cpu and self.enable_bf16: + cxx_args.append("-DBF16_AVAILABLE") + nvcc_args.append("-DBF16_AVAILABLE") + nvcc_args.append("-U__CUDA_NO_BFLOAT16_OPERATORS__") + nvcc_args.append("-U__CUDA_NO_BFLOAT162_OPERATORS__") + nvcc_args.append("-U__CUDA_NO_BFLOAT16_CONVERSIONS__") + + if self.is_rocm_pytorch(): + cxx_args.append("-D__HIP_PLATFORM_AMD__=1") + os.environ["PYTORCH_ROCM_ARCH"] = self.get_rocm_gpu_arch() + cxx_args.append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) + + op_module = load(name=self.name, + sources=self.strip_empty_entries(sources), + extra_include_paths=self.strip_empty_entries(extra_include_paths), + extra_cflags=cxx_args, + extra_cuda_cflags=nvcc_args, + extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), + with_cuda=True if (isinstance(self, CUDAOpBuilder) and not self.build_for_cpu) else None, + verbose=verbose) + + build_duration = time.time() - start_build + if verbose: + print(f"Time to load {self.name} op: {build_duration} seconds") + + # Reset arch list so we are not silently removing it for other possible use cases + if torch_arch_list: + os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list + + __class__._loaded_ops[self.name] = op_module + + return op_module + + +class CUDAOpBuilder(OpBuilder): + + def compute_capability_args(self, cross_compile_archs=None): + """ + Returns nvcc compute capability compile flags. + + 1. `TORCH_CUDA_ARCH_LIST` takes priority over `cross_compile_archs`. + 2. If neither is set default compute capabilities will be used + 3. Under `jit_mode` compute capabilities of all visible cards will be used plus PTX + + Format: + + - `TORCH_CUDA_ARCH_LIST` may use ; or whitespace separators. Examples: + + TORCH_CUDA_ARCH_LIST="6.1;7.5;8.6;9.0;10.0" pip install ... + TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 9.0 10.0+PTX" pip install ... + + - `cross_compile_archs` uses ; separator. + + """ + ccs = [] + if self.jit_mode: + # Compile for underlying architectures since we know those at runtime + for i in range(torch.cuda.device_count()): + CC_MAJOR, CC_MINOR = torch.cuda.get_device_capability(i) + cc = f"{CC_MAJOR}.{CC_MINOR}" + if cc not in ccs: + ccs.append(cc) + ccs = sorted(ccs) + ccs[-1] += '+PTX' + else: + # Cross-compile mode, compile for various architectures + # env override takes priority + cross_compile_archs_env = os.environ.get('TORCH_CUDA_ARCH_LIST', None) + if cross_compile_archs_env is not None: + if cross_compile_archs is not None: + print( + f"{WARNING} env var TORCH_CUDA_ARCH_LIST={cross_compile_archs_env} overrides cross_compile_archs={cross_compile_archs}" + ) + cross_compile_archs = cross_compile_archs_env.replace(' ', ';') + else: + if cross_compile_archs is None: + cross_compile_archs = get_default_compute_capabilities() + ccs = cross_compile_archs.split(';') + + ccs = self.filter_ccs(ccs) + if len(ccs) == 0: + raise RuntimeError( + f"Unable to load {self.name} op due to no compute capabilities remaining after filtering") + + args = [] + self.enable_bf16 = True + for cc in ccs: + num = cc[0] + cc[1].split('+')[0] + args.append(f'-gencode=arch=compute_{num},code=sm_{num}') + if cc[1].endswith('+PTX'): + args.append(f'-gencode=arch=compute_{num},code=compute_{num}') + + if int(cc[0]) <= 7: + self.enable_bf16 = False + + return args + + def filter_ccs(self, ccs: List[str]): + """ + Prune any compute capabilities that are not compatible with the builder. Should log + which CCs have been pruned. + """ + return [cc.split('.') for cc in ccs] + + def version_dependent_macros(self): + # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456 + version_ge_1_1 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0): + version_ge_1_1 = ['-DVERSION_GE_1_1'] + version_ge_1_3 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2): + version_ge_1_3 = ['-DVERSION_GE_1_3'] + version_ge_1_5 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): + version_ge_1_5 = ['-DVERSION_GE_1_5'] + return version_ge_1_1 + version_ge_1_3 + version_ge_1_5 + + def is_compatible(self, verbose=False): + return super().is_compatible(verbose) + + def builder(self): + try: + if not self.is_rocm_pytorch(): + assert_no_cuda_mismatch(self.name) + self.build_for_cpu = False + except MissingCUDAException: + self.build_for_cpu = True + + if self.build_for_cpu: + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + else: + from torch.utils.cpp_extension import CUDAExtension as ExtensionBuilder + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} if self.build_for_cpu else \ + {'cxx': self.strip_empty_entries(self.cxx_args()), \ + 'nvcc': self.strip_empty_entries(self.nvcc_args())} + + if not self.build_for_cpu and self.enable_bf16: + compile_args['cxx'].append("-DBF16_AVAILABLE") + compile_args['nvcc'].append("-DBF16_AVAILABLE") + + if self.is_rocm_pytorch(): + compile_args['cxx'].append("-D__HIP_PLATFORM_AMD__=1") + #cxx compiler args are required to compile cpp files + compile_args['cxx'].append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) + #nvcc compiler args are required to compile hip files + compile_args['nvcc'].append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) + if self.get_rocm_gpu_arch(): + os.environ["PYTORCH_ROCM_ARCH"] = self.get_rocm_gpu_arch() + + cuda_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args, + extra_link_args=self.strip_empty_entries(self.extra_ldflags())) + + if self.is_rocm_pytorch(): + # hip converts paths to absolute, this converts back to relative + sources = cuda_ext.sources + curr_file = Path(__file__).parent.parent # ds root + for i in range(len(sources)): + src = Path(sources[i]) + if src.is_absolute(): + sources[i] = str(src.relative_to(curr_file)) + else: + sources[i] = str(src) + cuda_ext.sources = sources + return cuda_ext + + def hipify_extension(self): + if self.is_rocm_pytorch(): + from torch.utils.hipify import hipify_python + hipify_python.hipify( + project_directory=os.getcwd(), + output_directory=os.getcwd(), + header_include_dirs=self.include_paths(), + includes=[os.path.join(os.getcwd(), '*')], + extra_files=[os.path.abspath(s) for s in self.sources()], + show_detailed=True, + is_pytorch_extension=True, + hipify_extra_files_only=True, + ) + + def cxx_args(self): + if sys.platform == "win32": + return ['-O2'] + else: + return ['-O3', '-std=c++17', '-g', '-Wno-reorder'] + + def nvcc_args(self): + if self.build_for_cpu: + return [] + args = ['-O3'] + if self.is_rocm_pytorch(): + ROCM_MAJOR, ROCM_MINOR = self.installed_rocm_version() + args += [ + '-std=c++17', '-U__HIP_NO_HALF_OPERATORS__', '-U__HIP_NO_HALF_CONVERSIONS__', + '-U__HIP_NO_HALF2_OPERATORS__', + '-DROCM_VERSION_MAJOR=%s' % ROCM_MAJOR, + '-DROCM_VERSION_MINOR=%s' % ROCM_MINOR + ] + else: + try: + nvcc_threads = int(os.getenv("DS_NVCC_THREADS", "")) + if nvcc_threads <= 0: + raise ValueError("") + except ValueError: + nvcc_threads = min(os.cpu_count(), 8) + + cuda_major, cuda_minor = installed_cuda_version() + if cuda_major > 10: + if cuda_major == 12 and cuda_minor >= 5: + std_lib = '-std=c++20' + else: + std_lib = '-std=c++17' + else: + std_lib = '-std=c++14' + args += [ + '-allow-unsupported-compiler' if sys.platform == "win32" else '', '--use_fast_math', std_lib, + '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__', + f'--threads={nvcc_threads}' + ] + if os.environ.get('DS_DEBUG_CUDA_BUILD', '0') == '1': + args.append('--ptxas-options=-v') + args += self.compute_capability_args() + return args + + def libraries_args(self): + if self.build_for_cpu: + return [] + + if sys.platform == "win32": + return ['cublas', 'curand'] + else: + return [] + + +class TorchCPUOpBuilder(CUDAOpBuilder): + + def get_cuda_lib64_path(self): + import torch + if not self.is_rocm_pytorch(): + CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64") + if not os.path.exists(CUDA_LIB64): + CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib") + else: + CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.ROCM_HOME, "lib") + return CUDA_LIB64 + + def extra_ldflags(self): + if self.build_for_cpu: + return ['-fopenmp'] + + if not self.is_rocm_pytorch(): + ld_flags = ['-lcurand'] + if not self.build_for_cpu: + ld_flags.append(f'-L{self.get_cuda_lib64_path()}') + return ld_flags + + return [] + + def cxx_args(self): + args = [] + if not self.build_for_cpu: + CUDA_LIB64 = self.get_cuda_lib64_path() + + args += super().cxx_args() + args += [ + f'-L{CUDA_LIB64}', + '-lcudart', + '-lcublas', + '-g', + ] + + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + CUDA_ENABLE = self.get_cuda_compile_flag() + args += [ + CPU_ARCH, + '-fopenmp', + SIMD_WIDTH, + CUDA_ENABLE, + ] + + return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7084db8469f1dae201796c44f911f539459e4757 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .comm import CCLCommBuilder, ShareMemCommBuilder +from .fused_adam import FusedAdamBuilder +from .cpu_adam import CPUAdamBuilder +from .no_impl import NotImplementedBuilder +from .async_io import AsyncIOBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f3351925e39d041e72faa3dd8168f1d2ec174b5e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6489a0777eda56a12ae400415df66988830150a2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96b8fdf538eacd2f4d0264bb9204e5704dc0e087 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..10c284beb50a5390d1b665812525d90e684a1a64 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f4b8f66a497f2469e8a10cfe293a2751e474048 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..910a2f6beef0aafb3d9ee8e22e2dfee799d9ee45 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8d0955c878e510412665104f5a19f2fd34d04a5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..dcb9feabcfc3c7c3081d812fafb7d8816a031eb3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/async_io.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess + +from .builder import CPUOpBuilder + + +class AsyncIOBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def lib_sources(self): + src_list = [ + 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', + 'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp', + 'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', + 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' + ] + return src_list + + def sources(self): + return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp'] + + def include_paths(self): + return ['csrc/aio/py_lib', 'csrc/aio/common'] + + def cxx_args(self): + # -O0 for improved debugging, since performance is bound by I/O + args = super().cxx_args() + import torch + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): + args.remove('-std=c++17') + args.append('-std=c++14') + args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder'] + return args + + def extra_ldflags(self): + return ['-laio', '-fopenmp'] + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_submit', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..d881842ad0b18654bc26cb3b54d345a4cd160da4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/builder.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class CPUOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + args = ['-O3', '-g', '-Wno-reorder'] + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + args += [CPU_ARCH, '-fopenmp', SIMD_WIDTH] + return args + + def libraries_args(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/comm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..fec960b63b2ec16cb7ab6a92b509a5d97b901cef --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/comm.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from .builder import CPUOpBuilder + + +class CCLCommBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_CCL_COMM" + NAME = "deepspeed_ccl_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/ccl.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) + + 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" + ) + return [] + 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) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..0c8438aea40d4ffd588a496844d692c9a17faaa6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class CPUAdamBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..34b43825b09024136afdc44e916349ea3f5ce5ad --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class FusedAdamBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/cpu/adam/fused_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..69d114a9f1c0b0defc482c1fb143c261fc466125 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class NotImplementedBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + 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 load(self, verbose=True): + raise ValueError("This op had not been implemented on CPU backend.") + + def sources(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..c05f7148895000c10cdfebef7461858a1f987fbb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class CPUAdagradBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..7f4c0847a8c4f32b5f3cacd40f88eed8eea904f9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adam.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class CPUAdamBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..9a60d99773b31252c9d002a0cf3282185d91a07c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_lion.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class CPULionBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_LION" + NAME = "cpu_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/dc.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/dc.py new file mode 100644 index 0000000000000000000000000000000000000000..15b25bf3393e45d3ea2936cf7ca57fee7c49adeb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/dc.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class DeepCompileBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_DEEP_COMPILE" + NAME = "dc" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/compile/deepcompile.cpp', 'csrc/compile/init.cpp', 'csrc/compile/z1.cpp', 'csrc/compile/z2.cpp', + 'csrc/compile/z3.cpp', 'csrc/compile/util.cpp' + ] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + import os + import torch + if self.build_for_cpu: + CUDA_INCLUDE = [] + elif not self.is_rocm_pytorch(): + CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] + else: + CUDA_INCLUDE = [ + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"), + ] + return ['csrc/includes', 'csrc/compile'] + CUDA_INCLUDE diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/evoformer_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/evoformer_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..7f68ccf8729085028d445fb3d7f76d6f51556bfe --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/evoformer_attn.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder, installed_cuda_version +import os + + +class EvoformerAttnBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_EVOFORMER_ATTN" + NAME = "evoformer_attn" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + self.cutlass_path = os.environ.get('CUTLASS_PATH') + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def sources(self): + src_dir = 'csrc/deepspeed4science/evoformer_attn' + return [f'{src_dir}/attention.cpp', f'{src_dir}/attention_back.cu', f'{src_dir}/attention_cu.cu'] + + def nvcc_args(self): + args = super().nvcc_args() + try: + import torch + except ImportError: + self.warning("Please install torch if trying to pre-compile kernels") + return args + major = torch.cuda.get_device_properties(0).major #ignore-cuda + minor = torch.cuda.get_device_properties(0).minor #ignore-cuda + args.append(f"-DGPU_ARCH={major}{minor}") + return args + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile kernels") + return False + if self.cutlass_path is None: + if verbose: + self.warning("Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH") + return False + if os.path.exists(f'{self.cutlass_path}/CHANGELOG.md'): + with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f: + if '3.1.0' not in f.read(): + if verbose: + self.warning("Please use CUTLASS version >= 3.1.0") + return False + else: + # pip install nvidia-cutlass package + try: + import cutlass + except ImportError: + if verbose: + self.warning("Please pip install nvidia-cutlass if trying to pre-compile kernels") + return False + cutlass_major, cutlass_minor = cutlass.__version__.split('.')[:2] + cutlass_compatible = (int(cutlass_major) >= 3 and int(cutlass_minor) >= 1) + if not cutlass_compatible: + if verbose: + self.warning("Please use CUTLASS version >= 3.1.0") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 7: + if verbose: + self.warning("Please use a GPU with compute capability >= 7.0") + cuda_okay = False + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("Please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def include_paths(self): + includes = [f'{self.cutlass_path}/include', f'{self.cutlass_path}/tools/util/include'] + return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fp_quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fp_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..2b962ac2c1fea18d1024359623f3faf82228f253 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fp_quantizer.py @@ -0,0 +1,120 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + from packaging import version as pkg_version +except ImportError: + pkg_version = None + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class FPQuantizerBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_FP_QUANTIZER" + NAME = "fp_quantizer" + + 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.fp_quantizer.{self.NAME}_op' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 8: + if verbose: + self.warning("NVIDIA Inference is only supported on Ampere and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + + try: + import triton + except ImportError: + if verbose: + self.warning( + f"please install triton==2.3.0, 2.3.1 or 3.0.0 if you want to use the FP Quantizer Kernels") + return False + + # triton 2.3.{0,1} and 3.0.0 are ok. + allowed_versions = ("2.3", "3.0", "3.1", "3.2") + if pkg_version: + allowed = (pkg_version.parse(v) for v in allowed_versions) + installed_triton = pkg_version.parse(triton.__version__) + triton_mismatch = all(installed_triton.major != a.major or installed_triton.minor != a.minor + for a in allowed) + else: + installed_triton = triton.__version__ + major, minor, _ = installed_triton.split(".") + allowed = (v.split(".") for v in allowed_versions) + triton_mismatch = all(major != v[0] or minor != v[1] for v in allowed) + + if triton_mismatch: + if verbose: + self.warning( + f"FP Quantizer is using an untested triton version ({installed_triton}), only 2.3.{0,1} and 3.0.0 are known to be compatible with these kernels" + ) + return False + + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 8: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def sources(self): + return [ + "csrc/fp_quantizer/fp_quantize.cu", + "csrc/fp_quantizer/fp_quantize.cpp", + ] + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def include_paths(self): + return ['csrc/fp_quantizer/includes', 'csrc/includes'] + + @staticmethod + def get_default_quant_dtype(): + import torch + return torch.uint8 + + @staticmethod + def get_quant_range(q_bits=None): + if q_bits == 8: + return 480 + elif q_bits == 6: + return 28. + elif q_bits == 12: + return 510. + else: + assert (0), \ + "Please specify the right quantization range for the selected precision!" diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6e4eeaaea5d9b2c2ee70de3d4261c6348abe94 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_adam.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + +import sys + + +class FusedAdamBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/fused_adam_frontend.cpp', 'csrc/adam/multi_tensor_adam.cu'] + + def include_paths(self): + return ['csrc/includes', 'csrc/adam'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() + + def nvcc_args(self): + nvcc_flags = ['-O3'] + self.version_dependent_macros() + if not self.is_rocm_pytorch(): + nvcc_flags.extend( + ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + + self.compute_capability_args()) + return nvcc_flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lamb.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lamb.py new file mode 100644 index 0000000000000000000000000000000000000000..f0cb557706b3f7e6d4ab73d7a999a70aebecbeb3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lamb.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + +import sys + + +class FusedLambBuilder(CUDAOpBuilder): + BUILD_VAR = 'DS_BUILD_FUSED_LAMB' + NAME = "fused_lamb" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lamb.{self.NAME}_op' + + def sources(self): + return ['csrc/lamb/fused_lamb_cuda.cpp', 'csrc/lamb/fused_lamb_cuda_kernel.cu'] + + def include_paths(self): + return ['csrc/includes'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() + + def nvcc_args(self): + nvcc_flags = ['-O3'] + self.version_dependent_macros() + if self.is_rocm_pytorch(): + ROCM_MAJOR, ROCM_MINOR = self.installed_rocm_version() + nvcc_flags += ['-DROCM_VERSION_MAJOR=%s' % ROCM_MAJOR, '-DROCM_VERSION_MINOR=%s' % ROCM_MINOR] + else: + nvcc_flags.extend( + ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + + self.compute_capability_args()) + return nvcc_flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..b900a8f2369df14d65a73235d360db21583ee4ab --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lion.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + +import sys + + +class FusedLionBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_LION" + NAME = "fused_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/fused_lion_frontend.cpp', 'csrc/lion/multi_tensor_lion.cu'] + + def include_paths(self): + return ['csrc/includes', 'csrc/lion'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() + + def nvcc_args(self): + nvcc_flags = ['-O3'] + self.version_dependent_macros() + if not self.is_rocm_pytorch(): + nvcc_flags.extend( + ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + + self.compute_capability_args()) + return nvcc_flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/gds.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/gds.py new file mode 100644 index 0000000000000000000000000000000000000000..d815ccea387a78a0a6256b6f981b81346a3c230b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/gds.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from .async_io import AsyncIOBuilder + + +class GDSBuilder(AsyncIOBuilder): + BUILD_VAR = "DS_BUILD_GDS" + NAME = "gds" + + def __init__(self): + super().__init__() + + def absolute_name(self): + return f'deepspeed.ops.gds.{self.NAME}_op' + + def lib_sources(self): + src_list = ['csrc/gds/py_lib/deepspeed_py_gds_handle.cpp', 'csrc/gds/py_lib/deepspeed_gds_op.cpp'] + return super().lib_sources() + src_list + + def sources(self): + return self.lib_sources() + ['csrc/gds/py_lib/py_ds_gds.cpp'] + + def cxx_args(self): + return super().cxx_args() + ['-lcufile'] + + def include_paths(self): + import torch + CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] + return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE + + def extra_ldflags(self): + return super().extra_ldflags() + ['-lcufile'] + + def is_compatible(self, verbose=False): + if self.is_rocm_pytorch(): + if verbose: + self.warning(f'{self.NAME} is not compatible with ROCM') + return False + + try: + import torch.utils.cpp_extension + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile GDS") + return False + + CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME + if CUDA_HOME is None: + if verbose: + self.warning("Please install torch CUDA if trying to pre-compile GDS with CUDA") + return False + + CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") + gds_compatible = self.has_function(funcname="cuFileDriverOpen", + libraries=("cufile", ), + library_dirs=( + CUDA_HOME, + CUDA_LIB64, + ), + verbose=verbose) + + return gds_compatible and super().is_compatible(verbose) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad1b9a7f891eae0ceec8c26109ba2a235ef6349 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .cpu_adam import CPUAdamBuilder +from .fused_adam import FusedAdamBuilder +from .transformer_inference import InferenceBuilder +from .no_impl import NotImplementedBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5398ff40de82847f3354ade3e38478276e47fdeb Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..690a1ae30990ba5595a96aba83f3ecbfb79fc955 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efdd52aadd0a9fd4de4c6ed3a70e8871e274dbdc Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49284bf77d00a456973111371bf943c1290374d4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40ed757197be177e85b2d09582ba16ad321a5a85 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11843bd9d64b245718c12f44ab31095737d50cd7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..023127cfea9bd52a20960b766640f6b23541022b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..11e710a8ee4833bc66e82c4ed3241f7fe8a2e977 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/builder.py @@ -0,0 +1,38 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class CPUOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + args = ['-O3', '-g', '-Wno-reorder'] + return args + + def libraries_args(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..58eea2698ebb9eb6a8269a527a902792f1c7a6bc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py @@ -0,0 +1,28 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class CPUAdamBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c74affb55045a167f8dd3bb51b54f1487a815587 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py @@ -0,0 +1,86 @@ +# Copyright (c) 2024 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class FPQuantizerBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_FP_QUANTIZER" + NAME = "fp_quantizer" + + 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.fp_quantizer.{self.NAME}_op' + + def sources(self): + return [] + + def load(self, verbose=True): + return FPQuantizer + + @staticmethod + def get_default_quant_dtype(): + return torch.float8_e4m3fn + + @staticmethod + def get_quant_range(q_bits=None): + import habana_frameworks.torch.utils.experimental as htexp + if htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi2: + dtype = torch.float8_e4m3fnuz + else: + dtype = torch.float8_e4m3fn + return torch.finfo(dtype).max + + +class FPQuantizer: + CUDA_IMPL = False + + @classmethod + def selective_dequantize(cls, val_q, scales, indexes, group_size, q_mantisa_bits, q_exponent_bits): + assert False, "Selective dequantize isn't implemented for HPU!" + + @classmethod + def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits): + orig_shape = fp_out.shape + orig_dtype = fp_out.dtype + dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape) + fp_out.copy_(dequant_out) + return fp_out + + @classmethod + def quantize(cls, out, val, scale, group_size, stochastic_rounding, q_bits, q_mantisa_bits): + assert q_bits == 8, "Quantize on HPU only supports quantization to FP8" + assert q_mantisa_bits == 3, "Quantize on HPU only supports q_mantissa_bits = 3" + assert out.dtype.is_floating_point, "Quantization on HPU is only to float dtypes" + + num_groups, group_size = out.shape + + # Reshape the tensor + val_reshaped = val.view(num_groups, group_size).float() + # Calculate the scale + max_vals = val_reshaped.abs().max(dim=1, keepdim=True)[0] + q_range = torch.finfo(out.dtype).max + tmp_scale = q_range / max_vals + scale.copy_(tmp_scale) + # Copy quantized + quant, _ = torch.ops.hpu.cast_to_fp8_v2(val_reshaped, scale, stochastic_rounding, dtype=out.dtype) + out.copy_(quant) + + return out + + @classmethod + def get_scales(cls, out, num_groups): + return out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..5acb121668e37d1f7b8b660b831d35ae863d2a1a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py @@ -0,0 +1,105 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + +try: + import torch + import math +except ImportError as e: + pass + + +class HPUFusedAdam: + htcore = None + is_lazy_mode = None + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + + if HPUFusedAdam.htcore is None: + from habana_frameworks.torch import core as htcore + from habana_frameworks.torch.utils.internal import is_lazy + HPUFusedAdam.htcore = htcore + HPUFusedAdam.is_lazy_mode = is_lazy() + + htcore = HPUFusedAdam.htcore + + htcore.step_closure._mark_step_if_lazy() + step_size = lr + if bias_correction: + bias_correction1 = 1.0 - pow(beta1, step) + bias_correction2 = 1.0 - pow(beta2, step) + step_size = step_size * math.sqrt(bias_correction2) / bias_correction1 + + neg_step = -step_size + neg_step_t = (torch.tensor([neg_step], dtype=torch.float, + requires_grad=False).to(tensor_lists[1][0].dtype).to(tensor_lists[1][0].device, + non_blocking=True)) + + weight_decay = weight_decay if adam_w_mode else 0 + + # since lr is fed into the kernel as tensor, perform the scalar multiplication of wd here + # NOTE: TODO if lr is updated every step, then we need to convert it as tensor and + # perform weight decay unconditonally. + modified_wd = 1.0 - weight_decay * lr + + if HPUFusedAdam.is_lazy_mode: + torch.ops.hpu.optimizer_adamw( + tensor_lists[0], + tensor_lists[1], + tensor_lists[2], + tensor_lists[3], + neg_step_t, + beta1, + beta2, + epsilon, + modified_wd, + ) + else: + modified_wd_t = (torch.tensor([modified_wd], dtype=torch.float, requires_grad=False).to( + tensor_lists[1][0].dtype).to(tensor_lists[1][0].device, non_blocking=True)) + torch.ops.hpu.optimizer_adamw( + tensor_lists[0], + tensor_lists[1], + tensor_lists[2], + tensor_lists[3], + neg_step_t, + beta1, + beta2, + epsilon, + modified_wd_t, + modified_wd != 1.0, + ) + + htcore.step_closure._mark_step_if_lazy() + + +class FusedAdamBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return HPUFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..140d65b48defbea439483b63fd545e52de92209e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class NotImplementedBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + 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 load(self, verbose=True): + raise ValueError("This op had not been implemented on HPU backend.") + + def sources(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..e397c99200ec9155c933bc0691050b1f1e85cf7b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py @@ -0,0 +1,39 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +import importlib + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class InferenceBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=self.NAME) + + def absolute_name(self): + return f"deepspeed.ops.transformer.inference.{self.NAME}_op" + + def sources(self): + return [] + + def load(self, verbose=True): + if self.name in __class__._loaded_ops: + return __class__._loaded_ops[self.name] + + from deepspeed.git_version_info import installed_ops # noqa: F401 + if installed_ops.get(self.name, False): + op_module = importlib.import_module(self.absolute_name()) + __class__._loaded_ops[self.name] = op_module + return op_module diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_core_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_core_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b6665ebb76180c533f9700db222fca45314ecd19 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_core_ops.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class InferenceCoreBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_INFERENCE_CORE_OPS" + NAME = "inference_core_ops" + + 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.inference.v2.kernels{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 6: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/kernels/core_ops/core_ops.cpp", + "inference/v2/kernels/core_ops/bias_activations/bias_activation.cpp", + "inference/v2/kernels/core_ops/bias_activations/bias_activation_cuda.cu", + "inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.cpp", + "inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu", + "inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.cpp", + "inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu", + "inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.cpp", + "inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels_cuda.cu", + "inference/v2/kernels/core_ops/cuda_linear/linear_kernels.cpp", + "inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.cu", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + return [] + + def include_paths(self): + sources = [ + 'inference/v2/kernels/core_ops/bias_activations', + 'inference/v2/kernels/core_ops/blas_kernels', + 'inference/v2/kernels/core_ops/cuda_layer_norm', + 'inference/v2/kernels/core_ops/cuda_rms_norm', + 'inference/v2/kernels/core_ops/gated_activations', + 'inference/v2/kernels/core_ops/cuda_linear', + 'inference/v2/kernels/includes', + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + + return sources diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..a4a607288ca857f6812fdd925413493822bfdcb9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py @@ -0,0 +1,95 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class InferenceCutlassBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_CUTLASS_OPS" + NAME = "cutlass_ops" + + 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.inference.v2.kernels.cutlass_ops.{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 8: + # Only support Ampere and newer + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/kernels/cutlass_ops/cutlass_ops.cpp", + "inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.cu", + "inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.cu", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + import dskernels + lib_path = dskernels.library_path() + prefix = self.get_prefix() + lib_path = os.path.join(prefix, lib_path) + lib_path = self.deepspeed_src_path(lib_path) + + args = [f'-L{lib_path}', '-ldeepspeedft'] + if self.jit_load: + args.append(f'-Wl,-rpath,{lib_path}') + return args + + def include_paths(self): + sources = [ + 'inference/v2/kernels/includes', + 'inference/v2/kernels/cutlass_ops/mixed_gemm', + 'inference/v2/kernels/cutlass_ops/moe_gemm', + 'inference/v2/kernels/cutlass_ops/shared_resources/', + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db12afbbf20ec7a5329b95110dbceaf5e8607d10 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +# MLU related operators will be added in the future. +from .no_impl import NotImplementedBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .cpu_adam import CPUAdamBuilder +from .fused_adam import FusedAdamBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ef60702ac2ebb5661c99104be0932a984bbb036 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39e4b026a6da1f48c4596fea31aa1795dc2331a7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e430f3f9f67bd361a8c4f9d05cf543a11b90d144 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb948533db205fb5915702a716946b58feddcd5f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a29a1de8bf10090a68678dac91b262522fa7c8a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..123ffdb541aa216f2bc18664cf92c080b08ea94a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..17b9723ffcc1e3c81ea2e594d20b352b0a2a5bdf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/builder.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class MLUOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=self.strip_empty_entries(self.include_paths()), + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + return ['-O3', '-g', '-Wno-reorder'] + + def libraries_args(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..68b7bbe514eea6d7714a6f8b6135ebd622a1b7f9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + + +class CPUAdagradBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..b3c8e476bf3960c3ff467f362a72947cf73babf8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + + +class CPUAdamBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..2343e55568c489e1ff5d2354e88b676889c7211c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + +try: + import torch +except ImportError as e: + pass + + +class MLUFusedAdam: + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + + torch.ops.torch_mlu.fused_adam(noop_flag_buffer, tensor_lists[0], tensor_lists[1], tensor_lists[2], + tensor_lists[3], lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction, + weight_decay) + + +class FusedAdamBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return MLUFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..375c148b4a5e06e53fb87cddbdf988708d425995 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + + +class NotImplementedBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + 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 load(self, verbose=True): + raise ValueError("This op had not been implemented on MLU backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e512e25a2a29e1b463bdc373722a64c755b89098 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .fused_adam import FusedAdamBuilder +from .async_io import AsyncIOBuilder +from .no_impl import NotImplementedBuilder +from .cpu_adam import CPUAdamBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .cpu_lion import CPULionBuilder +from .inference import InferenceBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..163b776004947a0fcd43aaaece695ab6663a1c88 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22b411e53b4e03b9733a7900b4971907c7188d9f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b63f9280cbe007d5420f751b607d975d2114c0d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..682d69e5c95ffb79af54327a274424ceac2362e6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ddc97d0c66f39282c6d390ee5c8ce11cd2e8444 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..482f1b0df3694d15814abb805d3077ed2d4e7f91 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75b92c9246cd28a18a8aeda5e1fa456c422c9bd4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ec55a3d8a4658f45b24bcbbf787e07f1d82bdea Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64a1bc262eab3ef11b59ae46b902533387945a91 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..9a574ffe2dfbd2b52f42296460e348c6f619207c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/async_io.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess + +from .builder import NPUOpBuilder + + +class AsyncIOBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/aio/py_lib/deepspeed_py_copy.cpp', 'csrc/aio/py_lib/py_ds_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio.cpp', 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', + 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', 'csrc/aio/common/deepspeed_aio_utils.cpp', + 'csrc/aio/common/deepspeed_aio_common.cpp', 'csrc/aio/common/deepspeed_aio_types.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' + ] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/aio/py_lib', 'csrc/aio/common'] + return args + + def cxx_args(self): + args = super().cxx_args() + # -O0 for improved debugging, since performance is bound by I/O + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + import torch # Keep this import here to avoid errors when building DeepSpeed wheel without torch installed + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1: + CPP_STD = '-std=c++17' + else: + CPP_STD = '-std=c++14' + return args + [ + '-g', + '-Wall', + '-O0', + CPP_STD, + '-shared', + '-fPIC', + '-Wno-reorder', + CPU_ARCH, + '-fopenmp', + SIMD_WIDTH, + '-laio', + ] + + def extra_ldflags(self): + args = super().extra_ldflags() + return args + ['-laio'] + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_pgetevents', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..0dea2e78915e2a939d5c1839c2e2d88e948e10f1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/builder.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import re +import os +try: + import torch_npu +except ImportError as e: + pass + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class NPUOpBuilder(OpBuilder): + _ascend_path = None + _torch_npu_path = None + _cann_version = None + + def __init__(self, name): + super().__init__(name) + self._ascend_path = self.installed_cann_path() + self._torch_npu_path = os.path.join(os.path.dirname(os.path.abspath(torch_npu.__file__))) + try: + self._cann_version = self.installed_cann_version(self.name) + except BaseException: + print(f"{self.name} ascend_cann is missing, npu ops cannot be compiled!") + + def cann_defs(self): + if self._cann_version: + return '-D__ENABLE_CANN__' + return '-D__DISABLE_CANN__' + + def installed_cann_path(self): + if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): + return os.environ["ASCEND_HOME_PATH"] + return None + + def installed_cann_version(self, name=""): + ascend_path = self.installed_cann_path() + assert ascend_path is not None, "CANN_HOME does not exist, unable to compile NPU op(s)" + cann_version = "" + for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): + if cann_version: + break + install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] + if install_files: + filepath = os.path.join(dirpath, install_files[0]) + with open(filepath, "r") as f: + for line in f: + if line.find("version") != -1: + cann_version = line.strip().split("=")[-1] + break + return cann_version + + def include_paths(self): + paths = super().include_paths() + paths += [os.path.join(self._ascend_path, 'include'), os.path.join(self._torch_npu_path, 'include')] + return paths + + def cxx_args(self): + args = super().cxx_args() + args += ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-fopenmp'] + args += ['-fstack-protector-all', '-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath'] + args += [ + self.cann_defs(), + self.cpu_arch(), + self.simd_width(), '-L' + os.path.join(self._ascend_path, 'lib64'), + '-L' + os.path.join(self._torch_npu_path, 'lib') + ] + return args + + def extra_ldflags(self): + flags = super().extra_ldflags() + flags += [ + '-L' + os.path.join(self._ascend_path, 'lib64'), '-lascendcl', + '-L' + os.path.join(self._torch_npu_path, 'lib'), '-ltorch_npu' + ] + return flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..161bc82efe1ca01660fdeedd30079a8f10f1d269 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPUAdagradBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..a4e9569c0f336122cd003a2df5e196527d84666c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPUAdamBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..6917e0fd03d08dec42e71479110224d577b55b5b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPULionBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_LION" + NAME = "cpu_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..d32103db7055669f8ecce578bb4ef7703e3c07ef --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + +try: + import torch_npu +except ImportError as e: + pass + + +class NPUFusedAdam: + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + bias_correction1 = beta1**(step - 1) + bias_correction2 = beta2**(step - 1) + + # iteration group['params'] + for i in range(len(tensor_lists[0])): + grad_flat = tensor_lists[0][i] + param_flat = tensor_lists[1][i] + m_flat = tensor_lists[2][i] + v_flat = tensor_lists[3][i] + + if adam_w_mode: + param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam_w( + bias_correction1, + bias_correction2, + lr, + weight_decay, + beta1, + beta2, + epsilon, + grad_flat, + None, # max_grad_norm + False, # amsgrad + False, # maximize + out=(param_flat.data, m_flat, v_flat)) + else: + param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam( + bias_correction1, + bias_correction2, + lr, + beta1, + beta2, + epsilon, + grad_flat, + False, # use_locking + False, # use_nesterov + out=(param_flat.data, m_flat, v_flat)) + + +class FusedAdamBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return NPUFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..46f28c0d401161f70431776a5a53387235ebb5ce --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/inference.py @@ -0,0 +1,307 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import IntEnum +from .builder import NPUOpBuilder + +try: + import torch + import torch_npu +except ImportError as e: + pass + + +class ActivationFuncType(IntEnum): + UNKNOWN = 0 + GELU = 1 + ReLU = 2 + GATED_GELU = 3 + GATED_SILU = 4 + + +class InferenceContext: + _workspace = None + + _seed = 42 + _curr_offset = 0 + _stream = 0 + _free_memory_size = 0 + _num_tokens = 1 + _attention_unfused_workspace_offset = 0 + _workSpaceSize = 0 + + workSpaceSize = 0 + kv_caches = None + + @staticmethod + def reset_tokens(initial_tokens=1): + InferenceContext._num_tokens = initial_tokens + + @staticmethod + def current_tokens(): + return InferenceContext._num_tokens + + @staticmethod + def GetWorkSpace(): + return InferenceContext._workspace + + +class NPUInference: + + @staticmethod + def layer_norm(inputs, gamma, beta, epsilon): + return torch.nn.functional.layer_norm(inputs, [inputs.shape[-1]], gamma, beta, eps=epsilon) + + @staticmethod + def _qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + inp_norm = torch.nn.functional.layer_norm(inputs, (inputs.shape[2], ), gamma, beta, eps) + weight = weight.t() if transpose else weight + tmp = torch.matmul(inp_norm, weight) + if add_bias: + tmp += bias + output = [tmp, inp_norm] + return output + + @staticmethod + def qkv_gemm_fp16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) + + @staticmethod + def qkv_gemm_bf16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) + + @staticmethod + def qkv_gemm_fp32(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) + + @staticmethod + def _bias_add_transform_0213(vals, bias, hidden_dim, seq_length, seq_offset, heads, num_kv, rotary_dim, + rotate_half, rotate_every_two, rope_theta): + bsz, _, _ = vals.shape + q = vals[..., :hidden_dim].reshape(bsz, seq_length, heads, -1) + k = vals[..., hidden_dim:hidden_dim + num_kv * (hidden_dim // heads)].reshape(bsz, seq_length, num_kv, -1) + v = vals[..., hidden_dim + num_kv * (hidden_dim // heads):] + + if rotary_dim > 0 and rotate_every_two: + # sin, cos may use cache + seq_id = torch.arange(0, seq_length).to("npu") + inv_freq = torch.arange(0, rotary_dim, 2) / rotary_dim + inv_freq = inv_freq.to("npu") + inv_freq = 1.0 / torch.pow(rope_theta, inv_freq) + inv_freq = torch.outer(seq_id, inv_freq) + sin = inv_freq.sin() + cos = inv_freq.cos() + # shape: [bsz=1, seq_len, heads=1, rotary_dim] + sin = sin.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1) + cos = cos.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1) + + q_pos, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_pos, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + + q_pos = torch_npu.npu_rotary_mul(q_pos, cos, sin) + q = torch.cat([q_pos, q_pass], dim=-1) + k_pos = torch_npu.npu_rotary_mul(k_pos, cos, sin) + k = torch.cat([k_pos, k_pass], dim=-1) + + output = q.reshape(bsz, seq_length, -1).contiguous() # [b, s, H] + k_cache = k.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d] + v_cache = v.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d] + return output, k_cache, v_cache + + @staticmethod + def _softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + bsz, seq_len, k = query_key_value.size() + k = k // (heads + 2 * (num_kv if num_kv > 0 else heads)) + hidden_dim = heads * k + + is_promt = seq_len > 1 + if not InferenceContext.kv_caches: + InferenceContext.kv_caches = [[None, None] for _ in range(num_layers)] + if is_promt: + InferenceContext.reset_tokens(seq_len) + InferenceContext.kv_caches[layer_id] = [None, None] + + soft_len = InferenceContext.current_tokens() + workspace = InferenceContext.GetWorkSpace() + seq_offset = 0 if is_promt else soft_len - 1 + + q, k, v = NPUInference._bias_add_transform_0213(vals=query_key_value, + bias=None, + hidden_dim=hidden_dim, + seq_length=seq_len, + seq_offset=seq_offset, + heads=heads, + num_kv=num_kv if num_kv > 0 else heads, + rotary_dim=rotary_dim, + rotate_half=rotate_half, + rotate_every_two=rotate_every_two, + rope_theta=rope_theta) + + if not is_promt: + k_cache, v_cache = InferenceContext.kv_caches[layer_id] + if k_cache is not None: + k = torch.cat([k_cache, k], dim=2) + v = torch.cat([v_cache, v], dim=2) + InferenceContext.kv_caches[layer_id] = [k, v] + seq_len = k.shape[2] + + layer_scale = max(1, layer_id) if len(alibi.size()) > 1 else 1.0 + alpha = norm_factor * norm_factor / layer_scale + + output = torch_npu.npu_fusion_attention(q, + k.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(), + v.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(), + heads, + "BSH", + pse=None, + padding_mask=None, + atten_mask=attn_mask.bool(), + scale=alpha, + pre_tockens=65536, + next_tockens=65536, + keep_prob=1, + inner_precise=0)[0] + + return output, k, v + + @staticmethod + def softmax_context_fp16(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, + heads, num_kv, norm_factor, triangular_masking, local_attention, + window_size, no_masking, layer_id, num_layers, alibi, rope_theta) + + @staticmethod + def softmax_context_bf16(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, + heads, num_kv, norm_factor, triangular_masking, local_attention, + window_size, no_masking, layer_id, num_layers, alibi, rope_theta) + + @staticmethod + def softmax_context_fp32(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, + heads, num_kv, norm_factor, triangular_masking, local_attention, + window_size, no_masking, layer_id, num_layers, alibi, rope_theta) + + @staticmethod + def _vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode): + if transposed_mode: + return torch.matmul(input, weight.t()) + return torch.matmul(input, weight) + + @staticmethod + def vector_matmul_fp16(input, weight, async_op, q_scale, q_int8, transposed_mode): + return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) + + @staticmethod + def vector_matmul_bf16(input, weight, async_op, q_scale, q_int8, transposed_mode): + return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) + + @staticmethod + def vector_matmul_fp32(input, weight, async_op, q_scale, q_int8, transposed_mode): + return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) + + @staticmethod + def _mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + if mlp_after_attn: + residual_add = torch.nn.functional.layer_norm(input + residual + input_bias, (input.shape[-1], ), gamma, + beta, eps) + else: + residual_add = torch.nn.functional.layer_norm(input, (input.shape[-1], ), gamma, beta, eps) + + weight_interm = weight_interm.t() if transpose else weight_interm + tmp = torch.matmul(residual_add, weight_interm) + if mlp_act_func_type == ActivationFuncType.GELU: + tmp = torch.nn.functional.gelu(tmp + bias) + elif mlp_act_func_type == ActivationFuncType.ReLU: + tmp = torch.nn.functional.relu(tmp + bias) + else: + raise Exception('Unsupported ActivationFuncType {}'.format(mlp_act_func_type)) + output = torch.matmul(tmp, weight_out.t()) + return output, residual_add + + @staticmethod + def mlp_gemm_fp16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, + mlp_act_func_type, transpose) + + @staticmethod + def mlp_gemm_bf16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, + mlp_act_func_type, transpose) + + @staticmethod + def mlp_gemm_fp32(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, + mlp_act_func_type, transpose) + + @staticmethod + def _residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + if mlp_after_attn: + if pre_layer_norm: + tmp = (residual.float() + attention_output.float() + attention_bias.float() + + final_bias.float()) / mp_size + hidden_state.float() + else: + tmp = residual.float() + hidden_state.float() + final_bias.float() + else: + if add_bias: + residual += attention_bias.float() + tmp = hidden_state.float() + attention_output.float() + (residual.float() + final_bias.float()) / mp_size + + input_dtype = hidden_state.dtype + residual.set_(tmp.to(input_dtype)) + + @staticmethod + def residual_add_bias_fp16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + mp_size, mlp_after_attn, add_bias, pre_layer_norm) + + @staticmethod + def residual_add_bias_bf16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + mp_size, mlp_after_attn, add_bias, pre_layer_norm) + + @staticmethod + def residual_add_bias_fp32(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + mp_size, mlp_after_attn, add_bias, pre_layer_norm) + + +class InferenceBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.transformer.inference.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return NPUInference diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..5b1771fabc22f6ad13161231900a737bba733e68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/no_impl.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class NotImplementedBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + 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 load(self, verbose=True): + raise ValueError("This op had not been implemented on NPU backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..0b5348e5af96e28e256524b38efaff5d50863d17 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/quantizer.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + + +class QuantizerBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_QUANTIZER" + NAME = "quantizer" + + 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.quantizer.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/quantization/pt_binding.cpp', + 'csrc/quantization/fake_quantizer.cu', + 'csrc/quantization/quantize.cu', + 'csrc/quantization/quantize_intX.cu', + 'csrc/quantization/dequantize.cu', + 'csrc/quantization/swizzled_quantize.cu', + 'csrc/quantization/quant_reduce.cu', + ] + + def include_paths(self): + return ['csrc/includes'] + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0df28cc2282a0beee59e4e5ae4d84c3a94ef9aa3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_ops.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class RaggedOpsBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_RAGGED_DEVICE_OPS" + NAME = "ragged_device_ops" + + 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.inference.v2.kernels.ragged_ops.{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 8: + # Blocked flash has a dependency on Ampere + newer + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/kernels/ragged_ops/ragged_ops.cpp", + "inference/v2/kernels/ragged_ops/atom_builder/atom_builder.cpp", + "inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.cpp", + "inference/v2/kernels/ragged_ops/embed/embed.cpp", + "inference/v2/kernels/ragged_ops/embed/embed_cuda.cu", + "inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cpp", + "inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu", + "inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cpp", + "inference/v2/kernels/ragged_ops/logits_gather/logits_gather_cuda.cu", + "inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cpp", + "inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter_cuda.cu", + "inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cpp", + "inference/v2/kernels/ragged_ops/moe_gather/moe_gather_cuda.cu", + "inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.cpp", + "inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cpp", + "inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating_cuda.cu", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + import dskernels + lib_path = dskernels.library_path() + + prefix = self.get_prefix() + lib_path = os.path.join(prefix, lib_path) + lib_path = self.deepspeed_src_path(lib_path) + + args = [f'-L{lib_path}', '-lblockedflash'] + if self.jit_load: + args.append(f'-Wl,-rpath,{lib_path}') + return args + + def include_paths(self): + sources = [ + 'inference/v2/kernels/includes', + 'inference/v2/kernels/ragged_ops', + 'inference/v2/kernels/ragged_ops/atom_builder', + 'inference/v2/kernels/ragged_ops/blocked_flash', + 'inference/v2/kernels/ragged_ops/embed', + 'inference/v2/kernels/ragged_ops/includes', + 'inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary', + 'inference/v2/kernels/ragged_ops/logits_gather', + 'inference/v2/kernels/ragged_ops/moe_gather', + 'inference/v2/kernels/ragged_ops/moe_scatter', + 'inference/v2/kernels/ragged_ops/ragged_helpers', + 'inference/v2/kernels/ragged_ops/top_k_gating', + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..208c9f833ebe97c39593e66c46a53b71190fedfd --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_utils.py @@ -0,0 +1,80 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class RaggedUtilsBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_RAGGED_OPS" + NAME = "ragged_ops" + + 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.inference.v2.{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 6: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/ragged/csrc/fast_host_buffer.cu", + "inference/v2/ragged/csrc/ragged_ops.cpp", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + return [] + + def include_paths(self): + include_dirs = ['inference/v2/ragged/includes', 'inference/v2/kernels/includes'] + prefix = self.get_prefix() + includes = [os.path.join(prefix, include_dir) for include_dir in include_dirs] + + return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/random_ltd.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/random_ltd.py new file mode 100644 index 0000000000000000000000000000000000000000..54af7150fb36f9eb8bd6a295648d8a11da8d8373 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/random_ltd.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + + +class RandomLTDBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_RANDOM_LTD" + NAME = "random_ltd" + + 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.{self.NAME}_op' + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def sources(self): + return [ + 'csrc/random_ltd/pt_binding.cpp', 'csrc/random_ltd/gather_scatter.cu', + 'csrc/random_ltd/slice_attn_masks.cu', 'csrc/random_ltd/token_sort.cu' + ] + + def include_paths(self): + includes = ['csrc/includes'] + return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a6eb8bbfa2c52c3e61ddbca734011c2ec5d630b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. + +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .no_impl import NotImplementedBuilder +from .cpu_adam import CPUAdamBuilder +from .fused_adam import FusedAdamBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a72c6e77032dc7c1510f512000ece69b9f15304 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad16d07da9757ba5c138fd4922d60b97fd7f70d2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cf4c2a96759c73f826fabf776775c391c185679 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea9ef88279f91862f7b4ae0d5808f533656e1fc5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62ec296ca461c110a9ed79b7ca22ab0138d59783 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..81f0e98c076837650bd5a3cd9a23493361e6b02e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/builder.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class SDAAOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=self.strip_empty_entries(self.include_paths()), + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + return ['-O3', '-g', '-Wno-reorder'] + + def libraries_args(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..b38a71a6275d3b962ed2bd7d8f625ba3298418dd --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +from .builder import SDAAOpBuilder + + +class CPUAdamBuilder(SDAAOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..73a2dff4145961a909258278e64a6fd094ca0206 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +from .builder import SDAAOpBuilder + +try: + import torch +except ImportError as e: + pass + + +class SDAAFusedAdam: + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + g_tensor_lis, p_tensor_lis, m_tensor_lis, v_tensor_lis = tensor_lists + torch.ops.sdaa.fused_adam(g_tensor_lis, p_tensor_lis, m_tensor_lis, v_tensor_lis, [], beta1, beta2, epsilon, + lr, weight_decay, adam_w_mode, step, bias_correction) + + +class FusedAdamBuilder(SDAAOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return SDAAFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..10a8b8f48652a8abbd9978fb444332036571acb3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +from .builder import SDAAOpBuilder + + +class NotImplementedBuilder(SDAAOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + 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 load(self, verbose=True): + raise ValueError("This op had not been implemented on SDAA backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sparse_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sparse_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..b8e5e8a6fd9aaa62372f72c0b78d250994b3e98e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sparse_attn.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import OpBuilder + +try: + from packaging import version as pkg_version +except ImportError: + pkg_version = None + + +class SparseAttnBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_SPARSE_ATTN" + NAME = "sparse_attn" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.sparse_attention.{self.NAME}_op' + + def sources(self): + return ['csrc/sparse_attention/utils.cpp'] + + def cxx_args(self): + return ['-O2', '-fopenmp'] + + def is_compatible(self, verbose=False): + # Check to see if llvm and cmake are installed since they are dependencies + #required_commands = ['llvm-config|llvm-config-9', 'cmake'] + #command_status = list(map(self.command_exists, required_commands)) + #deps_compatible = all(command_status) + + if self.is_rocm_pytorch(): + if verbose: + self.warning(f'{self.NAME} is not compatible with ROCM') + return False + + try: + import torch + except ImportError: + if verbose: + self.warning(f"unable to import torch, please install it first") + return False + + # torch-cpu will not have a cuda version + if torch.version.cuda is None: + cuda_compatible = False + if verbose: + self.warning(f"{self.NAME} cuda is not available from torch") + else: + major, minor = torch.version.cuda.split('.')[:2] + cuda_compatible = (int(major) == 10 and int(minor) >= 1) or (int(major) >= 11) + if not cuda_compatible: + if verbose: + self.warning(f"{self.NAME} requires CUDA version 10.1+") + + TORCH_MAJOR = int(torch.__version__.split('.')[0]) + TORCH_MINOR = int(torch.__version__.split('.')[1]) + torch_compatible = (TORCH_MAJOR == 1 and TORCH_MINOR >= 5) + if not torch_compatible: + if verbose: + self.warning( + f'{self.NAME} requires a torch version >= 1.5 and < 2.0 but detected {TORCH_MAJOR}.{TORCH_MINOR}') + try: + import triton + except ImportError: + # auto-install of triton is broken on some systems, reverting to manual install for now + # see this issue: https://github.com/deepspeedai/DeepSpeed/issues/1710 + if verbose: + self.warning(f"please install triton==1.0.0 if you want to use sparse attention") + return False + + if pkg_version: + installed_triton = pkg_version.parse(triton.__version__) + triton_mismatch = installed_triton != pkg_version.parse("1.0.0") + else: + installed_triton = triton.__version__ + triton_mismatch = installed_triton != "1.0.0" + + if triton_mismatch: + if verbose: + self.warning( + f"using untested triton version ({installed_triton}), only 1.0.0 is known to be compatible") + return False + + return super().is_compatible(verbose) and torch_compatible and cuda_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c5fa6611563bc0a45fe712e27f0f06f56bf2d0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class SpatialInferenceBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_SPATIAL_INFERENCE" + NAME = "spatial_inference" + + 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.spatial.{self.NAME}_op' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def sources(self): + return [ + 'csrc/spatial/csrc/opt_bias_add.cu', + 'csrc/spatial/csrc/pt_binding.cpp', + ] + + def include_paths(self): + return ['csrc/spatial/includes', 'csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..52b02a3c629e9d537b8e22139f48f5335396e4ff --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py @@ -0,0 +1,22 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .transformer import TransformerBuilder + + +class StochasticTransformerBuilder(TransformerBuilder): + BUILD_VAR = "DS_BUILD_STOCHASTIC_TRANSFORMER" + NAME = "stochastic_transformer" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.transformer.{self.NAME}_op' + + def nvcc_args(self): + args = super().nvcc_args() + args.append('-D__STOCHASTIC_MODE__') + return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..8db30fdc67919474679715014cbf8d3f470bedf4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + + +class TransformerBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER" + NAME = "transformer" + + 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.transformer.{self.NAME}_op' + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def sources(self): + return [ + 'csrc/transformer/ds_transformer_cuda.cpp', 'csrc/transformer/cublas_wrappers.cu', + 'csrc/transformer/transform_kernels.cu', 'csrc/transformer/gelu_kernels.cu', + 'csrc/transformer/dropout_kernels.cu', 'csrc/transformer/normalize_kernels.cu', + 'csrc/transformer/softmax_kernels.cu', 'csrc/transformer/general_kernels.cu' + ] + + def include_paths(self): + includes = ['csrc/includes'] + return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..642aed56a192e4b5170a62d8cc457bb451ae8646 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer_inference.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class InferenceBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + 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.transformer.inference.{self.NAME}_op' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 6: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def sources(self): + return [ + 'csrc/transformer/inference/csrc/pt_binding.cpp', + 'csrc/transformer/inference/csrc/gelu.cu', + 'csrc/transformer/inference/csrc/relu.cu', + 'csrc/transformer/inference/csrc/layer_norm.cu', + 'csrc/transformer/inference/csrc/rms_norm.cu', + 'csrc/transformer/inference/csrc/softmax.cu', + 'csrc/transformer/inference/csrc/dequantize.cu', + 'csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu', + 'csrc/transformer/inference/csrc/transform.cu', + 'csrc/transformer/inference/csrc/pointwise_ops.cu', + ] + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def include_paths(self): + return ['csrc/transformer/inference/includes', 'csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..927ff1b361bd0876c065b6c5f8904c6668fd2ace --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/utils.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import OpBuilder + + +class UtilsBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_UTILS" + NAME = "utils" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def sources(self): + return ['csrc/utils/tensor_cast.cpp', 'csrc/utils/py_ds_utils.cpp'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..093f90b30234359d37913bb36919f3259d4cab1c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .cpu_adam import CPUAdamBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .fused_adam import FusedAdamBuilder +from .async_io import AsyncIOBuilder +from .inference import InferenceBuilder +from .flash_attn import FlashAttentionBuilder +from .no_impl import NotImplementedBuilder +from .packbits import PackbitsBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48c6a39f32acaeaadaa8b7b18e00eb17bbe4ec20 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..956fa57de5d741b9fb34112f83260e3b8100abc5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..794726222ea34077fc50f11f2df36f2dfbcca70d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09d5c83cc0ddaf3021da2e2962105a9489e241dc Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..581a5afa69b7593d420fe43d33f0e81cd606b161 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd7a8adaba9df0125c674adca0a08a76d0b1fdf0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eade3a536ecb05b0ad584abd0e7f8489bf1bb573 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30840b286e6ac27c0de9feb2fa9e7d449f0a0d73 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05726b588cc3494a3f68654eb6f108d629f87d51 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cb3d61bbbf1711c1c848c4484bdbb40a935e0a7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..8ec030880368dff58337b9e78bfd80bad9f31f97 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/async_io.py @@ -0,0 +1,106 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess + +from .builder import OpBuilder + + +class AsyncIOBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/aio/py_lib/deepspeed_py_copy.cpp', + 'csrc/aio/py_lib/py_ds_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', + 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', + 'csrc/aio/common/deepspeed_aio_utils.cpp', + 'csrc/aio/common/deepspeed_aio_common.cpp', + 'csrc/aio/common/deepspeed_aio_types.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp', + 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', + 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', + 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', + ] + + def include_paths(self): + return ['csrc/aio/py_lib', 'csrc/aio/common'] + + def cxx_args(self): + import torch + # -O0 for improved debugging, since performance is bound by I/O + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1: + CPP_STD = '-std=c++17' + else: + CPP_STD = '-std=c++14' + return [ + '-g', + '-Wall', + '-O0', + CPP_STD, + '-shared', + '-fPIC', + '-Wno-reorder', + CPU_ARCH, + '-fopenmp', + SIMD_WIDTH, + '-laio', + ] + + def extra_ldflags(self): + return ['-laio'] + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_pgetevents', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..81b15f197f4375b6138f50cf865b2dd64194c629 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/builder.py @@ -0,0 +1,131 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import time +import importlib + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class SYCLOpBuilder(OpBuilder): + + def builder(self): + try: + from intel_extension_for_pytorch.xpu.cpp_extension import DPCPPExtension + except ImportError: + from intel_extension_for_pytorch.xpu.utils import DPCPPExtension + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + print("dpcpp sources = {}".format(self.sources())) + dpcpp_ext = DPCPPExtension(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + extra_compile_args={ + 'cxx': self.strip_empty_entries(self.cxx_args()), + }, + extra_link_args=self.strip_empty_entries(self.fixed_aotflags())) + return dpcpp_ext + + def version_dependent_macros(self): + try: + from op_builder.builder import TORCH_MAJOR, TORCH_MINOR + except ImportError: + from deepspeed.ops.op_builder.builder import TORCH_MAJOR, TORCH_MINOR + # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456 + version_ge_1_1 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0): + version_ge_1_1 = ['-DVERSION_GE_1_1'] + version_ge_1_3 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2): + version_ge_1_3 = ['-DVERSION_GE_1_3'] + version_ge_1_5 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): + version_ge_1_5 = ['-DVERSION_GE_1_5'] + return version_ge_1_1 + version_ge_1_3 + version_ge_1_5 + + def cxx_args(self): + cxx_flags = [ + '-fsycl', '-fsycl-targets=spir64_gen', '-g', '-gdwarf-4', '-O3', '-std=c++17', '-fPIC', '-DMKL_ILP64', + '-fno-strict-aliasing' + ] + if os.environ.get('USE_MKL_GEMM'): + cxx_flags.append('-DUSE_MKL_GEMM') + return cxx_flags + + def extra_ldflags(self): + return [ + '-fPIC', '-fsycl', '-fsycl-targets=spir64_gen', '-fsycl-max-parallel-link-jobs=8', + '-Xs "-options -cl-poison-unsupported-fp64-kernels,cl-intel-enable-auto-large-GRF-mode"', + '-Xs "-device pvc"', '-Wl,-export-dynamic' + ] + + def fixed_aotflags(self): + return [ + '-fsycl', '-fsycl-targets=spir64_gen', '-fsycl-max-parallel-link-jobs=8', '-Xs', + "-options -cl-poison-unsupported-fp64-kernels,cl-intel-enable-auto-large-GRF-mode", '-Xs', "-device pvc" + ] + + def load(self, verbose=True): + from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name # noqa: F401 + from deepspeed.accelerator import get_accelerator + if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: + return importlib.import_module(self.absolute_name()) + else: + return self.jit_load(verbose) + + def jit_load(self, verbose=True): + if not self.is_compatible(verbose): + raise RuntimeError( + f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" + ) + try: + import ninja # noqa: F401 + except ImportError: + raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") + + self.jit_mode = True + from intel_extension_for_pytorch.xpu.cpp_extension import load + + start_build = time.time() + # Recognize relative paths as absolute paths for jit load + + sources = [self.deepspeed_src_path(path) for path in self.sources()] + extra_include_paths = [self.deepspeed_src_path(path) for path in self.include_paths()] + + # Torch will try and apply whatever CCs are in the arch list at compile time, + # we have already set the intended targets ourselves we know that will be + # needed at runtime. This prevents CC collisions such as multiple __half + # implementations. Stash arch list to reset after build. + ''' + torch_arch_list = None + if "TORCH_CUDA_ARCH_LIST" in os.environ: + torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") + os.environ["TORCH_CUDA_ARCH_LIST"] = "" + ''' + + op_module = load( + name=self.name, + sources=self.strip_empty_entries(sources), + extra_include_paths=self.strip_empty_entries(extra_include_paths), + extra_cflags=self.strip_empty_entries(self.cxx_args()), + # extra_cuda_cflags=self.strip_empty_entries(self.nvcc_args()), + extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), + verbose=verbose) + + build_duration = time.time() - start_build + if verbose: + print(f"Time to load {self.name} op: {build_duration} seconds") + ''' + # Reset arch list so we are not silently removing it for other possible use cases + if torch_arch_list: + os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list + ''' + return op_module diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..18f80848e1b800de7f8826405bb9ea9f500db263 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import SYCLOpBuilder + + +class CPUAdagradBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/xpu/adagrad/cpu_adagrad.cpp', 'csrc/xpu/common/custom_cuda_kernel.dp.cpp'] + + def include_paths(self): + return ['csrc/xpu/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..9d5fdcd3e8ab51b84bb2fec137b95ee96fee8570 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import SYCLOpBuilder + + +class CPUAdamBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..c8c2674d5d27c87377d3a64bfbba401a44ae9d52 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class FlashAttentionBuilderObject(): + + def __init__(self): + pass + + # general functions + def flash_attn_func_v2(self, q, k, v, dropout_p, softmax_scale, is_causal): + try: + import torch + import intel_extension_for_pytorch # noqa + return torch.nn.functional.scaled_dot_product_attention(q, + k, + v, + dropout_p=dropout_p, + is_causal=is_causal, + scale=softmax_scale) + except ImportError: + raise ImportError( + "Please install pytorch and intel_extension_for_pytorch to include scaled dot product attention.") + + +class FlashAttentionBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_FlashAttention" + NAME = "flash_attn" + + 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.{self.NAME}_op' + + def sources(self): + return + + def include_paths(self): + return [] + + def extra_ldflags(self): + return [] + + def cxx_args(self): + return [] + + def load(self): + return FlashAttentionBuilderObject() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..0e0f1a66f8e648305d2ba6fd0ad85a384e6f8bd3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class FusedAdamBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/xpu/adam/fused_adam_frontend.cpp', 'csrc/xpu/adam/multi_tensor_adam.dp.cpp'] + + def include_paths(self): + return ['csrc/xpu/includes', 'csrc/xpu/adam'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..a9ac4f84c2ca9ba885e586934efd406b6469bd9e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/inference.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class InferenceBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + 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.transformer.inference.{self.NAME}_op' + + def sources(self): + return + + def include_paths(self): + return [] + + def extra_ldflags(self): + return [] + + def cxx_args(self): + return [] + + def load(self): + try: + import intel_extension_for_pytorch + if hasattr(intel_extension_for_pytorch, "deepspeed"): + return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference + else: + return intel_extension_for_pytorch.xpu.deepspeed + except ImportError: + raise ImportError("Please install intel-extension-for-pytorch >= 2.1.30 to include DeepSpeed kernels.") diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..8b294f70c2791db9686e8d463ad736ff2d7c90c5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import SYCLOpBuilder + + +class NotImplementedBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + 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.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on XPU backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/packbits.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/packbits.py new file mode 100644 index 0000000000000000000000000000000000000000..cf5b5ebc59e407ae20f478151846e282680fe4ae --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/packbits.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class PackbitsBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_PACK_BITS" + NAME = "pack_bits" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def sources(self): + return ['csrc/xpu/packbits/packing.cpp'] + + def include_paths(self): + return ['csrc/xpu/includes'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0a5d1da5e3ae0fa097e7313ddb1328c4f910801d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .quantizer import ds_quantizer diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6306206049784dca7dacb649eb9d2ced11d7bb5d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26af8c3833a98b313b5a2610089d29d7772369f5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..eb4bfd35700075f3b32db329c5f7026b39bef520 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/quantizer.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from deepspeed.ops.op_builder import QuantizerBuilder + +# Cuda modules will be imported if needed +quantizer_cuda_module = None + + +def ds_quantizer(input, groups=1, bit_num=8, sr=False, asym=False): + # Load cuda modules if needed + global quantizer_cuda_module + if quantizer_cuda_module is None: + quantizer_cuda_module = QuantizerBuilder().load() + if sr: + if asym: + quantize_func = quantizer_cuda_module.ds_sr_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_asym_fp32 + else: + quantize_func = quantizer_cuda_module.ds_sr_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_fp32 + else: + if asym: + quantize_func = quantizer_cuda_module.ds_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_asym_fp32 + else: + quantize_func = quantizer_cuda_module.ds_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_fp32 + return quantize_func(input, groups, bit_num) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1ade5811ce76d5b7bf51881db5ff0824c1fd379d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .dropping_utils import gpt_sample_tokens, bert_sample_tokens, GatherTokens, ScatterTokens diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef9ef6640ab54329bf9783fa46d9e67dc179513e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b78227479c82388b4267866246d669e0da5a8c7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/dropping_utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/dropping_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..dd36c94537f896831291458514b513ea840791e2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/dropping_utils.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from deepspeed.ops.op_builder import RandomLTDBuilder +""" +Returns: + sampled_indices: [layers, batch_size, reserved_length] + new_mask: [batch_size, 1, reserved_length, reserved_length] +""" + +random_ltd_module = None + + +def gpt_sample_tokens(reserved_length: int, + seq_length: int, + batch_size: int, + layers: int = 1, + device: str = 'cpu', + attn_mask: torch.Tensor = None): + + prob_dist = torch.ones((layers * batch_size, seq_length), device=device) + sampled_indices = torch.multinomial(prob_dist, reserved_length) + + sampled_indices = sampled_indices.reshape(layers, batch_size, reserved_length).to(torch.int32) + global random_ltd_module + if random_ltd_module is None: + random_ltd_module = RandomLTDBuilder().load() + sampled_indices = random_ltd_module.token_sort_(sampled_indices, seq_length) + + # Not certain the optimized kernel is actually better here, cause it kind of screws + # with alignment right if the sequence length is not divisible by like 16 + # new_mask = random_ltd_module.mask_gather_gpt(attn_mask, reserved_length) + if attn_mask is not None: + new_mask = attn_mask[:, :, :reserved_length, :reserved_length] + else: + new_mask = None + + return sampled_indices, new_mask + + +""" +Returns: + sampled_indices: [layers, batch_size, reserved_length] + new_mask: [layers, batch_size, 1, reserved_length, reserved_length] +""" + + +def bert_sample_tokens(reserved_length: int, + seq_length: int, + batch_size: int, + layers: int = 1, + device: str = 'cpu', + attn_mask: torch.Tensor = None): + assert attn_mask is not None + prob_dist = torch.ones((layers * batch_size, seq_length), device=device) + sampled_indices = torch.multinomial(prob_dist, reserved_length) + + sampled_indices = sampled_indices.reshape(layers, batch_size, reserved_length).to(torch.int32) + global random_ltd_module + if random_ltd_module is None: + random_ltd_module = RandomLTDBuilder().load() + + sampled_indices = random_ltd_module.token_sort_(sampled_indices, seq_length) + dtype = sampled_indices.dtype + + sampled_indices = sampled_indices.to(torch.long) + new_mask = [] + for l in range(layers): + tmp_mask_list = [] + for i in range(batch_size): + mask_tmp = attn_mask[i:i + 1, :, sampled_indices[l][i], :] + tmp_mask_list.append(mask_tmp[:, :, :, sampled_indices[l][i]]) + new_mask.append(torch.cat(tmp_mask_list, dim=0)) + + return sampled_indices.to(dtype), new_mask + + +class GatherTokens(torch.autograd.Function): + + @staticmethod + def forward(ctx, activations: torch.Tensor, sorted_indices: torch.Tensor, batch_first: bool): + global random_ltd_module + if random_ltd_module is None: + random_ltd_module = RandomLTDBuilder().load() + ctx.save_for_backward(activations, sorted_indices) + ctx.batch_first = batch_first + return activations, random_ltd_module.token_gather(activations, sorted_indices, batch_first) + + @staticmethod + def backward(ctx, a_gradients: torch.Tensor, g_gradients: torch.Tensor): + + g_gradients = g_gradients.contiguous() + global random_ltd_module + if random_ltd_module is None: + random_ltd_module = RandomLTDBuilder().load() + activations, sorted_indices = ctx.saved_tensors + batch_first = ctx.batch_first + + return random_ltd_module.token_scatter_(a_gradients, g_gradients, sorted_indices, batch_first), None, None + + +class ScatterTokens(torch.autograd.Function): + + @staticmethod + def forward(ctx, all_activations: torch.Tensor, layer_activations: torch.Tensor, sorted_indices: torch.Tensor, + batch_first: bool): + global random_ltd_module + if random_ltd_module is None: + random_ltd_module = RandomLTDBuilder().load() + scatter_results = random_ltd_module.token_scatter_(all_activations.clone(), layer_activations, sorted_indices, + batch_first) + + ctx.save_for_backward(sorted_indices) + ctx.batch_first = batch_first + return scatter_results + + @staticmethod + def backward(ctx, out_gradients: torch.Tensor): + + out_gradients = out_gradients.contiguous() + global random_ltd_module + if random_ltd_module is None: + random_ltd_module = RandomLTDBuilder().load() + sorted_indices, = ctx.saved_tensors + batch_first = ctx.batch_first + + ret_val = random_ltd_module.token_gather(out_gradients, sorted_indices, batch_first) + return out_gradients, ret_val, None, None diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e357257869f794a06d575bfa378769f8e6d3d43c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .sparsity_config import SparsityConfig, DenseSparsityConfig, FixedSparsityConfig, VariableSparsityConfig, BigBirdSparsityConfig, BSLongformerSparsityConfig, LocalSlidingWindowSparsityConfig +from .sparse_self_attention import SparseSelfAttention +from .bert_sparse_self_attention import BertSparseSelfAttention +from .sparse_attention_utils import SparseAttentionUtils diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24c275de33737fe3c0e358918721a5dfaeae2a59 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42d6388ddbc7cb88ce610a27435d1151125c6069 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcbd172a00bd7a657c5fbe7f047ea45af54d39b4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d3815a0fc1970f72dea70e85801f8aa5bb9a67a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..500a614b1c1f7bdf19c082044b9dcca25c279cff Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39a403c907f50f3a741bc8acbefd49ecc62d5b41 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9276165446a1f96afd9643de68a9abf474c7b004 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..37f065e4863111275aae82bcd0bc1ae8513a0896 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch import nn +from deepspeed.ops.sparse_attention import SparseSelfAttention, FixedSparsityConfig + + +class BertSparseSelfAttention(nn.Module): + """Implements Sparse Self Attention layer of Bert model based on https://github.com/deepspeedai/DeepSpeedExamples/blob/master/bing_bert/nvidia/modelingpreln.py#L373 + + For more information please see, TODO DeepSpeed Sparse Transformer. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial. + """ + + def __init__( + self, + config, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=FixedSparsityConfig(num_heads=4)): + """Initialize the bert sparse self attention layer. + + Note) you can use any of the provided sparsity configs or simply add yours! + + Arguments: + config: required: Bert model config + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on FixedSparsityConfig class. + """ + + super(BertSparseSelfAttention, self).__init__() + if config.hidden_size % config.num_attention_heads != 0: + raise ValueError("The hidden size (%d) is not a multiple of the number of attention " + "heads (%d)" % (config.hidden_size, config.num_attention_heads)) + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.sparse_self_attention = SparseSelfAttention(sparsity_config) + + def transpose_for_scores(self, x): + new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) + x = x.view(*new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward(self, hidden_states, attention_mask): + """Applies forward phase of bert sparse self attention + + Arguments: + hidden_states: required: hidden_states tensor of the bert model + attn_mask: required: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported + + Return: + context_layer: a dense tensor containing attention context + """ + mixed_query_layer = self.query(hidden_states) + mixed_key_layer = self.key(hidden_states) + mixed_value_layer = self.value(hidden_states) + + query_layer = self.transpose_for_scores(mixed_query_layer) + key_layer = self.transpose_for_scores(mixed_key_layer) + value_layer = self.transpose_for_scores(mixed_value_layer) + + context_layer = self.sparse_self_attention(query_layer, + key_layer, + value_layer, + key_padding_mask=attention_mask) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size, ) + context_layer = context_layer.view(*new_context_layer_shape) + return context_layer diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/matmul.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..b30028fffbaafaf6869156406e0f8d70c8e97538 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/matmul.py @@ -0,0 +1,819 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a +# https://github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py +import importlib +import torch + +import triton +import triton.language as tl +import triton._C.libtriton as libtriton +from deepspeed.accelerator import get_accelerator + + +@triton.jit +def _kernel(A, B, C, stride_za, stride_ha, stride_ma, stride_ka, stride_zb, stride_hb, stride_kb, stride_nb, stride_zc, + stride_hc, stride_mc, stride_nc, DS0, DS1, SDD_K, SDD_off_width, lut, locks, nlocks, **meta): + TM = meta['TM'] + TN = meta['TN'] + TK = meta['TK'] + TZ = meta['TZ'] + BLOCK = meta['BLOCK'] + #------------# + #- Prologue -# + #------------# + pid0 = tl.program_id(0) + pid1 = tl.program_id(1) + pidz = tl.program_id(2) + if meta['SDD']: + pid1 = pid1 + SDD_off_width + blockidm = tl.arange(0, TM) // BLOCK + blockidn = tl.arange(0, TN) // BLOCK + offlutm = blockidm * (TN // BLOCK) * 4 + offlutn = blockidn * 4 + header = lut + pid1 * (TM // BLOCK) * (TN // BLOCK) * 4 + z = tl.load(header + 0) + i = tl.load(header + 1 + offlutm) + j = tl.load(header + 2 + offlutn) + AS1 = SDD_K // TZ + lockid = tl.where(TZ > 1, 1, 0) + offka = pid0 * AS1 + offkb = pid0 * AS1 + offmc = 0 + offnc = 0 + offpa = 0 + offpb = 0 + maxid = TZ + offhc = 0 + offha = z + offhb = z + ram = i * BLOCK + (tl.arange(0, TM) % BLOCK) + rbn = j * BLOCK + (tl.arange(0, TN) % BLOCK) + else: + header = lut + pid0 * 6 + offset = tl.load(header + 0) + AS1 = tl.load(header + 1) + column = tl.load(header + 2) + depth = tl.load(header + 3) + lockid = tl.load(header + 4) + maxid = tl.load(header + 5) + pinc = lut + offset + offhc = depth + if meta['DSD']: + # output offset + offnc = pid1 * TN + offmc = column * TM + offpc = 0 + # dense input offset + offnb = pid1 * TN + offkb = tl.load(pinc) + offkb = tl.multiple_of(offkb, 8) # compiler hint + offpb = 0 + # sparse input offset + offma = 0 + offka = 0 + offpa = tl.load(pinc + 1) + offpa = tl.multiple_of(offpa, 8) # compiler hint + offpa = offpa * BLOCK * BLOCK + offha = 0 + offhb = depth + else: + # output offset + offmc = pid1 * TM + offnc = column * TN + offpc = 0 + # dense input offset + offma = pid1 * TM + offka = tl.load(pinc) + offka = tl.multiple_of(offka, 8) # compiler hint + offpa = 0 + # sparse input offset + offnb = 0 + offkb = 0 + offpb = tl.load(pinc + 1) + offpb = tl.multiple_of(offpb, 8) # compiler hint + offpb = offpb * BLOCK * BLOCK + offha = depth + offhb = 0 + ram = offma + tl.arange(0, TM) + rbn = offnb + tl.arange(0, TN) + + # initialize a, b pointers + rka = offka + tl.arange(0, TK) + rkb = offkb + tl.arange(0, TK) + pa = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, None] * stride_ma + rka[None, :] * stride_ka + pb = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[None, :] * stride_nb + rkb[:, None] * stride_kb + if meta['DDS']: + checkam = ram[:, None] < DS0 + else: + checkam = AS1 > 0 + if meta['DSD']: + checkbn = rbn[None, :] < DS0 + else: + checkbn = AS1 > 0 + a = tl.load(pa, mask=checkam, other=0.) + b = tl.load(pb, mask=checkbn, other=0.) + + ## ---------------- ## + ## Inner Loop ## + ## ---------------- ## + acc = tl.zeros((TM, TN), dtype=tl.float32) + for k in range(AS1, 0, -TK): + acc += tl.dot(a, b) + if meta['SDD']: + inc_a = TK * stride_ka + inc_b = TK * stride_kb + else: + pinc += 2 + if meta['DSD']: + inc_b = tl.load(pinc) + inc_a = tl.load(pinc + 1) + inc_b = tl.multiple_of(inc_b, 8) + inc_a = tl.multiple_of(inc_a, 8) + inc_b = inc_b * stride_kb + if meta['DDS']: + inc_a = tl.load(pinc) + inc_b = tl.load(pinc + 1) + inc_a = tl.multiple_of(inc_a, 8) + inc_b = tl.multiple_of(inc_b, 8) + inc_a = inc_a * stride_ka + pa += inc_a + pb += inc_b + # pre-fetch + checkak = k > TK + checkbk = k > TK + checka = checkam & checkak + checkb = checkbn & checkbk + a = tl.load(pa, mask=checka) + b = tl.load(pb, mask=checkb) + c = acc.to(C.dtype.element_ty) + + if meta['SDD']: + checkc = True + rr_blockidm = tl.arange(0, TM) // BLOCK + rr_blockidn = tl.arange(0, TN) // BLOCK + rr_offlutm = rr_blockidm * (TN // BLOCK) * 4 + rr_offlutn = rr_blockidn * 4 + off_bkid = 3 + rr_offlutm[:, None] + rr_offlutn[None, :] + bkid = tl.load(header + off_bkid) + offpc = bkid * BLOCK * BLOCK + rcm = tl.arange(0, TM) % BLOCK + rcn = tl.arange(0, TN) % BLOCK + else: + rcm = offmc + tl.arange(0, TM) + rcn = offnc + tl.arange(0, TN) + if meta['DSD']: + checkc = rcn[None, :] < DS0 + if meta['DDS']: + checkc = rcm[:, None] < DS0 + + pc = C + offpc + offhc * stride_hc + pidz * stride_zc + rcm[:, None] * stride_mc + rcn[None, :] * stride_nc + # write-back directly + if lockid == 0: + tl.store(pc, c, mask=checkc) + # accumulate partial results using spin-locks + else: + plock = locks + tl.program_id(2) * nlocks * tl.num_programs(1) + tl.program_id(1) * nlocks + lockid - 1 + pcount = plock + tl.num_programs(2) * tl.num_programs(1) * nlocks + while tl.atomic_cas(plock, 0, 1) == 1: + pass + count = tl.load(pcount) + if count == 0: + tl.store(pc, c, mask=checkc) + else: + d = tl.load(pc, mask=checkc) + tl.store(pc, d + c, mask=checkc) + tl.atomic_xchg(pcount, (count + 1) % maxid) + tl.atomic_xchg(plock, 0) + + +############## +# MAIN API # +############## +class _sparse_matmul(torch.autograd.Function): + + sdd_cache = dict() + dsd_cache = dict() + dds_cache = dict() + locks = dict() + + # Given an array sizes representing reduction size for each + # column of a block-mode matrix multiplication, + # performs load-balancing to achieve more smaller reductions + # between `seg_size` elements + @staticmethod + def load_balance(sizes, block): + #global triton + #if triton is None: + # triton = importlib.import_module('triton') + # segment size + # heuristics taken from OpenAI blocksparse code + # https://github.com/openai/blocksparse/blob/master/blocksparse/matmul.py#L95 + max_size = sizes.max() + min_size = sizes[sizes != 0].min() + #if max_size > min_size * 2.0: + # seg_max = max(triton.cdiv(max_size, 4), min_size*2) + #else: + # seg_max = max_size + seg_max = max_size + seg_min = max(triton.cdiv(seg_max, 4), 4) + # split reduction into segments + div = sizes // seg_max + rem = sizes % seg_max + packs = div + (sizes < seg_min).long() + (rem >= seg_min).long() + width = packs.sum() + segments = torch.empty(width, dtype=sizes.dtype) + column = torch.empty_like(segments) + lockid = torch.zeros_like(segments) + maxid = torch.zeros_like(segments) + nlocks = 0 + current = 0 + col_idx = 0 + for i in range(len(sizes)): + d, r = div[i], rem[i] + isempty = sizes[i] < seg_min + last = current + d + (r >= seg_min) + isempty + # column id + column[current:last] = col_idx + # lock id + if d > 1 or (d == 1 and r >= seg_min): + nlocks += 1 + lockid[current:last] = nlocks + maxid[current:last] = last - current + # segment size + segments[current:current + d] = seg_max + if r < seg_min and not isempty: + segments[current + d - 1] += r + if r >= seg_min or isempty: + segments[current + d] = r + current = last + col_idx += 1 + offsets = torch.zeros_like(segments) + offsets[1:] = torch.cumsum(segments[:-1], dim=0) + return segments, column, lockid, maxid, offsets + + @staticmethod + def get_locks(size, dev): + if dev not in _sparse_matmul.locks or \ + size > _sparse_matmul.locks[dev].size(0): + _sparse_matmul.locks[dev] = torch.zeros(size, dtype=torch.int32, device=dev) + return _sparse_matmul.locks[dev] + + ########################## + # SPARSE = DENSE x DENSE # + ########################## + + @staticmethod + def make_sdd_lut(layout, block, dtype, device): + #_sparse_matmul._load_utils() + #start_width = 64 // block + #segmented = _sparse_matmul.sdd_segment(layout.type(torch.int32), start_width) + start_width = (128 if block > 16 else 32) // block + layout = layout.type(torch.int32) + segmented = libtriton.superblock(layout.data_ptr(), layout.shape[0], layout.shape[1], layout.shape[2], + start_width) + luts, widths, packs = [], [], [] + for size, nnz in segmented: + """ width = nnz.shape[0] // (size * size) + h = nnz[:, 0] + i = nnz[:, 1] + j = nnz[:, 2] + b = nnz[:, 3] + lut = torch.stack((h, i, j, b), dim=1).view(-1).contiguous() + luts.append(lut.type(torch.int32).to(device)) + widths.append(width) + packs.append(size) """ + nnz = nnz.reshape(-1, 4) + width = nnz.shape[0] // (size * size) + luts.append(torch.from_numpy(nnz).type(torch.int32).to(device)) + widths.append(width) + packs.append(size) + # create locks + return luts, None, widths, packs + + @staticmethod + def _sdd_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, luts, num_locks, widths, packs, bench, time): + if trans_c: + a, b = b, a + trans_a, trans_b = not trans_b, not trans_a + AS0 = a.size(0) + # Shape check + a_dim = -2 if trans_a else -1 + b_dim = -1 if trans_b else -2 + a_inner, b_inner = a.shape[a_dim], b.shape[b_dim] + if a_inner != b_inner: + raise ValueError(f"Size of tensor A along the {a_dim} dim ({a_inner}) must match size " + f"of tensor B along the {b_dim} dim ({b_inner})") + if a_inner % 16 != 0: + raise ValueError('Reduction size for SDD must be a multiple of 16') + + batch_size = a.size(0) + a_outer = a.size(3 if trans_a else 2) + dtype = a.dtype + is_16_multiple = a_inner % 16 == 0 + is_32_multiple = a_inner % 32 == 0 + is_64_multiple = a_inner % 64 == 0 + if not is_16_multiple: + raise ValueError('Reduction size for SDD must be a multiple of 16') + device = a.device + # create kernel + total_width = sum([width * pack * pack for width, pack in zip(widths, packs)]) + c = torch.empty((batch_size, total_width, block, block), dtype=dtype, device=a.device) + for lut, width, pack in zip(luts, widths, packs): + F32TK = [8, 16] + F16TK = [16] + F16TK += [32] if is_32_multiple else [] + F16TK += [64] if is_64_multiple else [] + TK = {torch.float32: F32TK, torch.float16: F16TK}[dtype] + num_lock = 1 + meta = { + 'TM': block * pack, + 'TN': block * pack, + 'BLOCK': block, + 'TK': TK[0], + 'TZ': 1, + 'SDD': True, + 'DSD': False, + 'DDS': False + } + # create output + locks = _sparse_matmul.get_locks(2 * width * AS0 * num_lock, a.device) + # maximum grid size is 65535 + # so operation might be decomposed into multiple + # kernel calls + max_width = 49152 + total = 0 if bench else None + for off_width in range(0, width, max_width): + grid = lambda meta: [meta['TZ'], min(max_width, width - off_width), batch_size] + _kernel[grid](a, + b, + c, + a.stride(0), + a.stride(1), + a.stride(3 if trans_a else 2), + a.stride(2 if trans_a else 3), + b.stride(0), + b.stride(1), + b.stride(3 if trans_b else 2), + b.stride(2 if trans_b else 3), + c.stride(0), + c.stride(0), + c.stride(2), + c.stride(3), + a_outer, + a_outer, + a_inner, + off_width, + lut, + locks, + num_lock, + num_warps=4, + **meta) + # save for backward pass + return c + + ########################## + # DENSE = DENSE x SPARSE # + ########################## + + # Given a binary layout of 0s and 1s, + # Construct look-up table for efficient execution on GPUs + @staticmethod + def make_dxx_lut(layout, block, step, trans, device, transform=lambda idx: idx): + # load-balancing + _empty = torch.tensor([], dtype=torch.int64, device=layout.device) + segments = _empty.clone() + column = _empty.clone() + depth = _empty.clone() + lockid = _empty.clone() + maxid = _empty.clone() + offsets = _empty.clone() + current_offset = 0 + current_maxid = 0 + for z in range(layout.size(0)): + if trans: + sizes = torch.sum(layout[z, :, :], 1) + else: + sizes = torch.sum(layout[z, :, :], 0) + z_segments, z_column, z_lockid, z_maxid, z_offsets = _sparse_matmul.load_balance(sizes, block) + z_depth = z * torch.ones_like(z_segments) + z_lockid[z_lockid > 0] += current_maxid + current_maxid = z_lockid.max() + # concatenate depth + segments = torch.cat((segments, z_segments)) + column = torch.cat((column, z_column)) + depth = torch.cat((depth, z_depth)) + maxid = torch.cat((maxid, z_maxid)) + offsets = torch.cat((offsets, current_offset + z_offsets)) + lockid = torch.cat((lockid, z_lockid)) + current_offset += layout[z, :, :].sum() + segments *= step + # pointer increments + if trans: + nnz = layout.nonzero() + else: + nnz = layout.transpose(1, 2).nonzero() + num_blocks = nnz.size(0) + offsets = torch.min(offsets, (num_blocks - 1) * torch.ones_like(offsets)) + idx = transform(nnz[:, 2] * block) + xincs = idx.clone() + xincs[1:] -= idx[:-1] + # divide block into multiple steps + div = block // step + xincs = xincs.view(-1, 1).repeat(1, div) + xincs[:, 1:] = step + xincs[:, 0] -= (div - 1) * step + # first increment for each reduction is actually the offset + xincs[offsets[segments > 0], 0] = idx[offsets[segments > 0]] + xincs = xincs.view(-1) + # block-mode input increments + if trans: + widx = torch.arange(num_blocks) + else: + widx = _empty.clone() + current_offset = 0 + for z in range(layout.size(0)): + layoutw = layout[z, :, :].clone() + msum = layoutw.sum() + layoutw[layoutw > 0] = 1 + torch.arange(msum) + widx = torch.cat((widx, current_offset + layoutw.T[layoutw.T > 0] - 1)) + current_offset += msum + widx = widx + wincs = widx * block * block + wincs[1:] -= widx[:-1] * block * block + wincs = wincs.view(-1, 1).repeat(1, div) + if trans: + wincs[:, 1:] = step + wincs[:, 0] -= (div - 1) * step + else: + wincs[:, 1:] = step * block + wincs[:, 0] -= (div - 1) * step * block + wincs[offsets[segments > 0], 0] = widx[offsets[segments > 0]] + wincs = wincs.view(-1) + # adjust offset and segment size + offsets *= 2 * div + segments *= div + # create header + width = column.size(0) + offsets += 6 * width + header = torch.stack((offsets, segments, column, depth, lockid, maxid), dim=1).view(-1).contiguous() + incs = torch.stack((xincs, wincs), dim=1).view(-1).contiguous() + incs = torch.cat((incs, torch.zeros(2, device=incs.device, dtype=incs.dtype))) + # create lut + lut = torch.cat((header, incs)) + lut = lut.type(torch.int32).to(device) + # create locks + num_locks = max(1, lockid.max()) + return lut, num_locks, width, None + + @staticmethod + def _dds_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, lut, num_locks, width, packs, bench, time): + global triton + if triton is None: + triton = importlib.import_module('triton') + + # shapes / dtypes + AS0 = a.size(0) + AS1 = a.size(1) + AS2 = a.size(3 if trans_a else 2) + AS3 = a.size(2 if trans_a else 3) + BS0 = spdims[0] + BS1 = block * spdims[2 if trans_b else 1] + BS2 = block * spdims[1 if trans_b else 2] + dtype = a.dtype + # kernel + meta = {'TN': block, 'TM': 128, 'TK': 16, 'BLOCK': block, 'TZ': 1, 'SDD': False, 'DSD': False, 'DDS': True} + # output + CS0 = AS0 + CS1 = AS1 + CS2 = BS2 if trans_c else AS2 + CS3 = AS2 if trans_c else BS2 + locks = _sparse_matmul.get_locks(2 * AS0 * AS2 // 32 * num_locks, a.device) + c = torch.empty((CS0, CS1, CS2, CS3), dtype=dtype, device=a.device) + grid = lambda meta: [width, triton.cdiv(AS2, meta['TM']), AS0] + _kernel[grid](a, + b, + c, + a.stride(0), + a.stride(1), + a.stride(3 if trans_a else 2), + a.stride(2 if trans_a else 3), + b.stride(0), + b.stride(1), + b.stride(3 if trans_b else 2), + b.stride(2 if trans_b else 3), + c.stride(0), + c.stride(1), + c.stride(3 if trans_c else 2), + c.stride(2 if trans_c else 3), + AS2, + BS2, + 0, + 0, + lut, + locks, + num_locks, + num_warps=4, + **meta) + return c + + @staticmethod + def _dsd_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, lut, num_locks, width, packs, bench, time): + global triton + if triton is None: + triton = importlib.import_module('triton') + + # shapes / dtypes + AS0 = spdims[0] + AS1 = block * spdims[2 if trans_a else 1] + AS2 = block * spdims[1 if trans_a else 2] + BS0 = b.size(0) + BS1 = b.size(1) + BS2 = b.size(3 if trans_b else 2) + BS3 = b.size(2 if trans_b else 3) + dtype = a.dtype + # kernel + + meta = {'TM': block, 'TN': 128, 'TK': 16, 'BLOCK': block, 'TZ': 1, 'SDD': False, 'DSD': True, 'DDS': False} + # output + CS0 = BS0 + CS1 = BS1 + CS2 = BS3 if trans_c else AS1 + CS3 = AS1 if trans_c else BS3 + locks = _sparse_matmul.get_locks(2 * BS0 * BS3 // 32 * num_locks, a.device) + c = torch.empty((CS0, CS1, CS2, CS3), dtype=dtype, device=a.device) + grid = lambda meta: [width, triton.cdiv(BS3, meta['TN']), BS0] + _kernel[grid](a, + b, + c, + a.stride(0), + a.stride(1), + a.stride(3 if trans_a else 2), + a.stride(2 if trans_a else 3), + b.stride(0), + b.stride(1), + b.stride(3 if trans_b else 2), + b.stride(2 if trans_b else 3), + c.stride(0), + c.stride(1), + c.stride(2), + c.stride(3), + BS3, + AS1, + 0, + 0, + lut, + locks, + num_locks, + num_warps=4, + **meta) + return c + + fn = {'sdd': _sdd_matmul.__get__(object), 'dsd': _dsd_matmul.__get__(object), 'dds': _dds_matmul.__get__(object)} + + @staticmethod + def forward(ctx, a, b, trans_a, trans_b, trans_c, mode, spdims, block, c_lut, c_num_locks, c_width, c_packs, + c_bench, c_time, da_lut, da_num_locks, da_width, da_packs, da_bench, da_time, db_lut, db_num_locks, + db_width, db_packs, db_bench, db_time): + c = _sparse_matmul.fn[mode](a, b, trans_a, trans_b, trans_c, spdims, block, c_lut, c_num_locks, c_width, + c_packs, c_bench, c_time) + # save for backward + ctx.save_for_backward(a, b) + ctx.da_num_locks = da_num_locks + ctx.da_lut = da_lut + ctx.da_width = da_width + ctx.da_packs = da_packs + ctx.da_bench = da_bench + ctx.da_time = da_time + ctx.db_lut = db_lut + ctx.db_num_locks = db_num_locks + ctx.db_width = db_width + ctx.db_bench = db_bench + ctx.db_packs = db_packs + ctx.db_time = db_time + ctx.mode = mode + ctx.spdims = spdims + ctx.block = block + ctx.trans_a = trans_a + ctx.trans_b = trans_b + return c + + @staticmethod + def backward(ctx, dc): + # saved for backward + a, b = ctx.saved_tensors + mode = ctx.mode + # gradients w.r.t. a + if ctx.needs_input_grad[0]: + mode_da = mode[1] + mode[0] + mode[2] + da = _sparse_matmul.fn[mode_da](dc, b, False, not ctx.trans_b, ctx.trans_a, ctx.spdims, ctx.block, + ctx.da_lut, ctx.da_num_locks, ctx.da_width, ctx.da_packs, ctx.da_bench, + ctx.da_time) + # gradients w.r.t. b + if ctx.needs_input_grad[1]: + mode_db = mode[2] + mode[1] + mode[0] + db = _sparse_matmul.fn[mode_db](a, dc, not ctx.trans_a, False, ctx.trans_b, ctx.spdims, ctx.block, + ctx.db_lut, ctx.db_num_locks, ctx.db_width, ctx.db_packs, ctx.db_bench, + ctx.db_time) + return da, db, None, None, None,\ + None, None, None, None,\ + None, None, None, None, None, None,\ + None, None, None, None, None, None,\ + None, None, None, None, None, None + + +class MatMul: + """Block-Sparse MatMul class; this class handles three types of matrix-multiplication: + - sparse = dense X dense + - dense = sparse X dense + - dense = dense X sparse + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + """ + + def make_lut(self, dtype, device): + """Generates the sparsity layout/s used in block-sparse matmul + """ + key = (dtype, device) + if key in self.lut_cache: + return self.lut_cache[key] + # C look-up table + layout, block = self.layout, self.block + step = 16 + if self.mode == 'sdd': + c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) + elif self.mode == 'dsd': + c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_dxx_lut(layout, block, step, not self.trans_a, + device) + elif self.mode == 'dds': + c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_dxx_lut(layout, block, step, self.trans_b, + device) + # DA look-up table + if self.mode == 'sdd': + da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, True, device) + elif self.mode == 'dsd': + da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) + elif self.mode == 'dds': + da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, + not self.trans_b, device) + # DB look-up table + if self.mode == 'sdd': + db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_dxx_lut(layout, block, step, False, device) + elif self.mode == 'dsd': + db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_dxx_lut(layout, block, step, self.trans_a, + device) + elif self.mode == 'dds': + db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) + self.lut_cache[key] = (c_lut, c_num_locks, c_width, c_packs,\ + da_lut, da_num_locks, da_width, da_packs,\ + db_lut, db_num_locks, db_width, db_packs) + return self.lut_cache[key] + + def __init__(self, layout, block, mode, trans_a=False, trans_b=False, bench=False): + """Initialize the Block-Sparse MatMul class. + + Arguments: + layout: required: sparsity layout tensor + block: required: an integer determining the block size. + mode: required: a string determining type of matmul; ('sdd') sparse = dense X dense, ('dsd') dense = sparse X dense, ('dds') dense = dense X sparse + trans_a: optional: a boolean determining if multiplication needs to be applied on transpose of input a; default is false + trans_b: optional: a boolean determining if multiplication needs to be applied on transpose of input b; default is false + bench: optional: set if you want to do benchmarking + """ + + if mode not in ['sdd', 'dsd', 'dds']: + raise NotImplementedError('Supported modes are: sdd, dsd, dds') + # look-up table cache + self.lut_cache = dict() + # attributes + self.trans_a = trans_a + self.trans_b = trans_b + self.mode = mode + self.block = block + self.layout = layout + layout_dim = layout.ndim + assert layout_dim in (2, 3), "Layout should be a 2 or 3 dimensional tensor of 0s and 1s" + if not mode == 'sdd': + # Dims to be reduced on the 'inside' of the matmul, either -1 or -2 + trans_dense, trans_sparse, sparse_inner = (trans_b, trans_a, -1) if mode == 'dsd' else (trans_a, trans_b, + -2) + self.dense_inner_dim = -((sparse_inner % 2) + 1) if not trans_dense else sparse_inner + sparse_inner = sparse_inner if not trans_sparse else -((sparse_inner % 2) + 1) + + # Inner dim of the dense input should be equal to the inner dim of the sparse input + self.dense_inner_size = layout.shape[sparse_inner] * block + # Expected shape for sparse inputs + self.sparse_shape = (layout.sum().item(), block, block) + + # Support using the same layout across attention heads etc. + if layout_dim == 2: + layout = layout.unsqueeze(0) + + layout = layout.long() # Above code assumes the layout tensor is an integral type + + self.spdims = layout.shape + # timings + self.bench = bench + self.time_c = None + self.time_da = None + self.time_db = None + + # pad shapes of a tensor to make it + # compatible with kernel calls + @staticmethod + def _pad_shape(x, is_sparse): + max_dim = 3 if is_sparse else 4 + for i in range(max_dim - x.dim()): + x = x.unsqueeze(0) + return x + + def __call__(self, a, b): + """Applies Block-Sparse MatMul. + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + + Arguments: + a: required: a dense/block-sparse tensor; first input of mat-mul + b: required: a dense/block-sparse tensor; second input of mat-mul + + Return: + c: a dense/block-sparse tensor result of a X b + """ + + + c_lut, c_num_locks, c_width, c_packs,\ + da_lut, da_num_locks, da_width, da_packs,\ + db_lut, db_num_locks, db_width, db_packs = self.make_lut(a.dtype, a.device) + # timings + time_c = [None] + time_da = [None] + time_db = [None] + + original_dims = max(a.ndim, b.ndim) + a, b = self._validate_inputs(a, b) + + # pad shapes with ones + a = MatMul._pad_shape(a, self.mode == 'dsd') + b = MatMul._pad_shape(b, self.mode == 'dds') + # execute + + c = _sparse_matmul.apply(a, b, self.trans_a, self.trans_b, False, self.mode, self.spdims, self.block, c_lut, + c_num_locks, c_width, c_packs, self.bench, time_c, da_lut, da_num_locks, da_width, + da_packs, self.bench, time_da, db_lut, db_num_locks, db_width, db_packs, self.bench, + time_db) + + # This removes any leading singleton dimensions we may have added to the tensor that weren't in the input + dims_to_trim = c.ndim - original_dims + for _ in range(dims_to_trim): + c = c.squeeze(0) + + self.time_c = time_c[0] + self.time_da = time_da[0] + self.time_db = time_db[0] + return c + + def _validate_inputs(self, a, b): + if a.device != b.device: + raise ValueError(f"Inputs must be on the same device; got {a.device} for tensor A " + f"and {b.device} for tensor B") + if not get_accelerator().on_accelerator(a): + raise ValueError("Only GPU devices are supported for now") + + # When autocast is enabled, torch.matmul autocasts to float16, so we do the same here + if torch.is_autocast_enabled(): + a, b = a.half(), b.half() + elif a.dtype != b.dtype: + raise ValueError(f"Inputs must be the same dtype; got {a.dtype} for A and {b.dtype} for B") + + mode, trans_a, trans_b = self.mode, self.trans_a, self.trans_b + if mode != 'sdd': + # One input is sparse + dense, dense_name, sparse, sparse_name = (a, 'A', b, 'B') if mode == 'dds' else (b, 'B', a, 'A') + dense_inner = dense.shape[self.dense_inner_dim] + if dense_inner != self.dense_inner_size: + raise ValueError(f"Expected tensor {dense_name} to have size {self.dense_inner_size} at dim " + f"{self.dense_inner_dim % dense.ndim}, got {dense_inner}.") + + if sparse.shape[-len(self.sparse_shape):] != self.sparse_shape: + raise ValueError(f"Expected tensor with trailing dimensions of shape {self.sparse_shape} for argument " + f"{sparse_name}, got {sparse.shape}") + + def add_extra_dims(x): + # Add extra leading singleton dimensions if needed + dims_needed = 4 - x.ndim + if dims_needed > 0: + singletons = [1] * dims_needed + x = x.view(*singletons, *x.shape) + elif dims_needed < 0: + raise ValueError("Tensors with more than 4 dimensions are not currently supported") + + return x + + # Pad shapes with leading singleton dimensions + a = add_extra_dims(a) + b = add_extra_dims(b) + + return a, b diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/softmax.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/softmax.py new file mode 100644 index 0000000000000000000000000000000000000000..debee5688fe3f06699a710dbaa35c09601189a4f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/softmax.py @@ -0,0 +1,296 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a +# https://github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py + +import torch + +import triton +import triton.language as tl + + +def next_power_of_2(n): + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n += 1 + return n + + +def num_warps(n): + if n < 512: + return 4 + if n < 2048: + return 8 + return 16 + + +@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[6] * meta['BLOCK'])}) +@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[6] * meta['BLOCK'])}) +@triton.jit +def _forward(X, scale, LUT, RPE, KP_M, ATTN_M, sizemax, stride_zx, stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, + stride_zattnm, **meta): + TN = meta['TN'] + BLOCK = meta['BLOCK'] + pidhm = tl.program_id(0) + pidz = tl.program_id(1) + # create index ranges + rxm = pidhm % BLOCK + rbm = pidhm // BLOCK + rxn = tl.arange(0, TN) % BLOCK + rbn = tl.arange(0, TN) // BLOCK + # extract information from LUT + header = LUT + rbm * 2 + size = tl.load(header + 0) + offset = tl.load(header + 1) + check = rbn < size + rbmn = tl.where(check, rbn, size - 1) + # block id and column id + blockid = tl.load(LUT + offset + rbmn * 4 + 0) + columnid = tl.load(LUT + offset + rbmn * 4 + 1) + rowid = tl.load(LUT + offset + rbmn * 4 + 2) + headid = tl.load(LUT + offset + rbmn * 4 + 3) + # pointers to X + px = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn + x = tl.load(px, mask=check, other=-float('inf')) + x = x.to(tl.float32) + # apply scale + if meta['APPLY_SCALE']: + x = x * scale + # apply RPE + if meta['APPLY_RPE']: + prpe = RPE + pidz * stride_zrpe + headid * stride_hrpe + columnid * BLOCK + rowid * BLOCK * stride_srpe + rxm * stride_srpe + rxn + rpe = tl.load(prpe, mask=check, other=0) + x = x + rpe + # apply key-padding mask + if meta['APPLY_KP_MASK']: + pkp_m = KP_M + pidz * stride_zkpm + columnid * BLOCK + rxn + kp_m = tl.load(pkp_m, mask=check, other=-float('inf')) + if meta['KP_MASK_MUL']: + kp_m = tl.where(kp_m == 0, -float('inf'), 0.) + x = x + kp_m + # apply attention mask + if meta['APPLY_ATTN_MASK']: + pattn_m = ATTN_M + columnid * BLOCK + rowid * BLOCK * stride_zattnm + rxm * stride_zattnm + rxn + attn_m = tl.load(pattn_m, mask=check, other=-float('inf')) + if meta['ATTN_MASK_MUL']: + attn_m = tl.where(attn_m == 0, -float('inf'), 0.) + x = x + attn_m + # computation + x = tl.softmax(x) + tl.store(px, x, mask=check) + + +@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[4] * meta['BLOCK'])}) +@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[4]) * meta['BLOCK']}) +@triton.jit +def _backward(X, scale, DX, LUT, sizemax, stride_zx, stride_zdx, **meta): + pidhm = tl.program_id(0) + pidz = tl.program_id(1) + TN = meta['TN'] + BLOCK = meta['BLOCK'] + # create index ranges + rxm = pidhm % BLOCK + rbm = pidhm // BLOCK + rxn = tl.arange(0, TN) % BLOCK + rbn = tl.arange(0, TN) // BLOCK + # extract information from look-up table + header = LUT + rbm * 2 + size = tl.load(header + 0) + offset = tl.load(header + 1) + # bounds checking on lut + check = rbn < size + rbmn = tl.where(check, rbn, size - 1) + # initialize pointers to block-sparse input + blockid = tl.load(LUT + offset + rbmn * 4) + X = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn + DX = DX + pidz * stride_zdx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn + # compute fused softmax backward + x = tl.load(X, mask=check, other=0) + dx = tl.load(DX, mask=check, other=0) + x = x.to(tl.float32) + dx = dx.to(tl.float32) + y = x * (dx - tl.sum(x * dx, 0)) * scale + tl.store(DX, y, mask=check) + + +class _sparse_softmax(torch.autograd.Function): + + bwd_kernels = dict() + + @staticmethod + def make_lut(layout, block, device): + _empty = torch.tensor([], dtype=torch.int64, device=layout.device) + sizes = _empty.clone() + # sizes along rows + for h in range(layout.shape[0]): + sizes = torch.cat((sizes, layout[h, :, :].sum(-1))) + # offsets in block format + offsets = torch.zeros_like(sizes) + offsets[1:] = torch.cumsum(sizes[:-1], dim=0) + # block indices + idx = torch.arange(layout.sum()) + head = layout.nonzero()[:, 0] + rows = layout.nonzero()[:, 1] + columns = layout.nonzero()[:, 2] + core = torch.stack((idx, columns, rows, head), dim=1).view(-1) + # construct look-up table + offsets = offsets * 4 + 2 * sizes.numel() + header = torch.stack((sizes, offsets), dim=1).view(-1) + lut = torch.cat((header, core)).type(torch.int32).to(device) + return lut, int(sizes.max()) + + @staticmethod + def forward(ctx, x, scale, rpe, key_padding_mask, attn_mask, kp_mask_mode, attn_mask_mode, spdims, block, lut, + num_blocks, maxlut, bench, time): + + apply_scale = False if scale == 1.0 else True + + # handle None rpe + if rpe is None: + apply_rpe = False + stride_zrpe, stride_hrpe, stride_srpe = 0, 0, 0 + rpe = torch.empty(0, dtype=x.dtype, device=x.device) + else: + apply_rpe = True + stride_zrpe, stride_hrpe, stride_srpe = rpe.stride(0), rpe.stride(1), rpe.stride(2) + + # handle None key_padding_mask + if key_padding_mask is None: + apply_kp_mask = False + stride_zkpm = 0 + key_padding_mask = torch.empty(0, dtype=x.dtype, device=x.device) + else: + apply_kp_mask = True + stride_zkpm = key_padding_mask.stride(0) + + # handle None attention_mask + if attn_mask is None: + apply_attn_mask = False + stride_zattnm = 0 + attn_mask = torch.empty(0, dtype=x.dtype, device=x.device) + else: + apply_attn_mask = True + stride_zattnm = attn_mask.stride(0) + + # run kernel + M = x.shape[0] + meta = { + 'BLOCK': block, + 'APPLY_SCALE': apply_scale, + 'APPLY_RPE': apply_rpe, + 'APPLY_KP_MASK': apply_kp_mask, + 'APPLY_ATTN_MASK': apply_attn_mask, + 'KP_MASK_MUL': kp_mask_mode == 'mul', + 'ATTN_MASK_MUL': attn_mask_mode == 'mul', + } + grid = lambda opt: [spdims[0] * spdims[1] * block, M] + _forward[grid](x, scale, lut, rpe, key_padding_mask, attn_mask, maxlut, x.stride(0),\ + stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, stride_zattnm, **meta) + + # save to context + ctx.mark_dirty(x) + ctx.save_for_backward(x, lut) + ctx.spdims = spdims + ctx.block = block + ctx.maxlut = maxlut + ctx.scale = scale + ctx.apply_scale = apply_scale + ctx.apply_rpe = apply_rpe + ctx.apply_kp_mask = apply_kp_mask + ctx.apply_attn_mask = apply_attn_mask + ctx.kp_mask_mode = kp_mask_mode + ctx.attn_mask_mode = attn_mask_mode + return x + + @staticmethod + def backward(ctx, dx): + + # retrieve from context + x, lut = ctx.saved_tensors + # run kernel + M = x.shape[0] + grid = lambda opt: [ctx.spdims[0] * ctx.spdims[1] * ctx.block, M] + _backward[grid](x, ctx.scale, dx, lut, ctx.maxlut, x.stride(0), dx.stride(0), BLOCK=ctx.block) + return dx, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class Softmax: + """Block-Sparse Softmax class; this class computes softmax on a block sparse matrix. It is also able to apply either/all of the following masks: + - relative position embedding + - key padding mask + - attention mask + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + """ + + def sparse_softmax(*args, **kwargs): + return _sparse_softmax.apply(*args, **kwargs) + + def make_lut(self, device): + """Generates the sparsity layout used in block-sparse softmax + """ + key = (device, ) + if key not in self.lut_cache: + self.lut_cache[key] = _sparse_softmax.make_lut(self.layout, self.block, device) + return self.lut_cache[key] + + def __init__(self, layout, block, bench=False): + """Initialize the Block-Sparse Softmax class. + + Arguments: + layout: required: sparsity layout tensor + block: required: an integer determining the block size. + bench: optional: set if you want to do benchmarking + """ + + self.num_blocks = layout.sum().item() + self.spdims = layout.shape + self.layout = layout + self.block = block + self.bench = bench + self.lut_cache = dict() + + def __call__(self, + x, + scale=1., + rpe=None, + key_padding_mask=None, + attn_mask=None, + key_padding_mask_mode='add', + attn_mask_mode='add'): + """Applies softmax on a Block-Sparse input tensor. + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + + Arguments: + x: required: a block-sparse tensor that softmax is applied on it; computation will be in place and result will be returned in the same tensor + scale: optional: a float value; x values will be multiplied by this value before normalization. Default value is 1.0. + rpe: optional: a tensor same dimension as x that is used as relative position embedding + key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength) + attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported + key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied + attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied + + Return: + x: a block-sparse tensor contains normalized input x using softmax; and masks applied if given + """ + + time_y = [None] + if rpe is not None and rpe.dtype != x.dtype: + raise ValueError('relative position embedding must be %s' % x.dtype) + if attn_mask is not None and attn_mask.dtype != x.dtype: + raise ValueError('Attention mask must be %s' % x.dtype) + if key_padding_mask is not None and key_padding_mask.dtype != x.dtype: + raise ValueError('Key padding mask must be %s' % x.dtype) + lut, maxlut = self.make_lut(x.device) + x = Softmax.sparse_softmax(x, scale, rpe, key_padding_mask, attn_mask, key_padding_mask_mode, attn_mask_mode, + self.spdims, self.block, lut, self.num_blocks, maxlut, self.bench, time_y) + self.time_y = time_y[0] + return x diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ccb0f940dff65839beac579f81c4dfb7e499e6bb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py @@ -0,0 +1,208 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from torch.nn import functional as F +from deepspeed.ops.sparse_attention import BertSparseSelfAttention, SparsityConfig +''' +This file contains few utility functions to handle adapting pretrained model with sparse self-attention module. +''' + + +class SparseAttentionUtils: + """This class provides some utility functions that are use integrating sparse attention into transformer models. + Such utilities include extending position embeddings, replacing current self-attention layer with sparse attention, padding sequences to multiple of block size, etc. + + """ + + @staticmethod + def extend_position_embedding(model, max_position): + """This function extends the position embedding weights of a model loaded from a checkpoint. + It assumes the new max position is bigger than the original max length. + + Arguments: + model: required: a transformer model + max_position: required: an integer determining new position embedding size + Return: + model: updated model; in which position embedding weights have been extended based on new size + """ + + if hasattr(model, 'bert'): + original_max_position = model.bert.embeddings.position_embeddings.weight.size(0) + assert max_position > original_max_position + extend_multiples = max(1, max_position // original_max_position) + model.bert.embeddings.position_embeddings.weight.data = model.bert.embeddings.position_embeddings.weight.repeat( + extend_multiples, 1) + elif hasattr(model, 'roberta'): + # RoBERTa has positions 0 & 1 reserved, so embedding size is max position + 2 + original_max_position, embed_size = model.roberta.embeddings.position_embeddings.weight.shape + original_max_position -= 2 + extend_multiples = max(1, max_position // original_max_position) + assert max_position > original_max_position + max_position += 2 + extended_position_embedding = model.roberta.embeddings.position_embeddings.weight.new_empty( + max_position, embed_size) + k = 2 + for i in range(extend_multiples): + extended_position_embedding[k:( + k + original_max_position)] = model.roberta.embeddings.position_embeddings.weight[2:] + k += original_max_position + model.roberta.embeddings.position_embeddings.weight.data = extended_position_embedding + else: + raise ValueError( + 'Please extend \"extend_position_embedding\" function to support your model type. It currently only supports \"bert\" & \"roberta\"!' + ) + + model.config.max_position_embeddings = max_position + print(f'Extended position embeddings to {original_max_position * extend_multiples}') + + return model + + @staticmethod + def update_tokenizer_model_max_length(tokenizer, max_position): + """This function updates the position embedding length of a tokenizer to a new max position. + + Arguments: + tokenizer: required: a transformer tokenizer + max_position: required: an integer determining new position embedding size + Return: + tokenizer: updated tokenizer; in which model maximum length has been extended based on new size + """ + + tokenizer.model_max_length = max_position + tokenizer.init_kwargs['model_max_length'] = max_position + print(f'updated tokenizer model max imum length to {max_position}') + + return tokenizer + + @staticmethod + def replace_model_self_attention_with_sparse_self_attention( + model, + max_position, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=SparsityConfig(num_heads=4)): + """This function replaces the self attention layers in model encoder with sparse self attention. + It currently supports bert and roberta model and can be easily extended to any other models following similar steps here. + For sparsityConfig, refer to the config class. + + Arguments: + model: required: a transformer model + max_position: required: an integer determining new position embedding size + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class + + Return: + model: updated model; in which self attention layer has been replaced with DeepSpeed Sparse Self Attention layer. + """ + + if hasattr(model, 'bert'): + model.config.max_position_embeddings = max_position + model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, model.bert.encoder.layer, + sparsity_config) + elif hasattr(model, 'roberta'): + model.config.max_position_embeddings = max_position + 2 + model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, + model.roberta.encoder.layer, + sparsity_config) + else: + raise ValueError( + 'Please extend \"update_model_self_attention_to_sparse_self_attention\" function to support \ + your model type. It currently only supports \"bert\" & \"roberta\"!') + return model + + @staticmethod + def replace_self_attention_layer_with_sparse_self_attention_layer( + config, + layers, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=SparsityConfig(num_heads=4)): + """This function replaces the self attention layers in attention layer with sparse self attention. + For sparsityConfig, refer to the config class. + + Arguments: + config: required: transformer model config + layers: required: transformer model attention layers + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class + + Return: + layers: updated attention layers; in which self attention layers have been replaced with DeepSpeed Sparse Self Attention layer. + """ + + for layer in layers: + deepspeed_sparse_self_attn = BertSparseSelfAttention(config, sparsity_config) + deepspeed_sparse_self_attn.query = layer.attention.self.query + deepspeed_sparse_self_attn.key = layer.attention.self.key + deepspeed_sparse_self_attn.value = layer.attention.self.value + + layer.attention.self = deepspeed_sparse_self_attn + + return layers + + @staticmethod + def pad_to_block_size(block_size, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds, + pad_token_id, model_embeddings): + """This function pads input tokens and attention mask on sequence length dimension to be multiple of block size. + This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size. + It needs to be called in your model, such as BertModel, right before you calculate the embedding outputs. + Note) + 1- instead of passing your embedding layer to this function, you can simply add this function to your model. It can be more simplified if given attention_mask and/or token_type_ids are none. + 2- you need to call unpad function before returning your model output to unpad the encoder sequence output. + + Arguments: + block_size: required: an integer determining the block size of sparsity config. + pad_token_id: required: an integer determining the pad token from the model config; such as bert.config.pad_token_id. + input_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary + attention_mask: a torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. + token_type_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). + position_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the indices of positions of each input sequence tokens in the position embeddings. + inputs_embeds: an optional torch.FloatTensor of shape [batch_size, sequence_length, hidden_size] that contains embedded representation and can be passed instead of input_ids directly. + model_embeddings: an optional object. If inputs_embeds are not none, this will be your model embeddings such as BertEmbeddings from your model such as BertModel. You can move this function inside your model and use self.embeddings instead of passing this parameter. + + Return: + pad_len: an integer determining how much inputs have been padded to transfer sequence length dimension to multiple of block size. + input_ids: if input_ids are not none padded input_ids otherwise none. + attention_mask: if attention_mask is not none padded attention_mask otherwise none. + token_type_ids: if token_type_ids are not none padded token_type_ids otherwise none. + position_ids: if position_ids are not none padded position_ids otherwise none. + inputs_embeds: if inputs_embeds are not none padded inputs_embeds otherwise none. + """ + + batch_size, seq_len = input_ids.shape if input_ids is not None else inputs_embeds.shape[:-1] + + pad_len = (block_size - seq_len % block_size) % block_size + if pad_len > 0: + if inputs_embeds is not None: + pad_input_ids = inputs_embeds.new_full((batch_size, pad_len), pad_token_id, dtype=torch.long) + pad_inputs_embeds = model_embeddings(pad_input_ids) + inputs_embeds = torch.cat([inputs_embeds, pad_inputs_embeds], dim=-2) + # may not be needed as input_ids are not used if inputs_embeds are given + if input_ids is not None: + input_ids = F.pad(input_ids, (0, pad_len), value=pad_token_id) + if position_ids is not None: + # pad position_id with pad_token_id + position_ids = F.pad(position_ids, (0, pad_len), value=pad_token_id) + # pad attention mask without attention on the padding tokens + attention_mask = F.pad(attention_mask, (0, pad_len), value=False) + # pad token_type_ids with token_type_id = 0 + token_type_ids = F.pad(token_type_ids, (0, pad_len), value=0) + + return pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds + + @staticmethod + def unpad_sequence_output(pad_len, sequence_output): + """This function unpads sequence output if inputs of the model were padded. + This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size. + It needs to be called in your model, such as BertModel, right before you return the model outputs. + + Arguments: + pad_len: required: an integer determining how much model inputs have been padded to transfer sequence length dimension to multiple of block size. + sequence_output: required: sequence output of the encoder layer. + + Return: + sequence_output: unpaded sequence output of the encoder layer. + """ + + if (pad_len > 0): + sequence_output = sequence_output[:, :-pad_len] + return sequence_output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..b673c4561902e943981ca3008fae53ec73c0cd73 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py @@ -0,0 +1,149 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch.nn as nn +import torch +from torch import distributed as dist +from deepspeed.ops.sparse_attention import SparsityConfig + + +class SparseSelfAttention(nn.Module): + """Implements an efficient Sparse Self Attention of Transformer layer based on `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + + For more information please see, TODO DeepSpeed Sparse Transformer. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial. + """ + + def __init__( + self, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=SparsityConfig(num_heads=4), + key_padding_mask_mode='add', + attn_mask_mode='mul', + max_seq_length=2048): + """Initialize the sparse self attention layer. + Arguments: + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class. + key_padding_mask_mode: optional: a string determining if key padding mask needs to be added, `add`, or be multiplied, `mul`. + attn_mask_mode: optional: a string determining if attention mask needs to be added, `add`, or be multiplied, `mul`. + max_seq_length: optional: the maximum sequence length this sparse attention module will be applied to; it controls the size of the master_layout. + """ + super().__init__() + + # sparsity information + self.sparsity_config = sparsity_config + + # initialize sparse layout and register as buffer + master_layout = self.sparsity_config.make_layout(max_seq_length) + self.register_buffer("master_layout", master_layout) + self._need_layout_synchronization = True + + # mask modes + self.key_padding_mask_mode = key_padding_mask_mode + self.attn_mask_mode = attn_mask_mode + + ops = dict() + + def get_layout(self, L): + # if layout is never synchronized across GPUs, broadcast the layout from global rank 0 + if self._need_layout_synchronization and dist.is_initialized(): + dist.broadcast(self.master_layout, src=0) + self._need_layout_synchronization = False + + if (L % self.sparsity_config.block != 0): + raise ValueError( + f'Sequence Length, {L}, needs to be dividable by Block size {self.sparsity_config.block}!') + + num_blocks = L // self.sparsity_config.block + return self.master_layout[..., :num_blocks, :num_blocks].cpu() # layout needs to be a CPU tensor + + # add to cache + def get_ops(self, H, L): + from deepspeed.ops.sparse_attention.matmul import MatMul + from deepspeed.ops.sparse_attention.softmax import Softmax + if L not in SparseSelfAttention.ops: + sparsity_layout = self.get_layout(L) + sparse_dot_sdd_nt = MatMul(sparsity_layout, self.sparsity_config.block, 'sdd', trans_a=False, trans_b=True) + + sparse_dot_dsd_nn = MatMul(sparsity_layout, + self.sparsity_config.block, + 'dsd', + trans_a=False, + trans_b=False) + + sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block) + + SparseSelfAttention.ops[L] = (sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax) + return SparseSelfAttention.ops[L] + + def transpose_key_for_scores(self, x, L): + bsz, num_heads, seq_len, head_dim = x.size() + if seq_len != L: + return x.permute(0, 1, 3, 2) + return x + + def transpose_mask_for_sparse(self, qtype, x, is_key_padding_mask=False): + x = x.type(qtype) + if is_key_padding_mask: + xdim = x.dim() + for d in range(xdim - 1, 0, -1): + x = x.squeeze(dim=d) + return x + return x.squeeze() + + # forward pass + def forward(self, query, key, value, rpe=None, key_padding_mask=None, attn_mask=None): + """Applies forward phase of sparse self attention + + Arguments: + query: required: query tensor + key: required: key tensor + value: required: value tensor + rpe: optional: a tensor same dimension as x that is used as relative position embedding + key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength) + attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported + key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied + attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied + + Return: + attn_output: a dense tensor containing attention context + """ + assert query.dtype == torch.half, "sparse attention only supports training in fp16 currently, please file a github issue if you need fp32 support" + bsz, num_heads, tgt_len, head_dim = query.size() + + # transpose back key if it is already transposed + key = self.transpose_key_for_scores(key, tgt_len) + + # check that operation is supported + if query.shape != key.shape or key.shape != value.shape: + raise NotImplementedError('only self-attention is supported for now') + + # squeeze key_padding_mask if it is given + if key_padding_mask is not None: + key_padding_mask = self.transpose_mask_for_sparse(query.dtype, key_padding_mask, is_key_padding_mask=True) + + # squeeze attn_mask if it is given + if attn_mask is not None: + attn_mask = self.transpose_mask_for_sparse(query.dtype, attn_mask) + + # cache look-up table computations etc + sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax = self.get_ops(num_heads, tgt_len) + + scaling = float(head_dim)**-0.5 + + # attention scores + attn_output_weights = sparse_dot_sdd_nt(query, key) + attn_output_weights = sparse_softmax(attn_output_weights, + scale=scaling, + rpe=rpe, + key_padding_mask=key_padding_mask, + attn_mask=attn_mask, + key_padding_mask_mode=self.key_padding_mask_mode, + attn_mask_mode=self.attn_mask_mode) + + # outputs + attn_output = sparse_dot_dsd_nn(attn_output_weights, value) + return attn_output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py new file mode 100644 index 0000000000000000000000000000000000000000..1f59c4b469282ebfdb26fdb7c40c61d738231030 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py @@ -0,0 +1,727 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import random + + +class SparsityConfig: + """Abstract Configuration class to store `sparsity configuration of a self attention layer`. + It contains shared property of different block-sparse sparsity patterns. However, each class needs to extend it based on required property and functionality. + """ + + def __init__(self, num_heads, block=16, different_layout_per_head=False): + """Initialize the Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + """ + + self.num_heads = num_heads + self.block = block + self.different_layout_per_head = different_layout_per_head + self.num_layout_heads = num_heads if different_layout_per_head else 1 + + def setup_layout(self, seq_len): + """Create layout tensor for the given sequence length + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) for sparsity layout of all head; initialized with zero + """ + + if (seq_len % self.block != 0): + raise ValueError(f'Sequence Length, {seq_len}, needs to be dividable by Block size {self.block}!') + num_blocks = seq_len // self.block + # TODO Currently we allocate layout per head; needs to be updated if heads share a single layout. + layout = torch.zeros((self.num_heads, num_blocks, num_blocks), dtype=torch.int64) + return layout + + def check_and_propagate_first_head_layout(self, layout): + """If all heads require same sparsity layout, it propagate first head layout to all heads + + Arguments: + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head + """ + + if not self.different_layout_per_head: + layout[1:self.num_heads, :, :] = layout[0, :, :] + return layout + + +class DenseSparsityConfig(SparsityConfig): + """Configuration class to store `Dense` configuration. + In reality, this is not sparse and all blocks are used. We keep it for the sake of comparison and comprehension. + """ + + def __init__(self, num_heads, block=16, different_layout_per_head=False): + """Initialize the Dense Sparsity Pattern Config. + In reality, this is not sparse and all blocks are used. We keep it for the sake of comparison and comprehension. + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + seq_len: required: an integer determining number of attention heads of the layer. + different_layout_per_head: optional: this is just for the sake of consistency with other sparsity formats; can ignore it for DenseSparsityConfig + """ + + super().__init__(num_heads, block, different_layout_per_head) + + def make_layout(self, seq_len): + """Set 1 to all blocks of the layout meaning the pattern is dense; not sparse. + + Arguments: + seq_len: required: an integer determining the underling sequence length; must be <= max sequence length + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; for dense everything is 1 + """ + + layout = self.setup_layout(seq_len) + layout[:, :, :] = 1 + return layout + + +class FixedSparsityConfig(SparsityConfig): + """Configuration class to store `Fixed` sparsity configuration. + For more details about this sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509; this has been customized. + This class extends parent class of `SparsityConfig` and customizes it for `Fixed` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_local_blocks=4, + num_global_blocks=1, + attention='bidirectional', + horizontal_global_attention=False, + num_different_global_patterns=1): + """Initialize `Fixed` Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + num_local_blocks: optional: an integer determining the number of blocks in local attention window. + num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks. + num_different_global_patterns: optional: an integer determining number of different global attentions layouts. While global attention can be fixed by which block/s are representative of any local window, since there are multi-heads, each head can use a different global representative. For example, with 4 blocks local window and global attention size of 1 block, we can have 4 different versions in which the first, Second, third, or forth block of each local window can be global representative of that window. This parameter determines how many of such patterns we want. Of course, there is a limitation based on num_local_blocks and num_global_blocks. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_local_blocks = num_local_blocks + + if (num_local_blocks % num_global_blocks != 0): + raise ValueError( + f'Number of blocks in a local window, {num_local_blocks}, must be dividable by number of global blocks, {num_global_blocks}!' + ) + self.num_global_blocks = num_global_blocks + + if (attention != 'unidirectional' and attention != 'bidirectional'): + raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') + self.attention = attention + + if (attention != 'bidirectional' and horizontal_global_attention): + raise ValueError('only \"bi-directional\" attentions can support horizontal global attention!') + self.horizontal_global_attention = horizontal_global_attention + + if (num_different_global_patterns > 1 and not different_layout_per_head): + raise ValueError( + f'Number of different layouts cannot be more than one when you have set a single layout for all heads! Set different_layout_per_head to True.' + ) + if (num_different_global_patterns > (num_local_blocks // num_global_blocks)): + raise ValueError( + f'Number of layout versions (num_different_global_patterns), {num_different_global_patterns}, cannot be larger than number of local window blocks divided by number of global blocks, {num_local_blocks} / {num_global_blocks} = {num_local_blocks//num_global_blocks}!' + ) + self.num_different_global_patterns = num_different_global_patterns + + def set_local_layout(self, h, layout): + """Sets local attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local layout is set + """ + + num_blocks = layout.shape[1] + for i in range(0, num_blocks, self.num_local_blocks): + end = min(i + self.num_local_blocks, num_blocks) + for row in range(i, end): + for col in range(i, (row + 1 if self.attention == 'unidirectional' else end)): + layout[h, row, col] = 1 + return layout + + def set_global_layout(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Currently we set global blocks starting from the last block of a local window to the first one. That means if a local window consists of 4 blocks and global attention size is one block, we use block #4 in each local window as global. If we have different layout per head, then other heads will get #3, #2, and #1. And if we have more heads (and different layout has set) than num of global attentions, multiple head may have same global attentions. + Note) if horizontal_global_attention is set, global blocks will be set both horizontally and vertically. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + first_global_block_idx = self.num_local_blocks - ( + 1 + h % self.num_different_global_patterns) * self.num_global_blocks + + # set all global blocks except the last one if (in last local window) + end = num_blocks - (num_blocks % self.num_local_blocks) + for i in range(first_global_block_idx, end, self.num_local_blocks): + + # vertical global attention + first_row = 0 if self.attention == 'bidirectional' else i + #(((i // self.num_local_blocks) + 1) * self.num_local_blocks) + #if (first_row < num_blocks): + layout[h, first_row:, i:i + self.num_global_blocks] = 1 + + # horizontal global attention; only in bidirectional attention + if (self.horizontal_global_attention): + layout[h, i:i + self.num_global_blocks, :] = 1 + + # set last global blocks; handle possible short last local window + if (end < num_blocks): + start = min(end + first_global_block_idx, num_blocks - self.num_global_blocks) + end = start + self.num_global_blocks + + # vertical global attention + first_row = 0 if self.attention == 'bidirectional' else start + #(((start // self.num_local_blocks) + 1) * self.num_local_blocks) + #if (first_row < num_blocks): + layout[h, first_row:, start:end] = 1 + + # horizontal global attention + if (self.horizontal_global_attention): + layout[h, start:end, :] = 1 + return layout + + def make_layout(self, seq_len): + """Generates `Fixed` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `Fixed` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_local_layout(h, layout) + layout = self.set_global_layout(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class VariableSparsityConfig(SparsityConfig): + """Configuration class to store `Variable` sparsity configuration. + This layout is an extension of FixedSparsityConfig in which: + - user can set random layout; default value is zero means no random block + - user can provide a list of local block sizes + - user can provide a list of global block indices. + + For more details about `Fixed` sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509; this has been customized. + This class extends parent class of `SparsityConfig` and customizes it for `Fixed` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_random_blocks=0, + local_window_blocks=[4], + global_block_indices=[0], + global_block_end_indices=None, + attention='bidirectional', + horizontal_global_attention=False): + """Initialize `Variable` Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. Currently this sparsity config can only assign single layout to all heads; needs to be extended for different layout per head. + num_random_blocks: optional: an integer determining the number of random blocks in each block row. + local_window_blocks: optional: a list of integers determining the number of blocks in each local attention window. It assumes first number determines # of blocks in the first local window, second the second window, ..., and the last number determines the number of blocks in the remaining local windows. + global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window. + global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention. + num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_random_blocks = num_random_blocks + self.local_window_blocks = local_window_blocks + self.global_block_indices = global_block_indices + + if (global_block_end_indices is not None): + if (len(global_block_indices) != len(global_block_end_indices)): + raise ValueError( + f'Global block start indices length, {len(global_block_indices)}, must be same as global block end indices length, {len(global_block_end_indices)}!' + ) + for _, (start_idx, end_idx) in enumerate(zip(global_block_indices, global_block_end_indices)): + if start_idx >= end_idx: + raise ValueError( + f'Global block start index, {start_idx}, must be smaller than global block end index, {end_idx}!' + ) + self.global_block_end_indices = global_block_end_indices + + if (attention != 'unidirectional' and attention != 'bidirectional'): + raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') + self.attention = attention + + if (attention != 'bidirectional' and horizontal_global_attention): + raise ValueError('only \"bi-directional\" attentions can support horizontal global attention!') + self.horizontal_global_attention = horizontal_global_attention + + def set_random_layout(self, h, layout): + """Sets random attention layout used by the given head in the sparse attention. + Note) By default, it assumes there will be a unique random block layout for all heads; unless `different_layout_per_head` parameter is set in which each head can have a different random layout. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which random layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_random_blocks): + raise ValueError( + f'Number of random blocks, {self.num_random_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + for row in range(0, num_blocks): + rnd_cols = random.sample(range(0, num_blocks), self.num_random_blocks) + layout[h, row, rnd_cols] = 1 + return layout + + def set_local_layout(self, h, layout): + """Sets local attention layout used by the given head in the sparse attention. + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local layout is set + """ + + num_blocks = layout.shape[1] + start_block_idx = 0 + end_block_idx = 0 + for block_size in self.local_window_blocks: + end_block_idx += block_size + end_block_idx = min(end_block_idx, num_blocks) + for row in range(start_block_idx, end_block_idx): + for col in range(start_block_idx, (row + 1 if self.attention == 'unidirectional' else end_block_idx)): + layout[h, row, col] = 1 + start_block_idx += block_size + + # if there is any remaining not attended part, use the lats local window block size as local window for the remaining applicable local windows + for i in range(start_block_idx, num_blocks, block_size): + end_block_idx = min(i + block_size, num_blocks) + for row in range(i, end_block_idx): + for col in range(i, (row + 1 if self.attention == 'unidirectional' else end_block_idx)): + layout[h, row, col] = 1 + return layout + + def set_global_layout(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + if (self.global_block_end_indices is None): + for idx in self.global_block_indices: + # if global block idx is in the range of the sequence blocks + if (idx < num_blocks): + #global rows + if (self.horizontal_global_attention): + layout[h, idx, :] = 1 + + #global columns + first_row = 0 if self.attention == 'bidirectional' else idx + layout[h, first_row:, idx] = 1 + else: + for _, (start_idx, end_idx) in enumerate(zip(self.global_block_indices, self.global_block_end_indices)): + # if global block idx is in the range of the sequence blocks + if (start_idx < num_blocks): + end_idx = min(end_idx, num_blocks) + #global rows + if (self.horizontal_global_attention): + layout[h, start_idx:end_idx, :] = 1 + + #global columns + first_row = 0 if self.attention == 'bidirectional' else start_idx + layout[h, first_row:, start_idx:end_idx] = 1 + return layout + + def make_layout(self, seq_len): + """Generates `Variable` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `Variable` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_random_layout(h, layout) + layout = self.set_local_layout(h, layout) + layout = self.set_global_layout(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class BigBirdSparsityConfig(SparsityConfig): + """Configuration class to store `BigBird` sparsity configuration. + For more details about this sparsity config, please see `Big Bird: Transformers for Longer Sequences`: https://arxiv.org/pdf/2007.14062.pdf + This class extends parent class of `SparsityConfig` and customizes it for `BigBird` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_random_blocks=1, + num_sliding_window_blocks=3, + num_global_blocks=1, + attention='bidirectional'): + """Initialize the BigBird Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + num_random_blocks: optional: an integer determining the number of random blocks in each block row. + num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. + num_global_blocks: optional: an integer determining how many consecutive blocks, starting from index 0, are considered as global attention. Global block tokens will be attended by all other block tokens and will attend to all other block tokens as well. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_random_blocks = num_random_blocks + self.num_sliding_window_blocks = num_sliding_window_blocks + self.num_global_blocks = num_global_blocks + + if (attention != 'unidirectional' and attention != 'bidirectional'): + raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') + self.attention = attention + + def set_random_layout(self, h, layout): + """Sets random attention layout used by the given head in the sparse attention. + Note) By default, it assumes there will be a unique random block layout for all heads; unless `different_layout_per_head` parameter is set in which each head can have a different random layout. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which random layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_random_blocks): + raise ValueError( + f'Number of random blocks, {self.num_random_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + for row in range(0, num_blocks): + sample_range = range(0, num_blocks) if self.attention == 'bidirectional' else range(0, row + 1) + rnd_cols = random.sample(sample_range, self.num_random_blocks) + layout[h, row, rnd_cols] = 1 + return layout + + def set_sliding_window_layout(self, h, layout): + """Sets sliding local attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_sliding_window_blocks): + raise ValueError( + f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + w = self.num_sliding_window_blocks // 2 + for row in range(0, num_blocks): + start = max(0, row - w) + end = min(row + w + 1, num_blocks) + layout[h, row, start:end] = 1 + return layout + + def set_global_layout_itc(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_global_blocks): + raise ValueError( + f'Number of global blocks, {self.num_global_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + #global rows + layout[h, 0:self.num_global_blocks, :] = 1 + + #global columns + layout[h, :, 0:self.num_global_blocks] = 1 + + if self.attention == 'unidirectional': + # zero out anything attending to the future + layout = torch.tril(layout) + + return layout + + def make_layout(self, seq_len): + """Generates `BigBird` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BigBird` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_random_layout(h, layout) + layout = self.set_sliding_window_layout(h, layout) + layout = self.set_global_layout_itc(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class BSLongformerSparsityConfig(SparsityConfig): + """Configuration class to store edited `Longformer` sparsity configuration. + + Note) this is a block-sparse version of the Longformer which is slightly different than original Longformer; which is element-wise sparsity. + + For more details about this sparsity config, please see `Longformer: The Long-Document Transformer`: https://arxiv.org/pdf/2004.05150.pdf + This class extends parent class of `SparsityConfig` and customizes it for `Longformer` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_sliding_window_blocks=3, + global_block_indices=[0], + global_block_end_indices=None, + attention='bidirectional'): + """Initialize the edited `Longformer` Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + + num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. + global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window. + global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_sliding_window_blocks = num_sliding_window_blocks + self.global_block_indices = global_block_indices + self.attention = attention + + if (global_block_end_indices is not None): + if (len(global_block_indices) != len(global_block_end_indices)): + raise ValueError( + f'Global block start indices length, {len(global_block_indices)}, must be same as global block end indices length, {len(global_block_end_indices)}!' + ) + for _, (start_idx, end_idx) in enumerate(zip(global_block_indices, global_block_end_indices)): + if start_idx >= end_idx: + raise ValueError( + f'Global block start index, {start_idx}, must be smaller than global block end index, {end_idx}!' + ) + self.global_block_end_indices = global_block_end_indices + + def set_sliding_window_layout(self, h, layout): + """Sets sliding local attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_sliding_window_blocks): + raise ValueError( + f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + w = self.num_sliding_window_blocks // 2 + for row in range(0, num_blocks): + start = max(0, row - w) + end = min(row + w + 1, num_blocks) + layout[h, row, start:end] = 1 + return layout + + def set_global_layout(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + if (self.global_block_end_indices is None): + for idx in self.global_block_indices: + # if global block idx is in the range of the sequence blocks + if (idx < num_blocks): + #global rows + layout[h, idx, :] = 1 + + #global columns + layout[h, :, idx] = 1 + else: + for _, (start_idx, end_idx) in enumerate(zip(self.global_block_indices, self.global_block_end_indices)): + # if global block idx is in the range of the sequence blocks + if (start_idx < num_blocks): + end_idx = min(end_idx, num_blocks) + #global rows + layout[h, start_idx:end_idx, :] = 1 + + #global columns + layout[h, :, start_idx:end_idx] = 1 + if self.attention == 'unidirectional': + layout = torch.tril(layout) + return layout + + def make_layout(self, seq_len): + """Generates edited `Longformer` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BSLongformer` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_sliding_window_layout(h, layout) + layout = self.set_global_layout(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class LocalSlidingWindowSparsityConfig(SparsityConfig): + """Configuration class to store `Local Sliding Window` sparsity configuration - a purely-local sliding window attention. + This class extends parent class of `SparsityConfig` and customizes it for `Local` sparsity. + """ + + def __init__(self, num_heads, block=16, num_sliding_window_blocks=3, attention='unidirectional'): + """Initialize the Local Sliding Window Sparsity Pattern Config. + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + """ + + super().__init__(num_heads, block) + self.num_sliding_window_blocks = num_sliding_window_blocks + self.attention = attention + + def set_sliding_window_layout(self, h, layout): + """Sets sliding local attention layout used by the given head in the sparse attention. + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_sliding_window_blocks): + raise ValueError( + f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + w = self.num_sliding_window_blocks // 2 + for row in range(0, num_blocks): + start = max(0, row - w) + end = min(row + w + 1, num_blocks) if self.attention == "bidirectional" else row + 1 + layout[h, row, start:end] = 1 + return layout + + def make_layout(self, seq_len): + """Generates `Local Sliding Window` sparsity layout used by each head in the sparse attention. + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BigBird` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_sliding_window_layout(h, layout) + layout = self.check_and_propagate_first_head_layout(layout) + return layout diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..18d58b5af4bf387ffc8b4e90645e3bd578ee10fe --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import os + + +def _build_file_index(directory, suffix='.tr'): + """Build an index of source files and their basenames in a given directory. + + Args: + directory (string): the directory to index + suffix (string): index files with this suffix + + Returns: + list: A list of tuples of the form [(basename, absolute path), ...] + """ + + index = [] + + for fname in os.listdir(directory): + if fname.endswith(suffix): + basename = fname[:fname.rfind(suffix)] # strip the suffix + path = os.path.join(directory, fname) + index.append((basename, path)) + + return index + + +# Go over all local source files and parse them as strings +_module = sys.modules[_build_file_index.__module__] +_directory = os.path.dirname(os.path.realpath(__file__)) +for name, fname in _build_file_index(_directory): + with open(fname, 'r') as fin: + setattr(_module, name, fin.read()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c248736d3fddb49f3f34113bf5d9e71f5ad8e61e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr new file mode 100644 index 0000000000000000000000000000000000000000..9be6aa2aed7ccfb0d438d2e4d766d30b21d8f744 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr @@ -0,0 +1,208 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py +*/ + +__global__ void NAME (TYPE* A __readonly __noalias __aligned(16), + TYPE* B __readonly __noalias __aligned(16), + TYPE* C __noalias __aligned(16), + int lda __multipleof(8), + int ldb __multipleof(8), + int ldc __multipleof(8), + long stride_za __multipleof(8), + long stride_zb __multipleof(8), + long stride_zc __multipleof(8), + long stride_ha __multipleof(8), + long stride_hb __multipleof(8), + long stride_hc __multipleof(8), + int DS0, int DS1, + int SDD_K __multipleof(16), + int SDD_off_width, + int* lut, int* locks, int nlocks) { + /* ---------------- */ + /* Prologue */ + /* ---------------- */ + // program ids + int pid0 = get_program_id(0); + int pid1 = get_program_id(1); + int pidz = get_program_id(2); +#ifdef SDD + // load LUT header + pid1 = pid1 + SDD_off_width; + int blockidm[TM] = (0 ... TM) / BLOCK; + int blockidn[TN] = (0 ... TN) / BLOCK; + int offlutm[TM] = blockidm*(TN/BLOCK)*4; + int offlutn[TN] = blockidn*4; + int *header = lut + pid1 * (TM/BLOCK) * (TN/BLOCK) * 4; + int z = *(header + 0); + int i[TM] = *(header + 1 + offlutm); + int j[TN] = *(header + 2 + offlutn); + int AS1 = SDD_K / TZ; + int lockid = select(TZ > 1, 1, 0); + int offka = pid0 * AS1; + int offkb = pid0 * AS1; + int offmc = 0; + int offnc = 0; + int offpa = 0; + int offpb = 0; + int maxid = TZ; + int offhc = 0; + int offha = z; + int offhb = z; + int ram[TM] = i*BLOCK + ((0 ... TM) % BLOCK); + int rbn[TN] = j*BLOCK + ((0 ... TN) % BLOCK); +#else + // load LUT header + int *header = lut + pid0 * 6; + int offset = *(header + 0); + int AS1 = *(header + 1); + int column = *(header + 2); + int depth = *(header + 3); + int lockid = *(header + 4); + int maxid = *(header + 5); + int *pinc = lut + offset; + int offhc = depth; +#ifdef DSD + // output offset + int offnc = pid1 * TN; + int offmc = column * TM; + int offpc = 0; + // dense input offset + int offnb = pid1 * TN; + int offkb __multipleof(8) = *pinc; + int offpb = 0; + // sparse input offset + int offma = 0; + int offka = 0; + long offpa __multipleof(8) = *(pinc + 1); + offpa = offpa * BLOCK * BLOCK; + int offha = 0; + int offhb = depth; +#endif +#ifdef DDS + // output offset + int offmc = pid1 * TM; + int offnc = column * TN; + int offpc = 0; + // dense input offset + int offma = pid1 * TM; + int offka __multipleof(8) = *pinc; + int offpa = 0; + // sparse input offset + int offnb = 0; + int offkb = 0; + long offpb __multipleof(8) = *(pinc + 1); + offpb = offpb * BLOCK * BLOCK; + int offha = depth; + int offhb = 0; +#endif + int ram[TM] = offma + 0 ... TM; + int rbn[TN] = offnb + 0 ... TN; +#endif + // initialize a, b pointers + int rka[TK] = offka + 0 ... TK; + int rkb[TK] = offkb + 0 ... TK; + TYPE* pa[TM, TK] = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, newaxis] * STRIDE_AM + rka[newaxis, :] * STRIDE_AK; + TYPE* pb[TK, TN] = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[newaxis, :] * STRIDE_BN + rkb[:, newaxis] * STRIDE_BK; + // pre-fetch +#ifdef DDS + bool checkam[TM, TK] = ram[:, newaxis] < DS0; +#else + bool checkam[TM, TK] = AS1 > 0; +#endif +#ifdef DSD + bool checkbn[TK, TN] = rbn[newaxis, :] < DS0; +#else + bool checkbn[TK, TN] = AS1 > 0; +#endif + TYPE a[TM, TK] = checkam ? *pa : 0; + TYPE b[TK, TN] = checkbn ? *pb : 0; + + /* ---------------- */ + /* Inner Loop */ + /* ---------------- */ + // create result tile + float acc[TM, TN] = 0; + int step = TK; + for(int k = AS1; k > 0; k -= step) { + acc += a @ b; + // update pointers +#ifdef SDD + int inc_a = TK * STRIDE_AK; + int inc_b = TK * STRIDE_BK; +#else + pinc += 2; +#ifdef DSD + int inc_b __multipleof(8) = *pinc; + int inc_a __multipleof(8) = *(pinc + 1); + inc_b = inc_b * STRIDE_BK; +#endif +#ifdef DDS + int inc_a __multipleof(8) = *pinc; + int inc_b __multipleof(8) = *(pinc + 1); + inc_a = inc_a * STRIDE_AK; +#endif +#endif + pa += inc_a; + pb += inc_b; + // pre-fetch + bool checkak[TM, TK] = k > TK; + bool checkbk[TK, TN] = k > TK; + bool checka[TM, TK] = checkam && checkak; + bool checkb[TK, TN] = checkbk && checkbn; + a = *?(checka)pa; + b = *?(checkb)pb; + } + TYPE c[TM, TN] = acc; + + /* ---------------- */ + /* Epilogue */ + /* ---------------- */ + // initialize c pointers +#ifdef SDD + bool checkc[TM, TN] = 1; + // rematerialize + int rr_blockidm[TM] = (0 ... TM) / BLOCK; + int rr_blockidn[TN] = (0 ... TN) / BLOCK; + int rr_offlutm[TM] = rr_blockidm*(TN/BLOCK)*4; + int rr_offlutn[TN] = rr_blockidn*4; + int off_bkid[TM, TN] = 3 + rr_offlutm[:, newaxis] + rr_offlutn[newaxis, :]; + int bkid[TM, TN] = *(header + off_bkid); + long offpc[TM, TN] = bkid * BLOCK * BLOCK; + // range within blocks + int rcm[TM] = (0 ... TM) % BLOCK; + int rcn[TN] = (0 ... TN) % BLOCK; +#else + int rcm[TM] = offmc + 0 ... TM; + int rcn[TN] = offnc + 0 ... TN; +#ifdef DSD + bool checkc[TM, TN] = rcn[newaxis, :] < DS0; +#endif +#ifdef DDS + bool checkc[TM, TN] = rcm[:, newaxis] < DS0; +#endif +#endif + TYPE* pc[TM, TN] = C + offpc + offhc*stride_hc + pidz*stride_zc + rcm[:, newaxis]*STRIDE_CM + rcn[newaxis, :]*STRIDE_CN; + // write-back directly + if(lockid == 0) { + *?(checkc) pc = c; + } + // accumulate partial result using spin-locks + else { + int *plock = locks + get_program_id(2)*nlocks*get_num_programs(1) + get_program_id(1)*nlocks + lockid - 1; + int *pcount = plock + get_num_programs(2)*get_num_programs(1)*nlocks; + for(int repeat = 1; repeat == 1; repeat = atomic_cas(plock, 0, 1)); + int count = *pcount; + if(count == 0) + *?(checkc) pc = c; + else + *?(checkc) pc = c + *?(checkc)pc; + atomic_xchg(pcount, (count + 1) % maxid); + atomic_xchg(plock, 0); + } + } diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr new file mode 100644 index 0000000000000000000000000000000000000000..1a90f41d94945e1d6d6f52e6beaea94fa52cdda8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/softmax.py +*/ + +__global__ void softmax_bwd(TYPE * X __readonly __noalias __aligned(16), + float scale, + TYPE* DX __readonly __noalias __aligned(16), + int* LUT, + int sizemax, + long stride_zx __multipleof(BLOCK), + long stride_zdx __multipleof(BLOCK)) { + int pidhm = get_program_id(0); + int pidz = get_program_id(1); + + // create index ranges + int rxm = pidhm % BLOCK; + int rbm = pidhm / BLOCK; + int rxn[TN] = (0 ... TN) % BLOCK; + int rbn[TN] = (0 ... TN) / BLOCK; + + // extract information from look-up table + int* header = LUT + rbm * 2; + int size = *(header + 0); + int offset = *(header + 1); + + // bounds checking on lut + bool check[TN] = rbn < size; + int rbmn[TN] = check ? rbn : size - 1; + + // initialize pointers to block-sparse input + long blockid[TN] = *(LUT + offset + rbmn*4); + + TYPE* px[TN] = X + pidz * stride_zx + + blockid * BLOCK * BLOCK + + rxm * BLOCK + + rxn; + + TYPE* pdx[TN] = DX + pidz * stride_zdx + + blockid * BLOCK * BLOCK + + rxm * BLOCK + + rxn; + + // compute fused softmax backward + TYPE x[TN] = check ? *px : 0; + TYPE dx[TN] = check ? *pdx : 0; + float Fdx[TN] = dx; + float Fx[TN] = x; + float Fxdx[TN] = Fdx*Fx; + float Fxdxsum = Fxdx[+]; + float Fy[TN] = Fx * (Fdx - Fxdxsum) * scale; + TYPE y[TN] = Fy; + + // write-back + *? (check)pdx = y; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr new file mode 100644 index 0000000000000000000000000000000000000000..ebd317d9469b47f7e2ee3032d3aabf57b5620a73 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr @@ -0,0 +1,143 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/softmax.py +*/ + +__global__ void softmax_fwd(TYPE *X __readonly __noalias __aligned(16), + float scale, + int *LUT __readonly __noalias __aligned(16), + TYPE *RPE __readonly __noalias __aligned(16), + TYPE *KP_M __readonly __noalias __aligned(16), + TYPE *ATTN_M __readonly __noalias __aligned(16), + int num_blocks, + int sizemax, + long stride_zx __multipleof(BLOCK), + long stride_zrpe __multipleof(BLOCK), + int stride_hrpe __multipleof(BLOCK), + int stride_srpe __multipleof(BLOCK), + int stride_zkpm __multipleof(BLOCK), + int stride_zattnm __multipleof(BLOCK)){ + int pidhm = get_program_id(0); + int pidz = get_program_id(1); + + // create index ranges + int rxm = pidhm % BLOCK; + int rbm = pidhm / BLOCK; + int rxn[TN] = (0 ... TN) % BLOCK; + int rbn[TN] = (0 ... TN) / BLOCK; + + // extract information from look-up table + int* header = LUT + rbm * 2; + int size = *(header + 0); + int offset = *(header + 1); + + bool check[TN] = rbn < size; + int rbmn[TN] = check ? rbn : size - 1; + + // block id and column id + long blockid [TN] = *(LUT + offset + rbmn*4 + 0); + long columnid[TN] = *(LUT + offset + rbmn*4 + 1); + long rowid [TN] = *(LUT + offset + rbmn*4 + 2); + long headid [TN] = *(LUT + offset + rbmn*4 + 3); + + // pointers to X + TYPE* px[TN] = X + pidz * stride_zx + + blockid * BLOCK * BLOCK + + rxm * BLOCK + + rxn; +#ifdef APPLY_RPE + // pointers to relative position embedding + TYPE* prpe[TN] = RPE + pidz * stride_zrpe + + headid * stride_hrpe + + columnid * BLOCK + + rowid * BLOCK * stride_srpe + + rxm * stride_srpe + + rxn; +#endif + +#ifdef APPLY_KP_MASK + // pointers to key padding mask + TYPE* pkp_m[TN] = KP_M + pidz * stride_zkpm + + columnid * BLOCK + + rxn; +#endif + +#ifdef APPLY_ATTN_MASK + // pointers to attention mask + TYPE* pattn_m[TN] = ATTN_M + columnid * BLOCK + + rowid * BLOCK * stride_zattnm + + rxm * stride_zattnm + + rxn; +#endif + + // load input + TYPE x[TN] = check ? *px : -INFINITY; + +#ifdef APPLY_RPE + // load relative position embedding + TYPE rpe[TN] = check ? *prpe : 0; +#endif + +#ifdef APPLY_KP_MASK + // load key-padding mask + TYPE kp_m[TN] = check ? *pkp_m : -INFINITY; +#endif + +#ifdef APPLY_ATTN_MASK + // load attention mask + TYPE attn_m[TN] = check ? *pattn_m : -INFINITY; +#endif + + // compute softmax in float +#ifdef APPLY_RPE + float Frpe[TN] = rpe; +#endif + +#ifdef APPLY_KP_MASK + float Fkp_m[TN] = kp_m; +#endif + +#ifdef APPLY_ATTN_MASK + float Fattn_m[TN] = attn_m; +#endif + +#ifdef KP_MASK_MUL + Fkp_m = (Fkp_m == 0) ? (float[TN])-INFINITY : 0; +#endif + +#ifdef ATTN_MASK_MUL + Fattn_m = (Fattn_m == 0) ? (float[TN])-INFINITY : 0; +#endif + + float Fx[TN] = x; + +#ifdef APPLY_SCALE + Fx = Fx * scale; // apply scale +#endif + +#ifdef APPLY_RPE + Fx = Fx + Frpe; // apply relative position embedding +#endif + +#ifdef APPLY_KP_MASK + Fx = Fx + Fkp_m; // apply key padding mask +#endif + +#ifdef APPLY_ATTN_MASK + Fx = Fx + Fattn_m; // apply attention mask +#endif + + float Fxmax = Fx[max]; + float Fy[TN] = exp(Fx - Fxmax); + float Fysum = (check ? Fy : 0)[+]; + + // write-back in half/float + TYPE y[TN] = Fy; + TYPE ysum = Fysum; + *?(check)px = y / ysum; +} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b80fe2b4ba714611b6bcb652d9e559d87c7ed6fd --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig +from .inference.config import DeepSpeedInferenceConfig +from ...model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference +from .inference.moe_inference import DeepSpeedMoEInferenceConfig, DeepSpeedMoEInference diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d40e38a15fccf16639bf041eda206def43db479c Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25953a0a79bca2f45e8e8e33112ad68855089524 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b31a90eac2850fa7554214760c00df9815db25 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .config import DeepSpeedInferenceConfig +from ....model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference +from .moe_inference import DeepSpeedMoEInferenceConfig, DeepSpeedMoEInference diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fccdc794e975c1715426602a0131da16329bd6a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac10b402be0d1e377c8e1d04558b3f2c7c1a2638 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46c012eda8fd71bb8faa1c7720934e8d0b9acffa Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cca2957c9857a3f845085e0b6331c051e295afd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c65d09e74f5dc34c7c5f6805e4d02feebd18d63 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa024f67be7410ad10443fa4ca13efbfb52da562 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a3c9c97ef21b6fc25dc923ba32760493b2e2529 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c59b1dbc94c8fc8e3cc70248c90628563fd717c0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..227a21d78223723e650d4719beb6c00e169b5737 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..073603c404114efb97c3a8232e511a7946eee205 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/bias_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/bias_add.py new file mode 100644 index 0000000000000000000000000000000000000000..253784f001aeb3431a0d60812e4d2068c0bd5455 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/bias_add.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Optional +import torch +from deepspeed.ops.op_builder import SpatialInferenceBuilder + +spatial_cuda_module = None + + +def nhwc_bias_add(activation: torch.Tensor, + bias: torch.Tensor, + other: Optional[torch.Tensor] = None, + other_bias: Optional[torch.Tensor] = None) -> torch.Tensor: + global spatial_cuda_module + if spatial_cuda_module is None: + spatial_cuda_module = SpatialInferenceBuilder().load() + + if other is None: + return spatial_cuda_module.nhwc_bias_add(activation, bias) + elif other_bias is None: + return spatial_cuda_module.nhwc_bias_add_add(activation, bias, other) + else: + return spatial_cuda_module.nhwc_bias_add_bias_add(activation, bias, other, other_bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/config.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c0dd29f4f9627e79a294089323d8bb83728e13db --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/config.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import json +import torch +from deepspeed.utils.types import ActivationFuncType, NormType + + +class TransformerConfig(): + + def __init__(self, hidden_size, intermediate_size, heads, num_hidden_layers): + self.layer_id = -1 + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.heads = heads + self.num_hidden_layers = num_hidden_layers + + +class DeepSpeedInferenceConfig(TransformerConfig): + """Initialize the DeepSpeed Transformer Config. + Arguments: + hidden_size: The hidden size of the transformer layer + intermediate_size: The intermediate size of the feed-forward part of transformer layer + heads: The number of heads in the self-attention of the transformer layer + num_hidden_layers: The number of transformer layers + layer_norm_eps: The epsilon value for the layer norm + local_rank: Optional: The rank of GPU running the transformer kernel, it is not required + to use if the model already set the current device, otherwise need to set it + so that the transformer kernel can work on the right device + mp_size (optional): This argument is mainly used to create the parameters on the kernel side + using model-parallel architecture. If the client model already takes care of this, there is no + need to pass this argument. + pre_layer_norm: Select between Pre-LN or Post-LN transformer architecture + stochastic_mode: Enable for high performance, please note that this flag has some level of + non-determinism and can produce different results on different runs. However, we have seen + that by enabling it, the pretraining tasks such as BERT are not affected and can obtain + a high accuracy level. On the other hand, for the downstream tasks, such as fine-tuning, we recommend + to turn it off in order to be able to reproduce the same result through the regular kernel execution. + + scale_attention: If true, both q and k are scaled by 1/sqrt(attention_heads) before attention computation. + return_tuple: if True, returns the transformer output as a tuple, otherwise returns as a tensor + bigscience_bloom: This flag is added temporarily for supporting the BLOOM-176B model architecture. + use_triton: This flag is to enable triton kernels in inference or not. + invert_mask: If True, the attention mask is inverted when passed to attention block. + """ + + def __init__(self, + hidden_size=-1, + intermediate_size=-1, + heads=-1, + num_hidden_layers=-1, + layer_norm_eps=1e-12, + local_rank=-1, + mp_size=1, + dtype=torch.float16, + pre_layer_norm=True, + norm_type=NormType.LayerNorm, + stochastic_mode=False, + scale_attention=True, + triangular_masking=True, + local_attention=False, + window_size=256, + rotary_dim=-1, + rotate_half=False, + rotate_every_two=True, + return_tuple=True, + mlp_after_attn=True, + mlp_act_func_type=ActivationFuncType.GELU, + training_mp_size=1, + bigscience_bloom=False, + max_out_tokens=1024, + min_out_tokens=1, + enable_qkv_quantization=False, + use_mup=False, + scale_attn_by_inverse_layer_idx=False, + return_single_tuple=False, + set_empty_params=False, + transposed_mode=False, + use_triton=False, + triton_autotune=False, + num_kv=-1, + rope_theta=10000, + invert_mask=True): + super(DeepSpeedInferenceConfig, + self).__init__(hidden_size, (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, + num_hidden_layers) + self.dtype = dtype + self.pre_layer_norm = pre_layer_norm + self.norm_type = norm_type + self.local_rank = local_rank + self.stochastic_mode = stochastic_mode + self.epsilon = layer_norm_eps + self.mp_size = mp_size + self.scale_attention = scale_attention + self.triangular_masking = triangular_masking + self.local_attention = local_attention + self.window_size = window_size + self.rotary_dim = rotary_dim + self.rotate_half = rotate_half + self.rotate_every_two = rotate_every_two + self.return_tuple = return_tuple + self.mlp_after_attn = mlp_after_attn + self.mlp_act_func_type = mlp_act_func_type + self.training_mp_size = training_mp_size + self.bigscience_bloom = bigscience_bloom + self.max_out_tokens = max_out_tokens + self.min_out_tokens = min_out_tokens + self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx + self.enable_qkv_quantization = enable_qkv_quantization + self.use_mup = use_mup + self.return_single_tuple = return_single_tuple + self.set_empty_params = set_empty_params + self.transposed_mode = transposed_mode + self.use_triton = use_triton + self.triton_autotune = triton_autotune + self.num_kv = num_kv + self.rope_theta = rope_theta + self.invert_mask = invert_mask + + @classmethod + def from_dict(cls, json_object): + config = DeepSpeedInferenceConfig() + for key, value in json_object.items(): + config.__dict__[key] = value + return config + + @classmethod + def from_json_file(cls, json_file): + with open(json_file, "r", encoding='utf-8') as reader: + text = reader.read() + return cls.from_dict(json.loads(text)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_2d_transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_2d_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..fa4c6d53f871d326d0a51f2bc9ae7c71743befda --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_2d_transformer.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + + +class Diffusers2DTransformerConfig(): + + def __init__(self, int8_quantization=False): + self.int8_quantization = int8_quantization diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..3c2340ccfc6f9f2bc66fd1f1fa10d4a7c2b2aaf9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py @@ -0,0 +1,192 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +import torch +from torch.autograd import Function +import torch.nn as nn +from packaging import version as pkg_version +from deepspeed.utils.logging import log_dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp +from deepspeed.ops.transformer.inference.op_binding.softmax_context import SoftmaxContextOp +from deepspeed.ops.transformer.inference.op_binding import LinearOp +from deepspeed.ops.transformer.inference.op_binding.pad_transform import PadTransformOp + +minus_inf = -10000.0 +triton_flash_attn = None + + +def load_triton_flash_attn(): + global triton_flash_attn + try: + import triton + except ImportError: + raise ImportError("Please install triton 2.0+ or `pip install deepspeed[sd]`") + + if pkg_version.parse(triton.__version__) < pkg_version.parse("2.0"): + raise ImportError("Please install triton 2.0+ or `pip install deepspeed[sd]`") + + from .triton_ops import triton_flash_attn + + +class DeepSpeedDiffusersAttentionFunction(Function): + + @staticmethod + def forward(ctx, input, context, input_mask, config, attn_qkvw, attn_qw, attn_kw, attn_vw, attn_qkvb, + num_attention_heads_per_partition, norm_factor, hidden_size_per_partition, attn_ow, attn_ob, + do_out_bias, score_context_func, linear_func, pad_transform_func, triton_flash_attn_kernel, + rope_theta): + + def _transpose_for_context(x): + x = x.permute(0, 2, 1, 3) + new_x_layer_shape = x.size()[:-2] + \ + (hidden_size_per_partition,) + return x.reshape(*new_x_layer_shape) + + def _transpose_for_scores(x): + attention_head_size = x.shape[-1] // num_attention_heads_per_partition + new_x_shape = x.size()[:-1] + (num_attention_heads_per_partition, attention_head_size) + x = x.reshape(*new_x_shape) + x = x.permute(0, 2, 1, 3) + return x.contiguous() + + def selfAttention_fp(input, context, input_mask): + if config.dtype in [torch.half, torch.float16] and input.dtype == torch.float32: + input = input.half() + head_size = input.shape[-1] // config.heads + do_flash_attn = (head_size <= 128) + scale = (1 / norm_factor) * (1 / norm_factor) + if do_flash_attn and context is None: + qkv_out = linear_func(input, attn_qkvw, attn_qkvb if attn_qkvb is not None else attn_qkvw, attn_qkvb + is not None, do_flash_attn, config.heads, False, rope_theta) + + context_layer = triton_flash_attn_kernel(qkv_out[0], qkv_out[1], qkv_out[2], scale, + input.shape[-2] % 128 == 0) + context_layer = _transpose_for_context(context_layer[:, :, :, :head_size]) + + else: + do_flash_attn = False + if context is not None: + query = torch.matmul(input, attn_qw) + key = torch.matmul(context, attn_kw) + value = torch.matmul(context, attn_vw) + else: + qkv = torch.matmul(input, attn_qkvw) + query, key, value = qkv.chunk(3, dim=-1) + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + query, key, value = pad_transform_func(query, key, value, config.heads, do_flash_attn) + attention_scores = (torch.matmul(query, key.transpose(-1, -2)) * scale).softmax(dim=-1) + context_layer = _transpose_for_context(torch.matmul(attention_scores, value)) + + output = linear_func(context_layer, attn_ow, attn_ob, do_out_bias, False, config.heads, False, rope_theta) + return output + + output = selfAttention_fp(input, context, input_mask) + + return output + + @staticmethod + def backward(ctx, grad_output, grad_output1, grad_output2, grad_output3): + raise RuntimeError('You are running with DeepSpeed Inference mode. \ + Please switch to Training mode for running backward!') + + +class DeepSpeedDiffusersAttention(nn.Module): + """Initialize the DeepSpeed Transformer Layer. + Arguments: + layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, + layer_id will be 0,1,2...23 when each layer object is instantiated + config: An object of DeepSpeedInferenceConfig + """ + layer_id = 0 + + def __init__( + self, + config, + ): + super(DeepSpeedDiffusersAttention, self).__init__() + + self.config = config + self.config.layer_id = DeepSpeedDiffusersAttention.layer_id + DeepSpeedDiffusersAttention.layer_id += 1 + device = get_accelerator().current_device_name() if config.bigscience_bloom else 'cpu' + qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 + + data_type = self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype + + if DeepSpeedDiffusersAttention.layer_id == 1: + log_dist(f"DeepSpeed-Attention config: {self.config.__dict__}", [0]) + + self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, + qkv_size_per_partition, + dtype=data_type, + device=device), + requires_grad=False) + self.attn_kw = nn.Parameter(torch.empty(self.config.hidden_size, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + self.attn_vw = nn.Parameter(torch.empty(self.config.hidden_size, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + self.attn_qw = nn.Parameter(torch.empty(self.config.hidden_size, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), + requires_grad=False) + out_size_per_partition = self.config.hidden_size // self.config.mp_size + self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + + self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + self.do_out_bias = True + + if triton_flash_attn is None: + load_triton_flash_attn() + self.triton_flash_attn_kernel = triton_flash_attn() + self.num_attention_heads_per_partition = self.config.heads // self.config.mp_size + self.hidden_size_per_partition = self.config.hidden_size // self.config.mp_size + self.hidden_size_per_attention_head = self.config.hidden_size // self.config.heads + + self.norm_factor = math.sqrt(math.sqrt(self.config.hidden_size // self.config.heads)) + + if self.config.scale_attn_by_inverse_layer_idx is True: + self.norm_factor *= math.sqrt(self.config.layer_id + 1) + # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 + + self.workspace = WorkspaceOp(self.config) + self.score_context_func = SoftmaxContextOp(self.config) + self.linear_func = LinearOp(self.config) + self.pad_transform_func = PadTransformOp(self.config) + + def allocate_workspace(self, size): + # Allocate memory only on first layer forward + if self.config.layer_id == 0: + self.workspace.allocate_workspace(self.config.hidden_size, self.config.heads, size[1], size[0], + DeepSpeedDiffusersAttention.layer_id, self.config.mp_size, False, 0, + self.config.max_out_tokens, self.config.min_out_tokens) + + def forward(self, input, context=None, input_mask=None): + self.allocate_workspace(input.size()) + output = DeepSpeedDiffusersAttentionFunction.apply( + input, context, input_mask, self.config, self.attn_qkvw, self.attn_qw, self.attn_kw, self.attn_vw, + self.attn_qkvb, self.num_attention_heads_per_partition, self.norm_factor, self.hidden_size_per_partition, + self.attn_ow, self.attn_ob, self.do_out_bias, self.score_context_func, self.linear_func, + self.pad_transform_func, self.triton_flash_attn_kernel, self.config.rope_theta) + + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py new file mode 100644 index 0000000000000000000000000000000000000000..d01638f36e40141dc91276dcb484aeb96865048a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn as nn + +from deepspeed import module_inject +from .diffusers_attention import DeepSpeedDiffusersAttention +from .bias_add import nhwc_bias_add +from .diffusers_2d_transformer import Diffusers2DTransformerConfig +from deepspeed.utils.types import ActivationFuncType +from .op_binding.gated_activation import GatedActivationOp +from .op_binding.layer_norm import LayerNormOp + + +class DeepSpeedDiffusersTransformerBlock(nn.Module): + + def __init__(self, equivalent_module: nn.Module, config: Diffusers2DTransformerConfig): + super(DeepSpeedDiffusersTransformerBlock, self).__init__() + self.quantizer = module_inject.GroupQuantizer(q_int8=config.int8_quantization) + # Ensure ops are built by the time we start running + self.config = config + + self.ff1_w = self.quantizer.quantize( + nn.Parameter(equivalent_module.ff.net[0].proj.weight.data, requires_grad=False)) + self.ff1_b = nn.Parameter(equivalent_module.ff.net[0].proj.bias.data, requires_grad=False) + self.ff2_w = self.quantizer.quantize(nn.Parameter(equivalent_module.ff.net[2].weight.data, + requires_grad=False)) + self.ff2_b = nn.Parameter(equivalent_module.ff.net[2].bias.data, requires_grad=False) + + self.norm1_g = nn.Parameter(equivalent_module.norm1.weight.data, requires_grad=False) + self.norm1_b = nn.Parameter(equivalent_module.norm1.bias.data, requires_grad=False) + self.norm1_eps = equivalent_module.norm1.eps + + self.norm2_g = nn.Parameter(equivalent_module.norm2.weight.data, requires_grad=False) + self.norm2_b = nn.Parameter(equivalent_module.norm2.bias.data, requires_grad=False) + self.norm2_eps = equivalent_module.norm2.eps + + self.norm3_g = nn.Parameter(equivalent_module.norm3.weight.data, requires_grad=False) + self.norm3_b = nn.Parameter(equivalent_module.norm3.bias.data, requires_grad=False) + self.norm3_eps = equivalent_module.norm3.eps + + self.attn_1 = equivalent_module.attn1 + self.attn_2 = equivalent_module.attn2 + + # Pull the bias in if we can + if isinstance(self.attn_1, DeepSpeedDiffusersAttention): + self.attn_1.do_out_bias = False + self.attn_1_bias = self.attn_1.attn_ob + else: + self.attn_1_bias = nn.Parameter(torch.zeros_like(self.norm2_g), requires_grad=False) + + # Pull the bias in if we can + if isinstance(self.attn_2, DeepSpeedDiffusersAttention): + self.attn_2.do_out_bias = False + self.attn_2_bias = self.attn_2.attn_ob + else: + self.attn_2_bias = nn.Paramaeter(torch.zeros_like(self.norm3_g), requires_grad=False) + + self.gated_activation = GatedActivationOp() + self.layer_norm = LayerNormOp() + + def forward(self, hidden_states, context=None, timestep=None, **kwargs): + # In v0.12.0 of diffuser, several new kwargs were added. Capturing + # those with kwargs to maintain backward compatibility + + # In v0.11.0 of diffusers, the kwarg was changed from 'context' to 'encoder_hidden_states' + # This is so we can support older and newer versions of diffusers + if "encoder_hidden_states" in kwargs and kwargs["encoder_hidden_states"] is not None: + context = kwargs["encoder_hidden_states"] + + out_norm_1 = self.layer_norm(hidden_states, self.norm1_g, self.norm1_b, self.norm1_eps) + out_attn_1 = self.attn_1(out_norm_1) + + out_norm_2, out_attn_1 = self.layer_norm.layer_norm_residual_store_pre_ln_res( + out_attn_1, self.attn_1_bias, hidden_states, self.norm2_g, self.norm2_b, self.norm2_eps) + out_attn_2 = self.attn_2(out_norm_2, context=context) + out_norm_3, out_attn_2 = self.layer_norm.layer_norm_residual_store_pre_ln_res( + out_attn_2, self.attn_2_bias, out_attn_1, self.norm3_g, self.norm3_b, self.norm3_eps) + + out_ff1 = nn.functional.linear(out_norm_3, self.ff1_w) + out_geglu = self.gated_activation(out_ff1, self.ff1_b, ActivationFuncType.GATED_GELU) + + out_ff2 = nn.functional.linear(out_geglu, self.ff2_w) + return nhwc_bias_add(out_ff2, self.ff2_b, other=out_attn_2) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..24f710d224949a5b7684742b7ff7cf96eb007122 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_attention.py @@ -0,0 +1,301 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +import torch +import torch.nn as nn +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from .op_binding import LinearOp, VectorMatMulOp, SoftmaxContextOp, QKVGemmOp, SoftmaxOp + +minus_inf = -10000.0 + + +class DeepSpeedSelfAttention(nn.Module): + num_layers = 0 + _qkv_buffers = [] + + def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1): + super(DeepSpeedSelfAttention, self).__init__() + self.config = config + data_type = self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype + self.config.layer_id = DeepSpeedSelfAttention.num_layers + DeepSpeedSelfAttention.num_layers = DeepSpeedSelfAttention.num_layers + 1 + device = get_accelerator().current_device_name() #if config.bigscience_bloom else 'cpu' + if self.config.set_empty_params: + self.attn_qw = None + self.attn_qb = None + self.attn_kw = None + self.attn_kb = None + self.attn_vw = None + self.attn_vb = None + self.attn_qkvw = None + self.attn_qkvb = None + self.attn_ow = None + self.attn_ob = None + else: + qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 if config.num_kv < 0 else \ + ((self.config.heads + self.config.num_kv * 2) // self.config.mp_size) * (self.config.hidden_size // self.config.heads) + self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, + qkv_size_per_partition, + dtype=data_type, + device=device), + requires_grad=False) + self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), + requires_grad=False) + out_size_per_partition = self.config.hidden_size // self.config.mp_size + self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + + self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + + self.num_attention_heads_per_partition = self.config.heads // self.config.mp_size + self.num_kv_partition = self.config.num_kv // self.config.mp_size + self.hidden_size_per_partition = self.config.hidden_size // self.config.mp_size + self.hidden_size_per_attention_head = self.config.hidden_size // self.config.heads + + self.mp_group = mp_group + + # used for quantization + self.q_scales = q_scales + self.q_groups = q_groups + self.merge_count = int(math.log2(merge_count)) + + self.norm_factor = math.sqrt(self.config.hidden_size // self.config.heads) + if not config.use_mup: + self.norm_factor = math.sqrt(self.norm_factor) + + if self.config.scale_attn_by_inverse_layer_idx is True: + self.norm_factor *= math.sqrt(self.config.layer_id + 1) + # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 + + self.qkv_func = QKVGemmOp(config) + self.score_context_func = SoftmaxContextOp(config) + self.linear_func = LinearOp(config) + self.vector_matmul_func = VectorMatMulOp(config) + if len(DeepSpeedSelfAttention._qkv_buffers) == 0: + DeepSpeedSelfAttention._qkv_buffers = [ + torch.empty(self.hidden_size_per_partition * 3, + self.config.hidden_size, + dtype=data_type_fp, + device=device), + torch.empty(self.hidden_size_per_partition * 3, dtype=data_type_fp, device=device) + ] + + def compute_attention(self, qkv_out, input_mask, layer_past, alibi, is_prompt, token_idx, position_ids): + if isinstance(qkv_out, list) or isinstance(qkv_out, tuple): + qkv_out = qkv_out[0] + + no_masking = input_mask is None or input_mask is False + + if no_masking: + input_mask = torch.empty(1) + + attn_key_value = self.score_context_func( + query_key_value=qkv_out, + attn_mask=((1 - input_mask).to(qkv_out.dtype) * + minus_inf) if input_mask.dtype == torch.int64 else input_mask, + heads=self.num_attention_heads_per_partition, + num_kv=self.num_kv_partition, + norm_factor=(1 / self.norm_factor if self.config.scale_attention else 1.0), + no_masking=no_masking, + layer_id=self.config.layer_id, + num_layers=DeepSpeedSelfAttention.num_layers, + alibi=alibi, + is_prompt=is_prompt, + token_idx=token_idx, + position_ids=position_ids) + + context_layer, key_layer, value_layer = attn_key_value + return context_layer, key_layer, value_layer + + def _merge_qkv(self): + qvkw = DeepSpeedSelfAttention._qkv_buffers[0] + qvkw[:self.hidden_size_per_partition, :] = self.attn_qw # type: ignore + qvkw[self.hidden_size_per_partition:2 * self.hidden_size_per_partition, :] = self.attn_kw # type: ignore + qvkw[2 * self.hidden_size_per_partition:, :] = self.attn_vw # type: ignore + if self.attn_qb is not None: + qvkb = DeepSpeedSelfAttention._qkv_buffers[1] + qvkb[:self.hidden_size_per_partition] = self.attn_qb + qvkb[self.hidden_size_per_partition:2 * self.hidden_size_per_partition] = self.attn_kb # type: ignore + qvkb[2 * self.hidden_size_per_partition:] = self.attn_vb # type: ignore + return DeepSpeedSelfAttention._qkv_buffers + + def forward(self, + input, + input_mask, + head_mask=None, + layer_past=None, + get_present=False, + encoder_hidden_states=None, + encoder_attention_mask=None, + output_attentions=False, + norm_w=None, + norm_b=None, + alibi=None, + **kwargs): + if self.attn_qkvw is None: + self._attn_qkvw, self._attn_qkvb = self._merge_qkv() + else: + self._attn_qkvw = self.attn_qkvw + self._attn_qkvb = self.attn_qkvb + if not self.config.pre_layer_norm: + qkv_out = self.linear_func(input=input, + weight=self._attn_qkvw, + bias=self._attn_qkvb, + add_bias=self.attn_qkvb is not None, + do_flash_attn=False, + num_heads=self.num_attention_heads_per_partition, + num_layers=DeepSpeedSelfAttention.num_layers) + else: + qkv_out = self.qkv_func(input=input, + weight=self._attn_qkvw, + bias=self._attn_qkvb, + gamma=norm_w, + beta=norm_b) + + is_prompt = kwargs.get("first_token", qkv_out[0].shape[1] > 1) + token_idx = kwargs.get("token_idx", None) + position_ids = kwargs.get("position_ids", None) + + context_layer, key_layer, value_layer = self.compute_attention(qkv_out=qkv_out, + input_mask=input_mask, + layer_past=layer_past, + alibi=alibi, + is_prompt=is_prompt, + token_idx=token_idx, + position_ids=position_ids) + + output = self.vector_matmul_func(input=context_layer, weight=self.attn_ow) + inp_norm = qkv_out[-1] + + if self.config.mlp_after_attn and self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: + dist.all_reduce(output, group=self.mp_group) + return (output, key_layer, value_layer, context_layer, inp_norm) + + +class BloomSelfAttention(DeepSpeedSelfAttention): + + def __init__(self, *args, **kwargs): + super(BloomSelfAttention, self).__init__(*args, **kwargs) + self.softmax_func = SoftmaxOp(self.config) + + ########### This part is taken/modified form the HF modeling_bloom.py ################ + # Reference: https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py + + def _transpose_for_context(self, x): + x = x.permute(0, 2, 1, 3).contiguous() + new_x_layer_shape = x.size()[:-2] + \ + (self.hidden_size_per_partition,) + return x.view(*new_x_layer_shape).contiguous() + + def _split_tensor_along_last_dim(self, tensor, num_partitions, contiguous_split_chunks=True): + """Split a tensor along its last dimension. + + Args: + tensor: ([`torch.tensor`], *required*): + input tensor to split + num_partitions ([`int`], *required*): + number of partitions to split the tensor + contiguous_split_chunks ([`bool`], *optional*, default=`False`):: + If True, make each chunk contiguous in memory. + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + numerator, denominator = tensor.size()[last_dim], num_partitions + if not (numerator % denominator == 0): + raise ValueError(f"{numerator} is not divisible by {denominator}") + last_dim_size = numerator // denominator + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # Note: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + def compute_attention(self, qkv_out, input_mask, layer_past, alibi, is_prompt, token_idx, position_ids): + if isinstance(qkv_out, list) or isinstance(qkv_out, tuple): + qkv_out = qkv_out[0] + + no_masking = input_mask is None + + if no_masking: + input_mask = torch.empty(1) + + mixed_x_layer = qkv_out + alibi = alibi.to(get_accelerator().current_device_name()) + head_dim = self.hidden_size_per_partition // self.num_attention_heads_per_partition + new_tensor_shape = mixed_x_layer.size()[:-1] + (self.num_attention_heads_per_partition, 3 * head_dim) + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + query_layer, key_layer, value_layer = self._split_tensor_along_last_dim(mixed_x_layer, 3) + + # [batch_size, head_dim, q_length, k_length] + output_size = (query_layer.size(0), query_layer.size(2), query_layer.size(1), key_layer.size(1)) + # [batch_size, q_length, num_heads, head_dim] -> [q_length, batch_size * num_heads, head_dim] + query_layer = query_layer.transpose(1, 2).reshape(output_size[0] * output_size[1], output_size[2], -1) + # [batch_size, k_length, num_heads, head_dim] -> [k_length, batch_size * num_heads, head_dim] + key_layer = key_layer.transpose(1, 2).reshape(output_size[0] * output_size[1], output_size[3], + -1).transpose(-1, -2) + value_layer = value_layer.transpose(1, 2).reshape(output_size[0] * output_size[1], output_size[3], -1) + if layer_past is not None: + past_key, past_value = layer_past + # concatenate along seq_length dimension -> [batch_size, qk_length, num_heads, head_dim] + key_layer = torch.cat((past_key.type_as(key_layer), key_layer), dim=-1) + value_layer = torch.cat((past_value.type_as(value_layer), value_layer), dim=-2) + + presents = (key_layer, value_layer) + # Raw attention scores. [batch_size * num_heads, q_length, k_length] + matmul_result = torch.matmul(query_layer, key_layer) + # change view to [batch_size, num_heads, q_length, k_length] + attention_scores = matmul_result.view(output_size[0], output_size[1], output_size[2], -1) + + offset = dist.get_rank() * self.num_attention_heads_per_partition if dist.is_initialized() else 0 + target_dtype = torch.float16 if self.config.dtype == torch.int8 else self.config.dtype + + # When using the hybrid engine with BLOOM, input_mask needs to be converted from torch.bool -> torch.int64 + if input_mask.dtype == torch.bool: + input_mask = input_mask.long() + + # Invert input_mask per transformer implementation (eg, in BLOOM, it's already inverted) + if self.config.invert_mask: + input_mask = 1 - input_mask + + attention_probs = self.softmax_func(attn_scores=attention_scores, + attn_mask=input_mask.to(target_dtype) * minus_inf, + alibi=alibi, + triangular=(self.config.triangular_masking + and (attention_scores.shape[-2] > 1)), + recompute=False, + local_attention=False, + window_size=1, + async_op=False, + layer_scale=1 / (self.norm_factor * self.norm_factor), + head_offset=offset) + + # change view [batch_size x num_heads, q_length, k_length] + attention_probs_reshaped = attention_probs.view(*matmul_result.shape) + + # matmul: [batch_size * num_heads, q_length, head_dim] + context_layer = torch.bmm(attention_probs_reshaped, value_layer) + + # change view [batch_size, num_heads, q_length, head_dim] + context_layer = context_layer.view( + context_layer.size(0) // self.num_attention_heads_per_partition, self.num_attention_heads_per_partition, + context_layer.size(1), context_layer.shape[-1]) + + context_layer = self._transpose_for_context(context_layer) + key_layer = presents[0] + value_layer = presents[1] + + return context_layer, key_layer, value_layer + + ###################### End of HF modeling_bloom addition ######################## diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..36de06db920fb51adf31809e02d74a797494938c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +import torch +import torch.nn as nn +from deepspeed import comm as dist +from deepspeed.utils.types import GATED_ACTIVATION_TYPES +from deepspeed.accelerator import get_accelerator +from .op_binding import MLPGemmOp, VectorMatMulOp, GELUGemmOp, ResidualAddOp + + +class DeepSpeedMLP(nn.Module): + _inter_w_buffers = [] + + def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False): + super(DeepSpeedMLP, self).__init__() + + self.config = config + + data_type = torch.int8 if self.config.dtype == torch.int8 else self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype + device = get_accelerator().current_device_name() + + proj_factor = 2 if self.config.mlp_act_func_type in GATED_ACTIVATION_TYPES else 1 + self.config.intermediate_size = self.config.intermediate_size if self.config.intermediate_size > 0 else 4 * self.config.hidden_size + self.intm_w_sz_per_partition = self.config.intermediate_size * proj_factor // self.config.mp_size + self.intm_o_sz_per_partition = self.config.intermediate_size // self.config.mp_size + + if self.config.set_empty_params: + self.attn_nw = None + self.attn_nb = None + self.inter_w = None + self.inter_b = None + self.inter_up_w = None + self.inter_up_b = None + self.inter_gate_w = None + self.inter_gate_b = None + self.output_w = None + self.output_b = None + else: + self.attn_nw = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + self.attn_nb = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + + self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, + self.intm_w_sz_per_partition, + dtype=data_type, + device=device), + requires_grad=False) + self.inter_b = nn.Parameter(torch.empty(self.intm_w_sz_per_partition, dtype=data_type_fp, device=device), + requires_grad=False) + self.output_w = nn.Parameter(torch.empty(self.intm_o_sz_per_partition, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + self.output_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + + # used for quantization + self.q_scales = q_scales + self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups + self.merge_count = int(math.log2(merge_count)) + self.mp_group = mp_group + + self.mlp_gemm_func = MLPGemmOp(config) + self.vector_matmul_func = VectorMatMulOp(config) + self.fused_gemm_gelu = GELUGemmOp(config) + self.residual_add_func = ResidualAddOp(config) + + if len(DeepSpeedMLP._inter_w_buffers) == 0: + DeepSpeedMLP._inter_w_buffers = [ + torch.empty(self.intm_w_sz_per_partition, self.config.hidden_size, dtype=data_type, device=device), + torch.empty(self.intm_w_sz_per_partition, dtype=data_type_fp, device=device) + ] + + def _merge_inter_w(self): + inter_w = DeepSpeedMLP._inter_w_buffers[0] + inter_w[:self.intm_w_sz_per_partition // 2, :] = self.inter_up_w # type: ignore + inter_w[self.intm_w_sz_per_partition // 2:, :] = self.inter_gate_w # type: ignore + if self.inter_up_b is not None: + inter_b = DeepSpeedMLP._inter_w_buffers[1] + inter_b[:self.intm_w_sz_per_partition // 2] = self.inter_up_b # type: ignore + inter_b[self.intm_w_sz_per_partition // 2:] = self.inter_gate_b # type: ignore + return DeepSpeedMLP._inter_w_buffers + + def forward(self, input, residual, residual_norm, bias): + if self.inter_w is None: + self._inter_w, self._inter_b = self._merge_inter_w() + else: + self._inter_w = self.inter_w + self._inter_b = self.inter_b + + residual_add = None + if self.attn_nw is None: + output = self.fused_gemm_gelu(input=residual_norm, + weight=self._inter_w, + bias=self._inter_b, + weight_out=self.output_w) + else: + output, residual_add = self.mlp_gemm_func(input=input, + residual=residual, + weight_interm=self._inter_w, + weight_out=self.output_w, + input_bias=bias, + bias=self._inter_b, + gamma=self.attn_nw, + beta=self.attn_nb) + + residual = self.residual_add_func(hidden_state=output, + residual=residual, + add_bias=bias is not None, + attention_output=input, + attention_bias=bias if bias is not None else self.output_b, + final_bias=self.output_b, + residual_add=residual_add) + if self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: + dist.all_reduce(residual, group=self.mp_group) + + return residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/moe_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/moe_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..3a9785985d196704abe162cd6cef6f75438c7e7a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/moe_inference.py @@ -0,0 +1,350 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import json +import math +import torch +from torch.autograd import Function +import torch.nn as nn +from .ds_attention import DeepSpeedSelfAttention +from .config import DeepSpeedInferenceConfig +from .op_binding import SoftmaxOp, VectorMatMulOp, GELUGemmOp +from .op_binding.bias_residual import BiasResidualOp +from .op_binding.einsum_sec_sm_ecm import EinsumSecSmEcmOp +from .op_binding.layer_norm import LayerNormOp +from ....moe.sharded_moe import TopKGate +from deepspeed import comm as dist +from .op_binding.moe_res_matmul import MoEResMatmulOp + + +class DeepSpeedMoEInferenceConfig(DeepSpeedInferenceConfig): + """Initialize the DeepSpeed Transformer Config. + Arguments: + hidden_size: The hidden size of the transformer layer + intermediate_size: The intermediate size of the feed-forward part of transformer layer + heads: The number of heads in the self-attention of the transformer layer + num_hidden_layers: The number of transformer layers + layer_norm_eps: The epsilon value for the layer norm + local_rank: Optional: The rank of GPU running the transformer kernel, it is not required + to use if the model already set the current device, otherwise need to set it + so that the transformer kernel can work on the right device + mp_size (optional): This argument is mainly used to create the parameters on the kernel side + using model-parallel architecture. If the client model already takes care of this, there is no + need to pass this argument. + fp16: Enable half-precision computation + bf16: Enable bf16 floating point computation + pre_layer_norm: Select between Pre-LN or Post-LN transformer architecture + stochastic_mode: Enable for high performance, please note that this flag has some level of + non-determinism and can produce different results on different runs. However, we have seen + that by enabling it, the pretraining tasks such as BERT are not affected and can obtain + a high accuracy level. On the other hand, for the downstream tasks, such as fine-tuning, we recommend + to turn it off in order to be able to reproduce the same result through the regular kernel execution. + + scale_attention: If true, both q and k are scaled by 1/sqrt(attention_heads) before attention computation. + return_tuple: if True, returns the transformer output as a tuple, otherwise returns as a tensor + """ + + def __init__(self, + hidden_size=-1, + intermediate_size=-1, + heads=-1, + num_hidden_layers=-1, + layer_norm_eps=1e-12, + local_rank=-1, + mp_size=1, + fp16=False, + bf16=False, + q_int8=False, + pre_layer_norm=True, + stochastic_mode=False, + scale_attention=True, + triangular_masking=True, + local_attention=False, + window_size=256, + return_tuple=True, + moe_experts=1, + global_experts=1, + k=1, + capacity_factor=1., + eval_capacity_factor=1., + min_capacity=1, + noisy_gate_policy=None, + drop_tokens=True, + use_rts=False, + mlp_type='standard', + scale_attn_by_inverse_layer_idx=False): + super(DeepSpeedMoEInferenceConfig, + self).__init__(hidden_size, (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, + num_hidden_layers, layer_norm_eps, local_rank, mp_size, fp16, bf16, q_int8, + pre_layer_norm, stochastic_mode, scale_attention, triangular_masking, local_attention, + window_size, return_tuple) + self.moe_experts = moe_experts + self.k = k + self.capacity_factor = capacity_factor + self.eval_capacity_factor = eval_capacity_factor + self.min_capacity = min_capacity + self.noisy_gate_policy = noisy_gate_policy + self.drop_tokens = drop_tokens + self.use_rts = use_rts + self.global_experts = global_experts + self.mlp_type = mlp_type + self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx + + @classmethod + def from_dict(cls, json_object): + config = DeepSpeedInferenceConfig() + for key, value in json_object.items(): + config.__dict__[key] = value + return config + + @classmethod + def from_json_file(cls, json_file): + with open(json_file, "r", encoding='utf-8') as reader: + text = reader.read() + return cls.from_dict(json.loads(text)) + + +class DeepSpeedMLPFunction(Function): + + @staticmethod + def forward(ctx, input, inter_w, inter_b, config, output_b, output_w, q_scales, q_groups, merge_count, mp_group, + async_op, gelu_gemm_func, vector_matmul_func): + if config.q_int8: + intermediate = gelu_gemm_func(input, inter_w, inter_b, config.epsilon, q_scales[2], + (q_groups * (2**merge_count)), config.pre_layer_norm) + output = vector_matmul_func(intermediate, output_w, q_scales[3], q_groups, (merge_count)) + else: + output = gelu_gemm_func(input, inter_w, inter_b, output_w, config.epsilon, config.pre_layer_norm, async_op) + if mp_group is not None and dist.get_world_size(group=mp_group) > 1: + dist.all_reduce(output, group=mp_group, async_op=async_op) + + return output + output_b + + @staticmethod + def backward(ctx, grad_output): + raise RuntimeError('You are running with DeepSpeed Inference mode. \ + Please switch to Training mode for running backward!') + + +class DeepSpeedMoEMLP(nn.Module): + + def __init__(self, config, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False, mp_group=None): + super(DeepSpeedMoEMLP, self).__init__() + + self.config = config + self.attn_nw = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.attn_nb = nn.Parameter(torch.Tensor(self.config.hidden_size)) + interm_size = self.config.intermediate_size // (1 if mp_group is None else dist.get_world_size(group=mp_group)) + self.inter_w = nn.Parameter(torch.Tensor(self.config.hidden_size, interm_size)) + self.inter_b = nn.Parameter(torch.Tensor(interm_size)) + self.output_w = nn.Parameter(torch.Tensor((interm_size), self.config.hidden_size)) + self.output_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) + + # used for quantization + self.q_scales = q_scales + self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups + self.merge_count = int(math.log2(merge_count)) + self.mp_group = mp_group + self.gelu_gemm_func = GELUGemmOp(self.config) + self.vector_matmul_func = VectorMatMulOp(self.config) + + def forward(self, input, async_op=False): + return DeepSpeedMLPFunction.apply(input, self.inter_w, self.inter_b, self.config, self.output_b, self.output_w, + self.q_scales, self.q_groups, self.merge_count, self.mp_group, async_op, + self.gelu_gemm_func, self.vector_matmul_func) + + +class DeepSpeedMoEInference(nn.Module): + """Initialize the DeepSpeed MoE Transformer Layer. + Arguments: + layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, + layer_id will be 0,1,2...23 when each layer object is instantiated + config: An object of DeepSpeedInferenceConfig + mp_group: Model parallelism group initialized on the modeling side. + quantize_scales: This argument groups all the layers' scales used for quantization + quantize_groups: Number of groups used for quantizing the model + merge_count: Shows the number of model-parallel checkpoints merged before running inference. + We use this argument to control the quantization scale for the model parameters if a bigger + quantize-grouping than 1 is used. + mlp_extra_grouping: This flag is used to show a 2x higher number of groups used for the MLP part + of a Transformer layer. We use this feature for quantization to reduce the convergence impact + for specific downstream tasks. + """ + layer_id = 0 + + def __init__(self, + config, + mp_group=None, + ep_group=None, + expert_mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super(DeepSpeedMoEInference, self).__init__() + + self.config = config + self.config.layer_id = DeepSpeedMoEInference.layer_id + + assert self.config.dtype != torch.bfloat16, "DeepSpeed MoE Transformer Inference not yet tested for bfloat support" + + DeepSpeedMoEInference.layer_id += 1 + self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count) + self.attn_nw = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.attn_nb = nn.Parameter(torch.Tensor(self.config.hidden_size)) + + self.norm_w = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.norm_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) + + if config.mlp_type == 'residual': + self.res_mlp = DeepSpeedMoEMLP(config, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping, + mp_group) + self.res_coef = nn.Parameter(torch.Tensor(self.config.hidden_size, 2)) + self.coef_func = SoftmaxOp(self.config) + self.vector_matmul_func = VectorMatMulOp(self.config) + + config.mp_size = 1 + self.mlp = nn.ModuleList( + DeepSpeedMoEMLP(config, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping, expert_mp_group) + for i in range(self.config.moe_experts)) + + self.moe_gate = TopKGate(self.config.hidden_size, self.config.global_experts, self.config.k, + self.config.capacity_factor, self.config.eval_capacity_factor, + self.config.min_capacity, self.config.noisy_gate_policy, self.config.drop_tokens, + self.config.use_rts, self.ep_group) + + self.ep_group = ep_group + self.mp_group = mp_group + self.expert_mp_group = expert_mp_group + + print("DeepSpeed MoE Transformer Inference config is ", self.config.__dict__) + + self.bias_residual_func = BiasResidualOp(self.config) + self.ds_layernorm = LayerNormOp(self.config) + self.einsum_sec_sm_ecm = EinsumSecSmEcmOp(self.config) + self.moe_res_matmul = MoEResMatmulOp(self.config) + + def res_coef_func(self, inp, async_op): + inp = self.vector_matmul_func(inp, self.res_coef, async_op) + return self.coef_func(inp, torch.empty(1), False, False, False, 256, async_op) + + def moe_gate_einsum(self, attention_output): + _, combined_weights, dispatch_mask, _ = self.moe_gate( + attention_output.view(-1, self.config.hidden_size), + None, + ) + dispatched_attention = self.einsum_sec_sm_ecm(dispatch_mask.type_as(attention_output), + attention_output.view(-1, self.config.hidden_size)) + return dispatched_attention, combined_weights + + def expert_exec(self, dispatched_input): + dispatched_input = dispatched_input.reshape(self.config.global_experts // self.config.moe_experts, + self.config.moe_experts, -1, self.config.hidden_size) + + chunks = dispatched_input.chunk(self.config.moe_experts, dim=1) + expert_outputs = torch.empty(( + self.config.moe_experts, + chunks[0].shape[0], + ) + chunks[0].shape[2:], + dtype=dispatched_input.dtype, + device=dispatched_input.device) + for chunk, expert in zip(chunks, range(len(self.mlp))): + expert_outputs[expert] = self.mlp[expert](chunk.view(-1, dispatched_input.shape[-2], + dispatched_input.shape[-1])) + return expert_outputs + + def _alltoall(self, dispatched_attention): + if dist.get_world_size(group=self.ep_group) > 1: + dispatched_input = torch.empty_like(dispatched_attention) + dist.all_to_all_single(dispatched_input, dispatched_attention, group=self.ep_group) + return dispatched_input + else: + return dispatched_attention + + def scale_expert_output(self, attention_output, expert_output, combined_weights): + combined_output = torch.matmul( + combined_weights.type_as(attention_output).reshape(combined_weights.shape[0], -1), + expert_output.reshape(-1, expert_output.shape[-1])) + return combined_output.reshape(attention_output.shape) + + def forward(self, + input, + input_mask=None, + attention_mask=None, + head_mask=None, + layer_past=None, + get_key_value=False, + get_present=False, + encoder_output=None, + enc_dec_attn_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + use_cache=False, + output_attentions=False): + get_present = (get_present or get_key_value or use_cache) + input_mask = input_mask if attention_mask is None else attention_mask + input_type = input.dtype + + if (self.config.dtype in [torch.float16, torch.int8]) and input_type == torch.float: + input = input.half() + + with torch.no_grad(): + attention_output = self.attention(input, input_mask, head_mask, layer_past, get_present, + encoder_hidden_states, encoder_attention_mask, output_attentions, + self.norm_w, self.norm_b) + + if get_present: + attention_output, p_key, p_value = attention_output[0:3] + presents = (p_key, p_value) + elif output_attentions: + attention_output, _, _, context_output = attention_output[0:4] + else: + attention_output = attention_output[0] + + residual_add = attention_output + self.attention.attn_ob + attention_output = self.ds_layernorm(residual_add, self.attn_nw, self.attn_nb, self.config.epsilon) + + if self.config.mlp_type == 'residual': + res_mlp_out = self.res_mlp(attention_output, async_op=True) + res_coef_out = self.res_coef_func(attention_output, async_op=True) + + if self.expert_mp_group is not None: + world_size = dist.get_world_size(group=self.expert_mp_group) + gather_buffer = torch.empty(world_size * attention_output.numel(), + dtype=attention_output.dtype, + device=attention_output.device) + dist.all_gather_into_tensor(gather_buffer, attention_output, group=self.expert_mp_group) + attention_output = gather_buffer.view(-1, *attention_output.size()[1:]) + + ############## MoE Gating + Experts ############### + dispatched_attention, combined_weights = self.moe_gate_einsum(attention_output) + dispatched_input = self._alltoall(dispatched_attention) + expert_outputs = self.expert_exec(dispatched_input) + expert_output = self._alltoall(expert_outputs) + output = self.scale_expert_output(attention_output, expert_output, combined_weights) + ################################################ + + if self.expert_mp_group is not None: + output = output.split(output.shape[0] // dist.get_world_size(group=self.expert_mp_group), + dim=0)[dist.get_rank(group=self.expert_mp_group)] + + if self.config.mlp_type == 'residual': + self.moe_res_matmul(res_mlp_out, res_coef_out, output) + + output = self.bias_residual_func(output, residual_add, torch.empty(1)) + + if not self.config.pre_layer_norm: + output = self.ds_layernorm(output, self.norm_w, self.norm_b, self.config.epsilon) + + if input_type != output.dtype: + output = output.to(input_type) + + if get_present: + output = (output, presents) + + if self.config.return_tuple: + return output if type(output) is tuple else (output, ) + else: + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..20b7bf12a917865a753d9db41851562edc8cb337 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .linear import LinearOp +from .vector_matmul import VectorMatMulOp +from .softmax_context import SoftmaxContextOp +from .qkv_gemm import QKVGemmOp +from .softmax import SoftmaxOp +from .mlp_gemm import MLPGemmOp +from .gelu_gemm import GELUGemmOp +from .residual_add import ResidualAddOp diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..012102c60704f4047de723a1ff4367950069d4ed Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f04241896075d4bcf6162c8a52e2c08bfc61b03a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2e81d0dfd5b72c207377fce956919de089fb482 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e6c8db957654add351ab1a0390932618365fea6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7d8f33064ee43aa838b594ec35db9df8a430f60 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b81cbd18ecc0f49e4cbfba4d85692d3688710ba Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb411252369571c7b94392eec6edc73d92e31417 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c07b4e5659091b4d14bc9f85f3492ca6d2875e6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58ca4bc87ea6eb511de28ef9dd5f01cabd833e58 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1864d0c499379053f6989b479282bbd6ded512f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..624c384dd3f8b08673a3af0c09d34d0e989afb1c Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df781f3207bbf022318d81b351dde3843fcc9114 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d57f9436b5b96b3534b27a226b2622192916c37 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33d94db7b710e52e4de0027303868526ff3372f9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac6fd8ce74310f49cbb928c24c23071a288dd1e8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e680b4c826475e552278871f7bae0b8059e1ba9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1bf4da17454efbbc98ec993561112b4d797f75d5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f99b0c38a99b83657f88899c9648c15f04a3dd59 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9debe159dc3f10ac29b6b28b7534f5c17d952d6f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76f4d12318e319bd0845b838333bf0f753dcb22a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e7e5d7860166e9eaa2ca1639044c95be466234e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b84447b84ffb34d479c4561d005a54327d3a7f2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b901a6197f3083afe5d889a2c764cc9b7b36db18 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/base.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/base.py new file mode 100644 index 0000000000000000000000000000000000000000..5a997f95d5cc53e3323180ad7cd70ed9e8210a1b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/base.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig + +from deepspeed.ops.op_builder import InferenceBuilder + + +class BaseOp(torch.nn.Module): + inference_module = None + + def __init__(self, config: DeepSpeedInferenceConfig): + super(BaseOp, self).__init__() + self.config = config + if BaseOp.inference_module is None: + builder = InferenceBuilder() + BaseOp.inference_module = builder.load() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_add.py new file mode 100644 index 0000000000000000000000000000000000000000..d2ae38f546eb431ccd368b5da730cfe6a901d431 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_add.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class BiasAddOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(BiasAddOp, self).__init__(config) + + try: + if self.config.dtype == torch.float16: + self.bias_add_func = self.inference_module.bias_add_fp16 + elif self.config.dtype == torch.bfloat16: + self.bias_add_func = self.inference_module.bias_add_bf16 + else: + self.bias_add_func = self.inference_module.bias_add_fp32 + except AttributeError: + self.bias_add_func = self.bias_add_fallback + + @classmethod + def bias_add_fallback(cls, input, bias): + return torch.add(input, bias) + + def forward(self, activation: torch.Tensor, bias: torch.Tensor): + return self.bias_add_func(activation, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_gelu.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_gelu.py new file mode 100644 index 0000000000000000000000000000000000000000..f0fee0b0d06ea906d3c3229d5b3b27cd15c4f2c3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_gelu.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class BiasGeluOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(BiasGeluOp, self).__init__(config) + + try: + if self.config.dtype == torch.float16: + self.bias_gelu_func = self.inference_module.bias_gelu_fp16 + elif self.config.dtype == torch.bfloat16: + self.bias_gelu_func = self.inference_module.bias_gelu_bf16 + else: + self.bias_gelu_func = self.inference_module.bias_gelu_fp32 + except AttributeError: + self.bias_gelu_func = self.bias_gelu_fallback + + @classmethod + def bias_gelu_fallback(cls, activations, bias): + # Expected behavior is that of casting to float32 internally and using the tanh approximation + return F.gelu(activations.to(torch.float32) + bias.to(torch.float32), approximate='tanh').to(activations.dtype) + + def forward(self, activation: torch.Tensor, bias: torch.Tensor): + return self.bias_gelu_func(activation, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_relu.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_relu.py new file mode 100644 index 0000000000000000000000000000000000000000..ccfade1d9524562fb30f9d758cadd33191f53a45 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_relu.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class BiasReluOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(BiasReluOp, self).__init__(config) + + try: + if self.config.dtype == torch.float16: + self.bias_relu_func = self.inference_module.bias_relu_fp16 + elif self.config.dtype == torch.bfloat16: + self.bias_relu_func = self.inference_module.bias_relu_bf16 + else: + self.bias_relu_func = self.inference_module.bias_relu_fp32 + except AttributeError: + self.bias_relu_func = self.bias_relu_fallback + + @classmethod + def bias_relu_fallback(cls, activations, bias): + # Expected behavior is that of casting to float32 internally + return F.relu(activations.to(torch.float32) + bias.to(torch.float32)).to(activations.dtype) + + def forward(self, activation: torch.Tensor, bias: torch.Tensor): + return self.bias_relu_func(activation, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_residual.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_residual.py new file mode 100644 index 0000000000000000000000000000000000000000..ecad50e10ffef489eea41547c036124df057acb2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_residual.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class BiasResidualOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(BiasResidualOp, self).__init__(config) + + try: + if self.config.dtype in [torch.float16, torch.int8]: + self.bias_residual_func = self.inference_module.bias_residual_fp16 + else: + self.bias_residual_func = self.inference_module.bias_residual_fp32 + except AttributeError: + self.bias_residual_func = self.bias_residual_fallback + + @classmethod + def bias_residual_fallback(cls, output, residual, bias): + raise NotImplementedError("bias residual fallback isn't implemented") + + def forward(self, output, residual, bias): + return self.bias_residual_func(output, residual, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py new file mode 100644 index 0000000000000000000000000000000000000000..f34b10f786d11b87c0c8943ecba375c079451484 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class EinsumSecSmEcmOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(EinsumSecSmEcmOp, self).__init__(config) + + try: + if self.config.dtype in [torch.float16, torch.int8]: + self.einsum_sec_sm_ecm_func = self.inference_module.einsum_sec_sm_ecm_fp16 + else: + self.einsum_sec_sm_ecm_func = self.inference_module.einsum_sec_sm_ecm_fp32 + except AttributeError: + self.einsum_sec_sm_ecm_func = self.einsum_sec_sm_ecm_fallback + + @classmethod + def einsum_sec_sm_ecm_fallback(cls, Q, W): + raise NotImplementedError("einsum sec sm ecm fallback isn't implemented") + + def forward(self, Q, W): + return self.einsum_sec_sm_ecm_func(Q, W) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gated_activation.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gated_activation.py new file mode 100644 index 0000000000000000000000000000000000000000..d28d818ce4b33fbe5429ed9dd2c3e458baa52316 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gated_activation.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from deepspeed.utils.types import ActivationFuncType +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class GatedActivationOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + super(GatedActivationOp, self).__init__(config) + try: + self.gated_activation_func = self.inference_module.gated_activation + except AttributeError: + self.gated_activation_func = self.gated_activation_fallback + + @classmethod + def gated_activation_fallback(cls, activation, bias, activation_func_type): + # Expected behavior is that of casting to float32 internally + # Explicitly using the default GeLU + activation_func = None + activations = activation + bias.reshape(1, 1, -1) + hidden_states, gate = activations.chunk(2, dim=-1) + + if activation_func_type == ActivationFuncType.GATED_SILU: + activation_func = F.silu + elif activation_func_type == ActivationFuncType.GATED_GELU: + activation_func = F.gelu + + return hidden_states * activation_func(gate.to(torch.float32)).to(activations.dtype) + + def forward(self, activation: torch.Tensor, bias: torch.Tensor, activation_func_type: ActivationFuncType): + return self.gated_activation_func(activation, bias, activation_func_type) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..60bbb4b48bdb3977f037ac56f5a3008e3bea2af0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py @@ -0,0 +1,52 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +import deepspeed + + +class GELUGemmOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(GELUGemmOp, self).__init__(config) + try: + if self.config.dtype == torch.int8: + self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_int8 + elif self.config.dtype == torch.float16: + if deepspeed.HAS_TRITON and self.config.use_triton and self.config.dtype == torch.float16: + from deepspeed.ops.transformer.inference.triton.ops import fused_gemm_gelu as _triton_fused_gemm_gelu + self.fused_gemm_gelu = _triton_fused_gemm_gelu # type: ignore + else: + self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_bf16 # type: ignore + else: + self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_fp32 # type: ignore + except AttributeError: + self.fused_gemm_gelu = self.gelu_gemm_fallback + + def gelu_gemm_fallback(self, input, weight, scale, bias, out, out_scale, dtype, transpose): + tmp = torch.matmul(input, weight) + tmp = F.gelu(tmp.to(torch.float32) + bias.to(torch.float32), approximate="tanh").to(tmp.dtype) + output = torch.matmul(tmp, out) + + return output + + def forward(self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, weight_out: torch.Tensor): + + output = self.fused_gemm_gelu( + input, + weight, + weight.scale if hasattr(weight, 'scale') else torch.empty(1), # type: ignore + bias, + weight_out, + weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore + self.config.dtype == torch.int8, + self.config.transposed_mode) + + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/layer_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/layer_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..31219a58ac3c820ddb2d9a301860274ad84e5755 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/layer_norm.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class LayerNormOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + super(LayerNormOp, self).__init__(config) + try: + if config is None: + self.layer_norm_func = self.inference_module.layer_norm + elif self.config.dtype in [torch.float16, torch.int8]: + self.layer_norm_func = self.inference_module.layer_norm_fp16 + else: + self.layer_norm_func = self.inference_module.layer_norm_fp32 + except AttributeError: + self.layer_norm_func = self.layer_norm_fallback + + @classmethod + def layer_norm_residual(cls, vals, bias, res, gamma, beta, epsilon): + channels = gamma.shape[0] + dtype = gamma.dtype + vals_f = vals.to(torch.float32) + bias_f = bias.to(torch.float32).reshape(1, 1, -1) + res_f = res.to(torch.float32) + gamma_f = gamma.to(torch.float32) + beta_f = beta.to(torch.float32) + return F.layer_norm(vals_f + bias_f + res_f, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) + + @classmethod + def layer_norm_residual_store_pre_ln_res(cls, vals, bias, res, gamma, beta, epsilon): + channels = gamma.shape[0] + dtype = gamma.dtype + vals_f = vals.to(torch.float32) + bias_f = bias.to(torch.float32).reshape(1, 1, -1) + res_f = res.to(torch.float32) + gamma_f = gamma.to(torch.float32) + beta_f = beta.to(torch.float32) + res_output = vals_f + bias_f + res_f + norm_output = F.layer_norm(res_output, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) + return norm_output, res_output.to(dtype) + + @classmethod + def layer_norm_fallback(cls, vals, gamma, beta, epsilon): + channels = gamma.shape[0] + dtype = gamma.dtype + vals_f = vals.to(torch.float32) + gamma_f = gamma.to(torch.float32) + beta_f = beta.to(torch.float32) + return F.layer_norm(vals_f, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) + + def forward(self, vals, gamma, beta, epsilon): + return self.layer_norm_func(vals, gamma, beta, epsilon) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/linear.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..b8decb6dc5ea4333dd2ae117db352936de9b6c63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/linear.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +import deepspeed + + +class LinearOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(LinearOp, self).__init__(config) + try: + if self.config.dtype in [torch.float16, torch.int8]: + if deepspeed.HAS_TRITON and self.config.use_triton and self.config.dtype == torch.float16: + from deepspeed.ops.transformer.inference.triton.ops import linear_func as _triton_linear_func + self.linear_func = _triton_linear_func + triton_autotune = config.triton_autotune and config.layer_id == 0 + if triton_autotune: + __class__._triton_autotune(2, self.config.max_out_tokens, self.config.hidden_size) + else: + self.linear_func = self.inference_module.linear_layer_fp16 + self.linear_func = self.inference_module.linear_layer_fp16 + elif self.config.dtype == torch.bfloat16: + self.linear_func = self.inference_module.linear_layer_bf16 + else: + self.linear_func = self.inference_module.linear_layer_fp32 + except AttributeError: + self.linear_func = self.linear_fallback + + def linear_fallback(self, input, weight, bias, add_bias, do_flash_attn, num_heads, transpose, rope_theta): + raise NotImplementedError + + def forward(self, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + add_bias: bool, + do_flash_attn: bool, + num_heads: int, + external_cache: bool = None, + num_layers: int = None): + qkv_out = self.linear_func(input, weight, bias, add_bias, do_flash_attn, num_heads, + self.config.transposed_mode, self.config.rope_theta) + return qkv_out + + @staticmethod + def _triton_autotune(min_seqlen, max_seqlen, hidden_size, dtype=torch.float16): + from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, matmul + seqlen = [(min_seqlen + i) + for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] + Fp16Matmul._read_autotune_table() + for N in seqlen: + A = torch.randn((N, hidden_size), dtype=dtype, device='cuda') + B = torch.randn((hidden_size, 3 * hidden_size), dtype=dtype, device='cuda') + matmul(A, B) + Fp16Matmul._update_autotune_table() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..5f1f915ec02165b74ba5a8b7ec6dc238f9534d72 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Optional + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +from deepspeed.utils.types import NormType +from .pre_rms_norm import PreRMSNormOp + + +class MLPGemmOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(MLPGemmOp, self).__init__(config) + try: + if self.config.norm_type == NormType.LayerNorm: + if self.config.dtype in [ + torch.float16, torch.int8 + ]: # non-triton cuda kernel has a higher performance in MLP than mlp_gemm_func in triton.ops + self.mlp_gemm_func = self.inference_module.mlp_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.mlp_gemm_func = self.inference_module.mlp_gemm_bf16 + else: + self.mlp_gemm_func = self.inference_module.mlp_gemm_fp32 # type: ignore + elif self.config.norm_type == NormType.RMSNorm: + if self.config.dtype in [torch.float16, torch.int8]: + self.mlp_gemm_func = self.inference_module.rms_mlp_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.mlp_gemm_func = self.inference_module.rms_mlp_gemm_bf16 + else: + self.mlp_gemm_func = self.inference_module.rms_mlp_gemm_fp32 # type: ignore + except AttributeError: + if self.config.norm_type == NormType.LayerNorm: + self.mlp_gemm_func = self.mlp_gemm_fallback + elif self.config.norm_type == NormType.RMSNorm: + self.mlp_gemm_func = self.rms_mlp_gemm_fallback + self.pre_rms_norm = PreRMSNormOp() + + def mlp_gemm_fallback(self, input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, + transpose): + if mlp_after_attn: + residual_add = F.layer_norm(input + residual + input_bias, (input.shape[2], ), gamma, beta, eps) + tmp = torch.matmul(residual_add, weight_interm.t() if transpose else weight_interm) + tmp = F.gelu(tmp + bias) + output = torch.matmul(tmp, weight_out.t() if transpose else weight_out) + + return output, residual_add + else: + raise NotImplementedError + + def rms_mlp_gemm_fallback(self, input, residual, weight_interm, weight_out, gamma, eps, interm_scale, out_scale, + dtype, mlp_act_func_type, transpose): + inp_norm, residual = self.pre_rms_norm(input, residual, gamma, eps) + tmp = torch.matmul(inp_norm.view([-1, inp_norm.size(2)]), weight_interm.t() if transpose else weight_interm) + up_proj, gate_proj = tmp.chunk(2, dim=1) + + from deepspeed.utils.types import ActivationFuncType + if mlp_act_func_type == ActivationFuncType.GELU: + intermediate = F.gelu(gate_proj) + elif mlp_act_func_type == ActivationFuncType.ReLU: + intermediate = F.relu(gate_proj) + elif mlp_act_func_type == ActivationFuncType.GATED_GELU: + intermediate = F.gelu(gate_proj) + elif mlp_act_func_type == ActivationFuncType.GATED_SILU: + intermediate = F.silu(gate_proj) + else: + raise f"rms_mlp_gemm_fallback not implemented for activation type {mlp_act_func_type}" + + intermediate = intermediate * up_proj + + output = torch.matmul(intermediate, weight_out.t() if transpose else weight_out) + output = output.view([input.size(0), input.size(1), -1]) + + return [output, residual] + + def forward(self, + input: torch.Tensor, + residual: torch.Tensor, + weight_interm: torch.Tensor, + weight_out: torch.Tensor, + input_bias: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + gamma: Optional[torch.Tensor] = None, + beta: Optional[torch.Tensor] = None): + if self.config.norm_type == NormType.LayerNorm: + output, residual_add = self.mlp_gemm_func( + input, + residual, + input_bias, + weight_interm, + weight_out, + bias, + gamma, + beta, + self.config.epsilon, + self.config.pre_layer_norm, + self.config.mlp_after_attn, + weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), # type: ignore + weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore + self.config.dtype == torch.int8, + self.config.mlp_act_func_type, + self.config.transposed_mode) + else: + if input_bias is not None: + input += input_bias + output, residual_add = self.mlp_gemm_func( + input, + residual, + weight_interm, + weight_out, + gamma, + self.config.epsilon, + weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), # type: ignore + weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore + self.config.dtype == torch.int8, + self.config.mlp_act_func_type, + self.config.transposed_mode) + return output, residual_add diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..ef3558c8bc889442159aa73eec22315d42481fb4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class MoEResMatmulOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + super(MoEResMatmulOp, self).__init__(config) + try: + self.moe_res_matmul_func = self.inference_module.moe_res_matmul + except AttributeError: + self.moe_res_matmul_func = self.moe_res_matmul_fallback + + @classmethod + def moe_res_matmul_fallback(cls, residual, coef, output): + coef_t = coef.transpose(1, 2).contiguous() + coef1, coef2 = torch.split(coef_t, split_size_or_sections=coef_t.shape[len(coef_t.shape) - 1] // 2, dim=-1) + return residual * coef1 + output * coef2 + + def forward(self, residual, coef, output): + return self.moe_res_matmul_func(residual, coef, output) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pad_transform.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pad_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..876fefc3bcfbd34e2f64a30703132649a96224ec --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pad_transform.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class PadTransformOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + super(PadTransformOp, self).__init__(config) + try: + self.pad_transform_func = self.inference_module.pad_transform_fp16 + except AttributeError: + self.pad_transform_func = self.pad_transform_fallback + + @staticmethod + def pad_transform_fallback(query, key, value, heads, do_flash_attn): + raise NotImplementedError("pad_transform fallback is not implemented.") + + def forward(self, query, key, value, heads, do_flash_attn): + return self.pad_transform_func(query, key, value, heads, do_flash_attn) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..7969d20f0527bd6fb1fce2f666bd5e5683296cb9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +from .rms_norm import RMSNormOp + + +class PreRMSNormOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + super(PreRMSNormOp, self).__init__(config) + try: + self.pre_rms_norm_func = self.inference_module.pre_rms_norm + except AttributeError: + self.pre_rms_norm_func = self.pre_rms_norm_fallback + + @staticmethod + def pre_rms_norm_fallback(vals, residual, gamma, epsilon): + residual = vals.to(torch.float32) + residual.to(torch.float32) + vals = residual + + return RMSNormOp.rms_norm_fallback(vals, gamma, epsilon), residual.to(gamma.dtype) + + def forward(self, vals, residual, gamma, epsilon): + return self.pre_rms_norm_func(vals, residual, gamma, epsilon) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..9ff5366fae5d6a267a44852f1aace7a98d3ad3dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +from .rms_norm import RMSNormOp +import deepspeed +from deepspeed.utils.types import NormType + + +class QKVGemmOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(QKVGemmOp, self).__init__(config) + try: + if self.config.norm_type == NormType.LayerNorm: + if self.config.dtype in [torch.float16, torch.int8]: + if deepspeed.HAS_TRITON and self.config.use_triton and self.config.dtype == torch.float16: + from deepspeed.ops.transformer.inference.triton.ops import qkv_gemm_func as _triton_qkv_gemm_func + self.qkv_gemm_func = _triton_qkv_gemm_func + triton_autotune = config.triton_autotune and config.layer_id == 0 + if triton_autotune: + __class__._triton_autotune(2, self.config.max_out_tokens, self.config.hidden_size) + else: + self.qkv_gemm_func = self.inference_module.qkv_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.qkv_gemm_func = self.inference_module.qkv_gemm_bf16 + else: + self.qkv_gemm_func = self.inference_module.qkv_gemm_fp32 # type: ignore + elif self.config.norm_type == NormType.RMSNorm: + if self.config.dtype in [torch.float16, torch.int8]: + self.qkv_gemm_func = self.inference_module.rms_qkv_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.qkv_gemm_func = self.inference_module.rms_qkv_gemm_bf16 + else: + self.qkv_gemm_func = self.inference_module.rms_qkv_gemm_fp32 # type: ignore + except AttributeError: + if self.config.norm_type == NormType.LayerNorm: + self.qkv_gemm_func = self.qkv_gemm_fallback + elif self.config.norm_type == NormType.RMSNorm: + self.qkv_gemm_func = self.rms_qkv_gemm_fallback + + @staticmethod + def _triton_autotune(min_seqlen, max_seqlen, hidden_size, dtype=torch.float16): + from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, matmul + seqlen = [(min_seqlen + i) + for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] + Fp16Matmul._read_autotune_table() + for N in seqlen: + A = torch.randn((N, hidden_size), dtype=dtype, device='cuda') + B = torch.randn((hidden_size, 3 * hidden_size), dtype=dtype, device='cuda') + matmul(A, B) + Fp16Matmul._update_autotune_table() + + @staticmethod + def qkv_gemm_fallback(input, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + inp_norm = F.layer_norm(input, (input.shape[2], ), gamma, beta, eps) + tmp = torch.matmul(inp_norm, weight.t() if transpose else weight) + if add_bias: + tmp += bias + output = [tmp, inp_norm] + + return output + + @staticmethod + def rms_qkv_gemm_fallback(input, weight, q_scale, gamma, eps, q_int8, transpose): + inp_norm = RMSNormOp.rms_norm_fallback(input, gamma, eps) + tmp = torch.matmul(inp_norm, weight.t() if transpose else weight) + output = [tmp, inp_norm] + + return output + + def forward(self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, gamma: torch.Tensor, + beta: torch.Tensor): + + add_bias = bias is not None + bias = bias if add_bias else torch.empty(1) # type: ignore + q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) # type: ignore + q_int8 = self.config.dtype == torch.int8 + + if self.config.norm_type == NormType.LayerNorm: + output, norm = self.qkv_gemm_func(input, weight, q_scale, bias, gamma, beta, self.config.epsilon, add_bias, + q_int8, self.config.transposed_mode) + else: + output, norm = self.qkv_gemm_func(input, weight, q_scale, gamma, self.config.epsilon, q_int8, + self.config.transposed_mode) + if add_bias: + output += bias + + return output, norm diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/residual_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/residual_add.py new file mode 100644 index 0000000000000000000000000000000000000000..93b229c5d1ac0165f82ccc322769f11dd61d890e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/residual_add.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from typing import Optional + +from .vector_add import VectorAddOp +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class ResidualAddOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(ResidualAddOp, self).__init__(config) + try: + if self.config.dtype in [torch.float16, torch.int8]: + self.residual_add_func = self.inference_module.residual_add_bias_fp16 + elif self.config.dtype == torch.bfloat16: + self.residual_add_func = self.inference_module.residual_add_bias_bf16 + else: + self.residual_add_func = self.inference_module.residual_add_bias_fp32 + except AttributeError: + self.residual_add_func = self.residual_add_fallback + self.vector_add = VectorAddOp() + + @staticmethod + def res_add_bias(hidden_state, residual, attn_output, attn_bias, final_bias, add_attn_bias, mp_size): + hidden_state += attn_output + (residual + final_bias) / mp_size + if add_attn_bias: + hidden_state += attn_bias / mp_size + + return hidden_state + + @staticmethod + def residual_add_fallback(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + if mlp_after_attn: + if pre_layer_norm: + tmp = (residual.float() + attention_output.float() + attention_bias.float() + + final_bias.float()) / mp_size + hidden_state.float() + else: + tmp = residual.float() + hidden_state.float() + final_bias.float() + else: + tmp = ResidualAddOp.res_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + add_bias, mp_size) + residual.copy_(tmp.to(hidden_state.dtype)) + + return residual + + def forward(self, + hidden_state: torch.Tensor, + residual: torch.Tensor, + add_bias: bool, + attention_output: Optional[torch.Tensor] = None, + residual_add: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + final_bias: Optional[torch.Tensor] = None): + + if final_bias is None and attention_bias is None: + residual = self.vector_add(residual + attention_output, hidden_state, 1.0 / self.config.mp_size) + else: + if not self.config.pre_layer_norm and residual_add is not None: + # only use residual add if its set and we are not pre layer norm + residual = residual_add + + self.residual_add_func(hidden_state, residual, attention_output, attention_bias, final_bias, + self.config.mp_size, self.config.mlp_after_attn, add_bias, + self.config.pre_layer_norm) + + return residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/rms_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/rms_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..128883ce5d438b03718496abe1e26134d9e7f440 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/rms_norm.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class RMSNormOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + super(RMSNormOp, self).__init__(config) + try: + self.rms_norm_func = self.inference_module.rms_norm + except AttributeError: + self.rms_norm_func = self.rms_norm_fallback + + @staticmethod + def rms_norm_fallback(vals, gamma, epsilon): + variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True) + vals = vals * torch.rsqrt(variance + epsilon) + + if gamma.dtype in [torch.float16, torch.bfloat16]: + vals = vals.to(gamma.dtype) + + return gamma * vals + + def forward(self, vals, gamma, epsilon): + return self.rms_norm_func(vals, gamma, epsilon) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax.py new file mode 100644 index 0000000000000000000000000000000000000000..2e08541596fa8dc89f8ac893254dc773ed5ee284 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext + + +class SoftmaxOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(SoftmaxOp, self).__init__(config) + self.num_attention_heads_per_partition = config.heads // config.mp_size + try: + if self.config.dtype in [torch.float16, torch.int8]: + self.softmax_func = self.inference_module.softmax_fp16 + elif self.config.dtype == torch.bfloat16: + self.softmax_func = self.inference_module.softmax_bf16 + else: + self.softmax_func = self.inference_module.softmax_fp32 + except AttributeError: + self.softmax_func = self.softmax_fallback + + @staticmethod + def softmax_fallback(attn_scores, attn_mask, alibi, triangular, recompute, local_attention, window_size, async_op, + layer_scale, head_offset, mp_size): + scores_len = len(attn_scores.size()) + heads = 1 + if scores_len > 1: + heads = attn_scores.size()[1] + num_attention_heads_per_partition = heads // mp_size + + if alibi is not None: + if len(alibi.shape) == 1: + alibi = None + else: + alibi = alibi[head_offset:head_offset + num_attention_heads_per_partition] + if attn_mask is not None and len(attn_mask.shape) == 1: + attn_mask = None + input_dtype = attn_scores.dtype + attn_scores *= layer_scale + + if alibi is not None: + attn_scores += alibi + if attn_mask is not None: + # expand atten_mask from two dim into 4 dim, insert two dims in the middle + if len(attn_mask.shape) == 2: + attn_mask = attn_mask[:, None, None, :] + attn_scores += attn_mask + if triangular: + if attn_scores.shape[2] == 1: # query using kv cache + token_idx = InferenceContext.Instance().current_tokens() + tri = torch.arange(attn_scores.shape[2], device=attn_scores.device).ge(token_idx) + else: + tri = ~torch.tril(torch.ones(attn_scores.size(), device=attn_scores.device)).to(bool) + attn_scores = torch.masked_fill(attn_scores, tri, float('-inf')) + output = F.softmax(attn_scores, dim=-1, dtype=torch.float32).to(input_dtype) + + return output + + def forward(self, attn_scores: torch.Tensor, attn_mask: torch.Tensor, alibi: torch.Tensor, triangular: bool, + recompute: bool, local_attention: bool, window_size: int, async_op: bool, layer_scale: float, + head_offset: int): + output = self.softmax_func(attn_scores, attn_mask, alibi, triangular, recompute, local_attention, window_size, + async_op, layer_scale, head_offset, self.config.mp_size) + + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax_context.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax_context.py new file mode 100644 index 0000000000000000000000000000000000000000..d745df678e93bf3e359288610156c7cb8392e2b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax_context.py @@ -0,0 +1,144 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +from .softmax import SoftmaxOp +from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext + + +class SoftmaxContextOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(SoftmaxContextOp, self).__init__(config) + try: + if self.config.dtype in [torch.float16, torch.int8]: + self.softmax_context_func = self.inference_module.softmax_context_fp16 + elif self.config.dtype == torch.bfloat16: + self.softmax_context_func = self.inference_module.softmax_context_bf16 + else: + self.softmax_context_func = self.inference_module.softmax_context_fp32 + except AttributeError: + self.softmax_context_func = self.softmax_context_fallback + + @staticmethod + def transform4d_0213(x, seq_length): + assert x.dim() == 3, F"Dim {x.dim()} is not supported" + batch_size, num_heads, seq_length_head_dim = x.shape + head_dim = seq_length_head_dim // seq_length + x = x.view(batch_size, num_heads, seq_length, head_dim) + x = x.permute(0, 2, 1, 3) + + return x + + @staticmethod + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep <= 1 or num_key_value_heads == 1: + return hidden_states + + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + @staticmethod + def bias_add_transform_0213(input, bias, num_heads, trans_count, perform_bias=False): + assert trans_count == 1 or trans_count == 3, F"Trans count {trans_count} is not supported" + assert input.dim() == 3, F"Dim {input.dim()} is not supported" + input_biased = torch.add(input, bias) if perform_bias else input + batch_size, seq_length, value_size = input_biased.shape + hid_dim = value_size // trans_count + head_dim = hid_dim // num_heads + + if trans_count == 1: + query_layer = input.view(batch_size, seq_length, num_heads, head_dim) + query_layer = query_layer.permute(0, 2, 1, 3) + key_layer = torch.zeros_like(query_layer) + value_layer = torch.zeros_like(query_layer) + return query_layer, key_layer, value_layer + + qkv_layers = input.view(batch_size, seq_length, 3, num_heads, head_dim) + query_layer, key_layer, value_layer = qkv_layers[..., 0, :, :], qkv_layers[..., 1, :, :], qkv_layers[..., + 2, :, :] + query_layer = query_layer.transpose(1, 2) + key_layer = key_layer.transpose(1, 2) + value_layer = value_layer.transpose(1, 2) + + return query_layer, key_layer, value_layer + + def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, + num_kv, norm_factor, triangular_masking, local_attention, window_size, no_masking, + layer_id, num_layers, alibi, rope_theta, is_prompt, token_idx, position_ids): + bat_0213_query, bat_0213_key, bat_0213_value = self.bias_add_transform_0213( + query_key_value, None, heads, 3, False) + + if rotary_dim > 0 and rotate_half: + from transformers.models.llama.modeling_llama import apply_rotary_pos_emb + + rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device) + cos, sin = rotary(bat_0213_value, InferenceContext.Instance().get_max_tokens_num()) + bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin, position_ids) + + bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt, + bat_0213_key, bat_0213_value) + + bat_0213_key = self.repeat_kv(bat_0213_key, num_kv) + bat_0213_value = self.repeat_kv(bat_0213_value, num_kv) + + bsz = query_key_value.shape[0] + head_dim = query_key_value.shape[2] // (heads * 3) + + bmm_output = torch.bmm(bat_0213_query.reshape(bsz * heads, bat_0213_query.shape[2], head_dim), + bat_0213_key.reshape(bsz * heads, bat_0213_key.shape[2], head_dim).transpose(1, 2)) + + layer_scale = 1.0 + if alibi is not None and len(alibi.shape) > 1: + layer_scale = max(1, layer_id).to(float) + + alpha = norm_factor * norm_factor / layer_scale + bmm_output *= alpha + bmm_output_reshape = bmm_output.reshape(bsz, heads, bmm_output.shape[1], bmm_output.shape[2]) + + recompute = is_prompt + if attn_mask is not None and len(attn_mask.shape) > 1 and attn_mask.shape[-1] < bmm_output_reshape.shape[3]: + attn_mask = torch.nn.functional.pad(attn_mask, (0, bmm_output_reshape.shape[3] - attn_mask.shape[-1]), + value=torch.finfo(attn_mask.dtype).min) + softmax_output = SoftmaxOp.softmax_fallback(bmm_output_reshape, attn_mask, alibi, triangular_masking, + recompute, local_attention, window_size, None, layer_scale, 0, 1) + + output = torch.bmm(softmax_output.reshape(bsz * heads, softmax_output.shape[2], softmax_output.shape[3]), + bat_0213_value.reshape(bsz * heads, bat_0213_value.shape[2], head_dim)) + + output = output.reshape(bsz, heads, output.shape[1], head_dim) + output = output.reshape(bsz, heads, output.shape[2] * head_dim) + input_seq_len = query_key_value.shape[1] + t4d_0123_output = self.transform4d_0213(output, input_seq_len) + t4d_0123_output = t4d_0123_output.reshape(bsz, t4d_0123_output.shape[1], heads * head_dim) + + if layer_id == num_layers - 1: + InferenceContext.Instance().advance_tokens() + + return t4d_0123_output, bat_0213_key, bat_0213_value + + def forward(self, query_key_value: torch.Tensor, attn_mask: torch.Tensor, heads: int, num_kv: int, + norm_factor: float, no_masking: bool, layer_id: int, num_layers: int, alibi: torch.Tensor, + is_prompt: bool, token_idx: torch.Tensor, position_ids: torch.Tensor): + + if alibi is not None: + batch_heads = query_key_value.shape[0] * heads + offset = dist.get_rank() * batch_heads if dist.is_initialized() else 0 + alibi = alibi[offset:batch_heads + offset, :, :] + else: + alibi = torch.empty(1) + + output = self.softmax_context_func(query_key_value, attn_mask, self.config.rotary_dim, self.config.rotate_half, + self.config.rotate_every_two, heads, num_kv, norm_factor, + self.config.triangular_masking, self.config.local_attention, + self.config.window_size, no_masking, layer_id, num_layers, alibi, + self.config.rope_theta, is_prompt, token_idx, position_ids) + + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_add.py new file mode 100644 index 0000000000000000000000000000000000000000..015340a1084b5a64557737928462c5a1a79ae2ba --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_add.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + + +class VectorAddOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + super(VectorAddOp, self).__init__(config) + try: + self.vector_add_func = self.inference_module._vector_add + except AttributeError: + self.vector_add_func = self.vector_add_fallback + + @classmethod + def vector_add_fallback(cls, a, b, gamma): + """Based on csrc/transformer/inference/csrc/pt_binding.cpp code of _vector_add""" + dtype = a.dtype + return (gamma * a.float() + b.float()).to(dtype) + + def forward(self, a, b, gamma): + return self.vector_add_func(a, b, gamma) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..cabab8d8c4ab146bde0f7d0d69a9056a299694c1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp +import deepspeed + + +class VectorMatMulOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig): + super(VectorMatMulOp, self).__init__(config) + try: + if self.config.dtype == torch.float16: + if deepspeed.HAS_TRITON and config.use_triton: + from deepspeed.ops.transformer.inference.triton.ops import vector_matmul_func as _triton_vector_matmul_func + self.vector_matmul_func = _triton_vector_matmul_func + triton_autotune = config.triton_autotune and config.layer_id == 0 + if triton_autotune: + __class__._triton_autotune(2, self.config.max_out_tokens, self.config.hidden_size) + else: + self.vector_matmul_func = self.inference_module.vector_matmul_fp16 + elif self.config.dtype == torch.int8: + self.vector_matmul_func = self.inference_module.vector_matmul_int8 + elif self.config.dtype == torch.bfloat16: + self.vector_matmul_func = self.inference_module.vector_matmul_bf16 + else: + self.vector_matmul_func = self.inference_module.vector_matmul_fp32 + except AttributeError: + self.vector_matmul_func = self.vector_matmul_fallback + + def vector_matmul_fallback(self, input, weight, async_op, q_scale, q_int8, transpose): + return torch.matmul(input, weight.t() if transpose else weight) + + def forward(self, input: torch.Tensor, weight: torch.Tensor, async_op: bool = False): + q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) + q_int8 = self.config.dtype == torch.int8 + output = self.vector_matmul_func(input, weight, async_op, q_scale, q_int8, self.config.transposed_mode) + return output + + @staticmethod + def _triton_autotune(min_seqlen, max_seqlen, hidden_size, dtype=torch.float16): + from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, matmul + seqlen = [(min_seqlen + i) + for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] + Fp16Matmul._read_autotune_table() + for N in seqlen: + A = torch.randn((N, hidden_size), dtype=dtype, device='cuda') + B = torch.randn((hidden_size, hidden_size), dtype=dtype, device='cuda') + matmul(A, B) + Fp16Matmul._update_autotune_table() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/workspace.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/workspace.py new file mode 100644 index 0000000000000000000000000000000000000000..19de7d9576afe44bfbf2e9b4c1f24a842a05ebbf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/workspace.py @@ -0,0 +1,222 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..config import DeepSpeedInferenceConfig +from .base import BaseOp + +minus_inf = -10000.0 +key_idx = 0 +value_idx = 1 + + +class InferenceContext: + + __instance = None + + def __init__(self): + self.kv_cache = None + self.kv_cache_elem_dtype = None + self.num_tokens = 1 + self.kv_cache_num_layers = None + self.kv_cache_size = None + self.max_out_tokens = None + self.rotary = None + self.allocate_called = False + self.static_shapes = True + + @classmethod + def Instance(cls): + if InferenceContext.__instance is None: + InferenceContext.__instance = InferenceContext() + return InferenceContext.__instance + + def gen_workspace(self, num_layers, num_heads, batch_size, prompt_len, hidden_dim, mp_size, external_cache, + elem_dtype, rank, max_out_tokens, min_out_tokens): + self.allocate_called = True + self.kv_cache = None + if not external_cache: + self.kv_cache_num_layers = num_layers + self.max_out_tokens = max_out_tokens + head_size = hidden_dim // num_heads + self.kv_cache_size = torch.Size([batch_size, (num_heads // mp_size), max_out_tokens, head_size]) + self.kv_cache_elem_dtype = elem_dtype + self.num_tokens = 0 + self.static_shapes = True + return True + + def retake_workspace(self): + return True + + def _retake_workspace(self): + assert self.allocate_called, "retake workspace called before allocate workspace" + + import deepspeed.accelerator as accelerator + if self.kv_cache is None: + self.kv_cache = [] + for layer in range(self.kv_cache_num_layers): + self.kv_cache.append((torch.zeros(self.kv_cache_size, + dtype=self.kv_cache_elem_dtype, + device=accelerator.get_accelerator().device_name()), + torch.zeros(self.kv_cache_size, + dtype=self.kv_cache_elem_dtype, + device=accelerator.get_accelerator().device_name()))) + + return True + + def update_cache(self, layer_id, token_idx, is_prompt, bat_0213_key, bat_0213_value): + has_workspace = self._retake_workspace() + assert has_workspace, "Could not allocate workspace" + + # Update current token + if is_prompt: + self.static_shapes = True + if token_idx is None: + self.static_shapes = False + InferenceContext.Instance().reset_tokens(bat_0213_key.shape[2]) + else: + InferenceContext.Instance().reset_tokens(token_idx) + + if token_idx is None: + token_idx = InferenceContext.Instance().current_tokens() + + bsz = bat_0213_key.shape[0] + + # Update cache content + if is_prompt: + cache_max_seq = self.kv_cache_size[2] + cache_max_head_dim = self.kv_cache_size[3] + seq = bat_0213_key.shape[2] + + mask = torch.arange(cache_max_seq, device=bat_0213_key.device) + mask = mask.ge(token_idx) + mask = mask.unsqueeze(-1) + mask = mask.expand([cache_max_seq, cache_max_head_dim]) + + self.kv_cache[layer_id][key_idx][:bsz, :, :seq, :].copy_(bat_0213_key) + self.kv_cache[layer_id][key_idx][:bsz, :].masked_fill_(mask, 0) + self.kv_cache[layer_id][value_idx][:bsz, :, :seq, :].copy_(bat_0213_value) + self.kv_cache[layer_id][value_idx][:bsz, :].masked_fill_(mask, 0) + else: + if self.static_shapes: + assert type(token_idx) == torch.Tensor, "token_idx is expected to be torch.Tensor" + self.kv_cache[layer_id][key_idx][:bsz].index_copy_(2, token_idx - 1, bat_0213_key) + self.kv_cache[layer_id][value_idx][:bsz].index_copy_(2, token_idx - 1, bat_0213_value) + else: + assert type(token_idx) == int, "token_idx is expected to be int" + self.kv_cache[layer_id][key_idx][:bsz, :, token_idx - 1:token_idx, :] = bat_0213_key + self.kv_cache[layer_id][value_idx][:bsz, :, token_idx - 1:token_idx, :] = bat_0213_value + + bat_0213_key = self.kv_cache[layer_id][key_idx][:bsz] + bat_0213_value = self.kv_cache[layer_id][value_idx][:bsz] + + if not self.static_shapes: + bat_0213_key = bat_0213_key[:, :, :token_idx, :] + bat_0213_value = bat_0213_value[:, :, :token_idx, :] + + return bat_0213_key, bat_0213_value + + def release_workspace(self): + self.kv_cache = None + self.rotary = None + + def reset_tokens(self, initial_tokens=1): + self.num_tokens = initial_tokens + + def current_tokens(self): + return self.num_tokens + + def advance_tokens(self): + self.num_tokens = self.num_tokens + 1 + + def get_kv_cache(self): + return self.kv_cache + + def get_rotary(self, rotary_dim, rope_theta, device=None): + if self.rotary is None: + from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding + + self.rotary = LlamaRotaryEmbedding(rotary_dim, base=rope_theta, device=device) + + return self.rotary + + def get_max_tokens_num(self): + return self.max_out_tokens + + +class WorkspaceOp(BaseOp): + + def __init__(self, config: DeepSpeedInferenceConfig = None): + if config is None: + config = DeepSpeedInferenceConfig() + self.inference_context = InferenceContext.Instance() + self._is_allocated = False + try: + super(WorkspaceOp, self).__init__(config) + if config.dtype == torch.float32: + self.allocate_workspace_func = self.inference_module.allocate_workspace_fp32 + elif config.dtype == torch.bfloat16: + self.allocate_workspace_func = self.inference_module.allocate_workspace_bf16 + else: + self.allocate_workspace_func = self.inference_module.allocate_workspace_fp16 + self.release_workspace_func = self.inference_module.release_workspace + self.retake_workspace_func = self.inference_module.retake_workspace + self.reset_cache_func = self.inference_module.reset_cache + except (ValueError, AttributeError) as e: + print(f"Using fallback functions in workspace because of {e}") + if config.dtype == torch.float32: + self.allocate_workspace_func = self.allocate_workspace_fp32_fallback + elif config.dtype == torch.bfloat16: + self.allocate_workspace_func = self.allocate_workspace_bf16_fallback + else: + self.allocate_workspace_func = self.allocate_workspace_fp16_fallback + self.release_workspace_func = self.release_workspace_fallback + self.retake_workspace_func = self.retake_workspace_fallback + self.reset_cache_func = self.reset_cache_fallback + + def allocate_workspace(self, *args, **kwargs): + self._is_allocated = True + return self.allocate_workspace_func(*args, **kwargs) + + def release_workspace(self): + self._is_allocated = False + return self.release_workspace_func() + + def reset_cache(self): + return self.reset_cache_func() if self.reset_cache_func else None + + def retake_workspace(self): + return self.retake_workspace_func() if self.retake_workspace_func else None + + def allocate_workspace_fp32_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, + external_cache, rank, max_out_tokens, min_out_tokens): + return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, + mp_size, external_cache, torch.float, rank, max_out_tokens, + min_out_tokens) + + def allocate_workspace_bf16_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, + external_cache, rank, max_out_tokens, min_out_tokens): + return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, + mp_size, external_cache, torch.bfloat16, rank, max_out_tokens, + min_out_tokens) + + def allocate_workspace_fp16_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, + external_cache, rank, max_out_tokens, min_out_tokens): + return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, + mp_size, external_cache, torch.half, rank, max_out_tokens, + min_out_tokens) + + def reset_cache_fallback(self): + return self.inference_context.reset_tokens() + + def release_workspace_fallback(self): + return self.inference_context.release_workspace() + + def retake_workspace_fallback(self): + return self.inference_context.retake_workspace() + + def is_allocated(self): + return self._is_allocated diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7d1968df62a99849992c8b1e93698d9e51cec30 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .residual_add import residual_add_bias +from .layer_norm import layer_norm, layer_norm_residual +from .gelu import gelu +from .softmax import softmax +from .ops import * +from .matmul_ext import fp16_matmul, matmul_4d, score_4d_matmul, context_4d_matmul diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e6a0176bfe1be058db64f53eae02a9382752b2d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e52b580672056b6eee52fcfc2f4a2faa0c6e3328 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57812e209d38a9a88276c453bdc6fc0b37cfc82d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f72d3f22e65db520fa6b9e99a51e0f13b347d739 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..882737cb3d1e647047fb379983a0f1869ea277ad Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f826ca6c54b90878f90cbb4dab4b2459092d04ba Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0153896a84a9fdea6257af8fa0dc8dc26d2df89d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfacdec3c8db9f24f8f16441c79ae40491c32a7e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..10c6794cd9703fc2caf9a29ac11ec0b935bbd207 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07988c76d8954d9ef3c2dc716f20d05554e5cd39 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..6845d91b06be91ffe3e0eedb7f59fd79609c9f92 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/attention.py @@ -0,0 +1,397 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +import torch +import torch.nn as nn +import triton +import triton.language as tl +from deepspeed.accelerator import get_accelerator +from deepspeed import comm as dist +from deepspeed.ops.transformer.inference.op_binding import LinearOp, VectorMatMulOp, SoftmaxContextOp, QKVGemmOp +from deepspeed.ops.transformer.inference.triton import ( + softmax, + score_4d_matmul, + context_4d_matmul, +) + +minus_inf = -10000.0 + + +class TritonSelfAttention(nn.Module): + num_layers = 0 + + def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, qkv_merging=False): + super(TritonSelfAttention, self).__init__() + self.config = config + data_type = self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype + assert data_type_fp == torch.half, "triton supports fp16 data_type_fp" + + self.config.layer_id = TritonSelfAttention.num_layers + TritonSelfAttention.num_layers = TritonSelfAttention.num_layers + 1 + device = get_accelerator().current_device_name() #if config.bigscience_bloom else 'cpu' + + assert config.mp_size == 1, "mp_size has to be 1 with triton attention yet" + if self.config.set_empty_params: + self.attn_qw = None + self.attn_qb = None + self.attn_kw = None + self.attn_kb = None + self.attn_vw = None + self.attn_vb = None + self.attn_qkvw = None + self.attn_qkvb = None + self.attn_ow = None + self.attn_ob = None + else: + qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 + self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, + qkv_size_per_partition, + dtype=data_type, + device=device), + requires_grad=False) + self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), + requires_grad=False) + # self-ouput weights + out_size_per_partition = self.config.hidden_size // self.config.mp_size + self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + + self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + + self.num_attention_heads_per_partition = self.config.heads // self.config.mp_size + self.hidden_size_per_partition = self.config.hidden_size // self.config.mp_size + self.hidden_size_per_attention_head = self.config.hidden_size // self.config.heads + + self.mp_group = mp_group + self.use_flash = False + # triton flash attention is enabled when the compute capability >= 8.0 + if get_accelerator().is_triton_supported(): + self.use_flash = True + + # used for quantization + self.q_scales = q_scales + self.q_groups = q_groups + self.merge_count = int(math.log2(merge_count)) + + self.norm_factor = math.sqrt(self.config.hidden_size // self.config.heads) + if not config.use_mup: + self.norm_factor = math.sqrt(self.norm_factor) + + if self.config.scale_attn_by_inverse_layer_idx is True: + self.norm_factor *= math.sqrt(self.config.layer_id + 1) + # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 + + triton_autotune = self.config.triton_autotune and self.config.layer_id == 0 + self.qkv_func = QKVGemmOp(config) + self.score_context_func = SoftmaxContextOp(config) + self.linear_func = LinearOp(config) + self.vector_matmul_func = VectorMatMulOp(config) + + self.hidden_size = config.hidden_size + self.head_size = config.hidden_size // config.heads + self.scale = (1 / self.norm_factor / self.norm_factor if self.config.scale_attention else 1.0 + ) # making it back to 1/sqrt(head_size) + self.triangular_masking = self.config.triangular_masking + + # triton autotune table update for score/context matmul + if triton_autotune: + print(f"running triton autotune for regular attention kernel") + __class__._triton_autotune(2, self.config.max_out_tokens, self.head_size, self.config.hidden_size, + self.triangular_masking, self.scale) + + @staticmethod + def _triton_autotune(min_seqlen, + max_seqlen, + head_size, + hidden_size, + triangular_masking, + scale, + dtype=torch.float16): + from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, score_4d_matmul, context_4d_matmul + seqlen = [(min_seqlen + i) + for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] + Fp16Matmul._read_autotune_table() + for N in seqlen: + qkv = torch.randn((1, N, 3 * hidden_size), dtype=dtype, device='cuda') + output = score_4d_matmul(qkv, head_size, triangular_masking, scale) + context_4d_matmul(output, qkv, head_size) + Fp16Matmul._update_autotune_table() + + def ds_compute_attention(self, qkv_out, input_mask, layer_past, alibi, is_prompt, token_idx, position_ids): + if isinstance(qkv_out, list): + qkv_out = qkv_out[0] + + no_masking = input_mask is None + + if no_masking: + input_mask = torch.empty(1) + + attn_key_value = self.score_context_func( + query_key_value=qkv_out, + attn_mask=((1 - input_mask).to(qkv_out.dtype) * + minus_inf) if input_mask.dtype == torch.int64 else input_mask, + heads=self.num_attention_heads_per_partition, + norm_factor=(1 / self.norm_factor if self.config.scale_attention else 1.0), + no_masking=no_masking, + layer_id=self.config.layer_id, + num_layers=TritonSelfAttention.num_layers, + alibi=alibi, + is_prompt=is_prompt, + token_idx=token_idx, + position_ids=position_ids) + + context_layer, key_layer, value_layer = attn_key_value + return context_layer, key_layer, value_layer + + def forward( + self, + input, + input_mask, + head_mask=None, + layer_past=None, + get_present=False, # not used + encoder_hidden_states=None, # not used + encoder_attention_mask=None, # not used + triangularutput_attentions=False, # not used + norm_w=None, + norm_b=None, + alibi=None, + use_triton_attention=True, + **kwargs): + + if not self.config.pre_layer_norm: + qkv_out = self.linear_func(input=input, + weight=self.attn_qkvw, + bias=self.attn_qkvb, + add_bias=self.attn_qkvb is not None, + do_flash_attn=False, + num_heads=self.num_attention_heads_per_partition, + num_layers=TritonSelfAttention.num_layers) + qkv = qkv_out + else: + qkv_out = self.qkv_func(input=input, + weight=self.attn_qkvw, + bias=(self.attn_qkvb if self.attn_qkvb is not None else norm_b), + gamma=norm_w, + beta=norm_b) + qkv = qkv_out[0] + + if use_triton_attention and (alibi is None): + context_layer = _triton_attention(qkv=qkv, + input_mask=input_mask, + scale=self.scale, + layer_past=layer_past, + alibi=alibi, + head_size=self.head_size, + use_triton_flash=self.use_flash, + use_cuda_flash=False, + triangular=self.triangular_masking) + key_layer, value_layer = qkv[:, :, self.hidden_size:2 * self.hidden_size], qkv[:, :, 2 * self.hidden_size:] + else: + is_prompt = kwargs.get("first_token", qkv_out[0].shape[1] > 1) + token_idx = kwargs.get("token_idx", None) + position_ids = kwargs.get("position_ids", None) + context_layer, key_layer, value_layer = self.ds_compute_attention(qkv_out=qkv_out, + input_mask=input_mask, + layer_past=layer_past, + alibi=alibi, + is_prompt=is_prompt, + toke_idx=token_idx, + position_ids=position_ids) + output = self.vector_matmul_func(input=context_layer, weight=self.attn_ow) + + inp_norm = qkv_out[-1] + + if self.config.mlp_after_attn and self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: + dist.all_reduce(output, group=self.mp_group) + + return (output, key_layer, value_layer, context_layer, inp_norm) + + +global inference_module + + +def _triton_attention(qkv, + input_mask, + layer_past, + alibi, + scale, + head_size, + triangular=False, + use_cuda_flash=False, + use_triton_flash=False, + use_ds_attention=False): + if isinstance(qkv, list): + qkv = qkv[0] + + assert alibi is None, "layer_past not supported in alibi yet" + + if use_triton_flash: + output = _triton_packed_flash(qkv, + head_size, + input_mask, + scale, + causal=triangular, + add_mask=(not triangular and input_mask is not None)) + else: + output = score_4d_matmul(qkv, head_size, triangular, scale) + if triangular: + output = softmax(output) + else: + output = softmax(output, input_mask) + output = context_4d_matmul(output, qkv, head_size) + + return output + + +''' +flash attention 2 +modified the triton kernel in +https://github.com/openai/triton/blob/08c16589573621fcb8cd5a9c3b8a0537077f876d/python/tutorials/06-fused-attention.py +''' + + +@triton.jit +def _flash_packed_kernel( + QKV, + mask, + ADD_MASK: tl.constexpr, + IS_CAUSAL: tl.constexpr, + sm_scale, + Out, + stride_qz, + stride_qn, + stride_qm, + stride_mz, + stride_oz, + stride_on, + Z, + H, + N_CTX, + P_SEQ, + hidden_size, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + start_m = tl.program_id(0) + off_hz = tl.program_id(1) + batch = off_hz // H + head = off_hz % H + + q_offset = batch * stride_qz + head * BLOCK_DMODEL + k_offset = q_offset + hidden_size + v_offset = k_offset + hidden_size + + # initialize offsets + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + + q_ptrs = QKV + q_offset + offs_m[:, None] * stride_qn + offs_d[None, :] + k_ptrs = QKV + hidden_size + q_offset + offs_n[:, None] * stride_qn + offs_d[None, :] + v_ptrs = QKV + 2 * hidden_size + q_offset + offs_n[:, None] * stride_qn + offs_d[None, :] + + # mask + off_mask = batch * stride_mz + offs_n[None, :] + mask_ptrs = mask + off_mask + + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + # scale sm_scale by log_2(e) and use + # 2^x instead of exp in the loop because CSE and LICM + # don't work as expected with `exp` in the loop + qk_scale = sm_scale * 1.44269504 + # load q: it will stay in SRAM throughout + q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0) + q = (q * qk_scale).to(tl.float16) + # loop over k, v and update accumulator + lo = 0 + hi = P_SEQ + (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX + P_SEQ + for start_n in range(lo, hi, BLOCK_N): + # -- load k, v -- + k = tl.load(k_ptrs + start_n * stride_qn, mask=(start_n + offs_n)[:, None] < N_CTX, other=0.0) + v = tl.load(v_ptrs + start_n * stride_qn, mask=(start_n + offs_n)[:, None] < N_CTX, other=0.0) + # -- compute qk --- + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float16) + + if ADD_MASK: + mask_val = tl.load(mask_ptrs) + mask_ptrs += BLOCK_N + qk = qk + mask_val.to(tl.float32) + + if IS_CAUSAL: + qk = tl.where(P_SEQ + offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) + + qk += tl.dot(q, tl.trans(k), out_dtype=tl.float16) + qk += tl.where((start_n + offs_n)[None, :] < N_CTX, 0, minus_inf) + # -- compute scaling constant --- + m_i_new = tl.maximum(m_i, tl.max(qk, 1)) + alpha = tl.math.exp2(m_i - m_i_new) + p = tl.math.exp2(qk - m_i_new[:, None]) + # -- scale and update acc -- + acc_scale = l_i * 0 + alpha # workaround some compiler bug + acc *= acc_scale[:, None] + acc += tl.dot(p.to(tl.float16), v.to(tl.float16)) + # -- update m_i and l_i -- + l_i = l_i * alpha + tl.sum(p, 1) + m_i = m_i_new + + # write back l and m + acc = acc / l_i[:, None] + o_offset = batch * stride_oz + head * BLOCK_DMODEL + out_ptrs = Out + o_offset + (offs_m[:, None] * stride_on + offs_d[None, :]) + tl.store(out_ptrs, acc.to(tl.float16), mask=offs_m[:, None] < N_CTX) + + +def _triton_packed_flash(qkv, head_size, mask, sm_scale, causal=False, add_mask=True): + heads = qkv.shape[-1] // 3 // head_size + hidden_size = qkv.shape[-1] // 3 + + BLOCK_M = 128 + BLOCK_N = 64 if head_size <= 64 else 32 + + o = torch.empty((qkv.shape[0], qkv.shape[1], hidden_size), device=qkv.device, dtype=torch.half) + if mask is None: + mask = torch.empty(0) + add_mask = False + + grid = (triton.cdiv(qkv.shape[1], BLOCK_M), qkv.shape[0] * heads, 1) + num_stages = 4 if head_size <= 64 else 3 + num_warps = 4 + P_SEQ = 0 + + _flash_packed_kernel[grid](qkv, + mask, + add_mask, + causal, + sm_scale, + o, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + mask.stride(1) if add_mask else 0, + o.stride(0), + o.stride(1), + qkv.shape[0], + heads, + qkv.shape[1], + P_SEQ, + hidden_size, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=head_size, + num_warps=num_warps, + num_stages=num_stages) + + return o diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/gelu.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/gelu.py new file mode 100644 index 0000000000000000000000000000000000000000..738d7d96a1c9d57c3de3558452922737b9b4f7b4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/gelu.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import triton +import triton.language as tl +from deepspeed.accelerator import get_accelerator + + +@triton.jit +def gelu_functor(x): + # Using approximation introduces greater parity errors. + # return tl.sigmoid(1.702 * x) * x + return x * 0.5 * (1.0 + tl.math.erf(x / 1.41421356237)) + + +@triton.jit +def gelu_kernel(x_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + output = gelu_functor(x) + tl.store(output_ptr + offsets, output, mask=mask) + + +def gelu(activations: torch.Tensor) -> torch.Tensor: + assert activations.is_contiguous() + assert get_accelerator().on_accelerator(activations) + + output = torch.empty_like(activations) + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + gelu_kernel[grid](activations, output, n_elements, BLOCK_SIZE=1024) + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/layer_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/layer_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..d3f313d2ac3d8205702dc2ceb82856154e9ddb2c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/layer_norm.py @@ -0,0 +1,249 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import triton +import triton.language as tl +''' +layer-normalization +modified the triton kernel in +https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/tutorials/05-layer-norm.py +''' + + +@triton.jit +def layer_norm_kernel( + Out, + A, + Weight, + Bias, + stride, + N, + eps, + BLOCK_SIZE: tl.constexpr, +): + # position of elements processed by this program + row = tl.program_id(0) + Out += row * stride + A += row * stride + # compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) + a = tl.where(cols < N, a - mean, 0.0) + _var += a * a + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # multiply by weight and add bias + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + weight = tl.load(Weight + cols, mask=mask) + bias = tl.load(Bias + cols, mask=mask) + a = tl.load(A + cols, mask=mask, other=0.0).to(tl.float32) + a_hat = (a - mean) * rstd + out = a_hat * weight + bias + # # write-back + tl.store(Out + cols, out, mask=mask) + + +@triton.jit +def layer_norm_residual_kernel( + Out, + A, + Residual, + ln_input, + Weight, + Bias, + stride, + N, + eps, + BLOCK_SIZE: tl.constexpr, +): + # position of elements processed by this program + row = tl.program_id(0) + Out += row * stride + A += row * stride + Residual += row * stride + ln_input += row * stride + # compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) + res = tl.load(Residual + cols, mask=cols < N, other=0.0).to(tl.float32) + a = a + res + tl.store(ln_input + cols, a, mask=cols < N) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(ln_input + cols, mask=cols < N, other=0.0).to(tl.float32) + a = tl.where(cols < N, a - mean, 0.0) + _var += a * a + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # multiply by weight and add bias + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + weight = tl.load(Weight + cols, mask=mask) + bias = tl.load(Bias + cols, mask=mask) + a = tl.load(ln_input + cols, mask=mask, other=0.0).to(tl.float32) + a_hat = (a - mean) * rstd + out = a_hat * weight + bias + # write-back + tl.store(Out + cols, out, mask=mask) + + +@triton.jit +def layer_norm_residual_bias_kernel( + Out, + A, + Residual, + InputBias, + ln_input, + Weight, + Bias, + stride, + N, + eps, + BLOCK_SIZE: tl.constexpr, +): + # position of elements processed by this program + row = tl.program_id(0) + Out += row * stride + A += row * stride + Residual += row * stride + ln_input += row * stride + # compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) + res = tl.load(Residual + cols, mask=cols < N, other=0.0).to(tl.float32) + b = tl.load(InputBias + cols, mask=cols < N, other=0.0).to(tl.float32) + a = a + b + res + tl.store(ln_input + cols, a, mask=cols < N) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(ln_input + cols, mask=cols < N, other=0.0).to(tl.float32) + a = tl.where(cols < N, a - mean, 0.0) + _var += a * a + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # multiply by weight and add bias + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + weight = tl.load(Weight + cols, mask=mask) + bias = tl.load(Bias + cols, mask=mask) + a = tl.load(ln_input + cols, mask=mask, other=0.0).to(tl.float32) + a_hat = (a - mean) * rstd + out = a_hat * weight + bias + # write-back + tl.store(Out + cols, out, mask=mask) + + +def layer_norm(a, weight, bias, eps): + assert a.is_contiguous() + assert weight.is_contiguous() + assert bias.is_contiguous() + + # allocate output + out = torch.empty_like(a) + # reshape input data into 2D tensor + a_arg = a.view(-1, a.shape[-1]) + M, N = a_arg.shape + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // a.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + BLOCK_SIZE = max(BLOCK_SIZE, 128) + BLOCK_SIZE = min(BLOCK_SIZE, 4096) + BLOCK_SIZE = BLOCK_SIZE if N <= 4096 else 8192 + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + layer_norm_kernel[(M, )]( + out, + a_arg, + weight, + bias, + a_arg.stride(0), + N, + eps, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + return out + + +def layer_norm_residual(a, input_bias, residual, weight, bias, eps): + assert a.is_contiguous() + assert weight.is_contiguous() + assert bias.is_contiguous() + assert residual.is_contiguous() + + # allocate output and scratch-pad for residual addition + out = torch.empty_like(a) + ln_input = torch.empty_like(a) + # reshape input data into 2D tensor + a_arg = a.view(-1, a.shape[-1]) + residual = residual.view(-1, residual.shape[-1]) + M, N = a_arg.shape + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // a.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + BLOCK_SIZE = max(BLOCK_SIZE, 128) + BLOCK_SIZE = min(BLOCK_SIZE, 4096) + BLOCK_SIZE = BLOCK_SIZE if N <= 4096 else 8192 + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + if input_bias is None: + layer_norm_residual_kernel[(M, )]( + out, + a_arg, + residual, + ln_input, + weight, + bias, + a_arg.stride(0), + N, + eps, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + else: + layer_norm_residual_bias_kernel[(M, )]( + out, + a_arg, + residual, + input_bias, + ln_input, + weight, + bias, + a_arg.stride(0), + N, + eps, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + return out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py new file mode 100644 index 0000000000000000000000000000000000000000..9be4b0098c371d689581e3317bba701472a33ddc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py @@ -0,0 +1,480 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import triton +import os +from filelock import FileLock +import deepspeed.ops.transformer.inference.triton.triton_matmul_kernel as triton_matmul_kernel +import pickle +from io import open +import deepspeed +from pathlib import Path +import atexit +import subprocess + + +# ----------------------------------------------------------------------------- +# util class/functions for triton +def is_nfs_path(path): + if os.name == 'nt': + return False + + # Normalize the path to get the absolute path + path = os.path.abspath(path) + + # Use the 'df' command to find the file system type for the given path + try: + output = subprocess.check_output(['df', '-T', path], encoding='utf-8') + except subprocess.CalledProcessError: + return False # Command failed + + # Process the output of 'df -T' to check for 'nfs' in the filesystem type column + lines = output.strip().split('\n') + if len(lines) > 1: # The first line is headers + fs_type = lines[1].split()[1].lower() # File system type is the second column + return 'nfs' in fs_type + return False + + +class TritonCacheDir: + _warning_printed = False + + @staticmethod + def warn_if_nfs(cache_dir): + if is_nfs_path(cache_dir) and not TritonCacheDir._warning_printed: + print( + f"Warning: The cache directory for DeepSpeed Triton autotune, {cache_dir}, appears to be on an NFS system. While this is generally acceptable, if you experience slowdowns or hanging when DeepSpeed exits, it is recommended to set the TRITON_CACHE_DIR environment variable to a non-NFS path." + ) + TritonCacheDir._warning_printed = True + return + + @staticmethod + def default_cache_dir(): + tmp_path = os.path.join(Path.home(), ".triton", "autotune") + return tmp_path + + +def bias_add_activation(C, bias=None, activation=""): + if bias is not None: + C += bias + # activation + if activation == "relu": + relu = torch.nn.Relu() + C = relu(C) + elif activation == "leaky_relu": + leaky_relu = torch.nn.LeakyReLU(0.01) + C = leaky_relu(C) + elif activation == "gelu": + sigmoid = torch.nn.Sigmoid() + C = sigmoid(1.702 * C) * C + elif activation == "sigmoid": + sigmoid = torch.nn.Sigmoid() + C = sigmoid(C) + return C + + +class AutotuneCacheManager: + """ + Cache manager for autotune + """ + + def __init__(self, key): + self.key = key + self.file_path = None + self.lock_path = None + # if caching is enabled, get the lock and bin path + self.cache_dir = os.environ.get('TRITON_CACHE_DIR', TritonCacheDir.default_cache_dir()) + TritonCacheDir.warn_if_nfs(self.cache_dir) + if self.cache_dir: + os.makedirs(self.cache_dir, exist_ok=True) + self.file_path = os.path.join(self.cache_dir, self.key + ".pickle") + self.lock_path = self.file_path + ".lock" + + def has_file(self): + return self.file_path and os.path.exists(self.file_path) + + def put(self, table): + if self.file_path: + assert self.lock_path is not None + with FileLock(self.lock_path): + with open(self.file_path + ".tmp", 'wb') as handle: + pickle.dump(table, handle) + os.replace(self.file_path + ".tmp", self.file_path) + + def load(self): + if os.path.exists(self.file_path): + with open(self.file_path, 'rb') as handle: + loaded_dict = pickle.load(handle) + return loaded_dict + else: + return None + + +# ----------------------------------------------------------------------------- +# triton matmul class + + +class MatmulExt(torch.autograd.Function): + """ + a wrapper class that can call different triton matmul kernels depending on the input parameters + """ + + @staticmethod + def forward(A, B, bias=None, activation="", use_triton=True, update_autotune_table=False): + """ + A: input, activation matrix A + B: input, weight matrix B + """ + matmul = None + quantize_activation = False + Batch = 0 + + if len(A.shape) == 3: # if A is 3d-tensor where batch index is given as 0-axis + assert A.is_contiguous(), "matrix A must be contiguous" + Batch, M, K = A.shape + A = A.view(-1, K) + + # fp16 activation and fp16 weight matmul into fp16 output + matmul = fp16_matmul + C = matmul.forward(A, B, use_triton=use_triton, bias=bias, activation=activation) + + if matmul and update_autotune_table: + matmul._update_autotune_table() + + if Batch > 0: + C = C.view(Batch, M, -1) + + return C + + +class TritonMatmul(torch.autograd.Function): + """ + triton matmul kernel superclass + """ + + def __init__(self): + pass + + @staticmethod + def _ref_forward(A, B, ref_dtype=torch.float32): + C = torch.matmul(A.type(ref_dtype), B.type(ref_dtype)) + return C + + @staticmethod + def _read_autotune_table(cache_key, triton_kernel): + cache_manager = AutotuneCacheManager(cache_key) + table = cache_manager.load() + if table: + triton_kernel.cache = table + + @staticmethod + def _write_autotune_table(cache_key, triton_kernel): + cache_manager = AutotuneCacheManager(cache_key) + cache_manager.put(triton_kernel.cache) + + @staticmethod + def _update_autotune_table(cache_key, triton_kernel): + cache_manager = AutotuneCacheManager(cache_key) + autotune_table = cache_manager.load() + if autotune_table is None: + autotune_table = dict() + autotune_table.update(triton_kernel.cache) # always overwrite with the new autotune results + cache_manager = AutotuneCacheManager(cache_key) + cache_manager.put(autotune_table) + + @staticmethod + def forward( + A, + B, + ref_dtype=torch.float32, # fp32 only + bias=None, + activation=""): + C = torch.matmul(A.type(ref_dtype), B.type(ref_dtype)) + C = bias_add_activation(C, bias, activation) + return C + + +class Fp16Matmul(TritonMatmul): + """ + fp16 matrix multiplication kernel + dtypes: fp16 x fp16 = fp16 + """ + + _2d_kernel = triton_matmul_kernel._fp_matmul + _4d_kernel = triton_matmul_kernel.matmul_4d_kernel + _cache_stride = 32 + + def __init__(self, read_cache=True): + super().__init__() + if read_cache: + __class__._read_autotune_table() + + def skip_autotune(self): + __class__._2d_kernel.configs = [__class__._2d_kernel.configs[0]] + __class__._4d_kernel.configs = [__class__._4d_kernel.configs[0]] + + @staticmethod + def forward(A, B, use_triton=True, bias=None, activation=""): + if use_triton: + device = A.device + # handle non-contiguous inputs if necessary + if A.stride(0) > 1 and A.stride(1) > 1: + A = A.contiguous() + if B.stride(0) > 1 and B.stride(1) > 1: + B = B.contiguous() + # checks constraints + assert A.shape[1] == B.shape[0], "incompatible dimensions" + M, K = A.shape + _, N = B.shape + # allocates output + C = torch.empty((M, N), device=device, dtype=A.dtype) + # accumulator types + ACC_TYPE = triton.language.float32 if A.dtype in [torch.float16, torch.bfloat16, torch.float32 + ] else triton.language.int32 + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']) * triton.cdiv(N, META['BLOCK_N']), META['SPLIT_K']) + __class__._2d_kernel[grid](A, + B, + C, + M, + N, + K, + bias, + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(1), + C.stride(0), + C.stride(1), + M // __class__._cache_stride, + N // __class__._cache_stride, + K // __class__._cache_stride, + GROUP_M=8, + ACC_TYPE=ACC_TYPE, + BIAS_ADD=(0 if bias is None else 1), + ACTIVATION=activation) + else: + C = torch.matmul(A, B) + return C + + @staticmethod + def _matmul_4d(a, b): + assert a.shape[-1] == b.shape[-2], "incompatible dimensions" + assert a.is_contiguous(), "matrix A must be contiguous" + assert b.is_contiguous(), "matrix B must be contiguous" + + B, H, M, K = a.shape + B, H, K, N = b.shape + + assert K > 1, "inner-product dimension K should be larger than 1" + + c = torch.empty((B, H, M, N), device=a.device, dtype=a.dtype) + + grid = lambda META: ( + triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), + H, + B, + ) + + __class__._4d_kernel[grid]( + a, + b, + c, + M, + N, + K, + M // __class__._cache_stride, + N // __class__._cache_stride, + K // __class__._cache_stride, + a.stride(0), + a.stride(1), + a.stride(2), + a.stride(3), + b.stride(0), + b.stride(1), + b.stride(2), + b.stride(3), + c.stride(0), + c.stride(1), + c.stride(2), + c.stride(3), + scale=-1.0, + MASK=False, + ) + return c + + @staticmethod + def _score_4d_matmul(input, head_size, input_mask, scale=-1.0): + assert input.is_contiguous(), "matrix input must be contiguous" + + batches = input.shape[0] + d_model = input.shape[-1] // 3 + num_of_heads = d_model // head_size + + q = input[:, :, :d_model] + k = input[:, :, d_model:d_model * 2] + + q = q.view(batches, -1, num_of_heads, head_size) + k = k.view(batches, -1, num_of_heads, head_size) + + # checks constraints + assert q.shape == k.shape, "incompatible dimensions" + B, M, H, K = q.shape + B, N, H, K = k.shape + + assert K > 1, "inner-product dimension K should be larger than 1" + + # allocates output + output = torch.empty((B, H, M, N), device=q.device, dtype=q.dtype) + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + H, + B, + ) + __class__._4d_kernel[grid]( + q, + k, + output, + M, + N, + K, + M // __class__._cache_stride, + N // __class__._cache_stride, + K // __class__._cache_stride, + q.stride(0), + q.stride(2), + q.stride(1), + q.stride(3), + k.stride(0), + k.stride(2), + k.stride(3), + k.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + output.stride(3), + scale=scale, + MASK=False, + ) + return output + + @staticmethod + def _context_4d_matmul(prob, input, head_size): + assert prob.is_contiguous(), "matrix prob must be contiguous" + assert input.is_contiguous(), "matrix input must be contiguous" + + batches = input.shape[0] + d_model = input.shape[-1] // 3 + num_of_heads = d_model // head_size + + v = input[:, :, d_model * 2:] + + v = v.view(batches, -1, num_of_heads, head_size) + + # checks constraints + assert (prob.shape[0] == v.shape[0] and prob.shape[1] == v.shape[2] and prob.shape[2] == v.shape[1] + and prob.shape[3] == v.shape[1]), "incompatible dimensions" + B, H, M, K = prob.shape + B, K, H, N = v.shape + + assert K > 1, "inner-product dimension K should be larger than 1" + + # allocates output + output = torch.empty((B, M, H, N), device=v.device, dtype=v.dtype) + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + H, + B, + ) + + __class__._4d_kernel[grid]( + prob, + v, + output, + M, + N, + K, + M // __class__._cache_stride, + N // __class__._cache_stride, + K // __class__._cache_stride, + prob.stride(0), + prob.stride(1), + prob.stride(2), + prob.stride(3), + v.stride(0), + v.stride(2), + v.stride(1), + v.stride(3), + # Here we also transpose the output when writing to memory. + output.stride(0), + output.stride(2), + output.stride(1), + output.stride(3), + scale=-1, + MASK=False, + ) + return output.view(batches, -1, d_model) + + @staticmethod + def _ref_forward(A, B, ref_dtype=torch.float32, bias=None, activation=""): + C = torch.matmul(A.type(ref_dtype), B.type(ref_dtype)) + C = bias_add_activation(C, bias, activation) + return C + + @staticmethod + def _check_parity(A, + B, + output_dtype, + SA=None, + SB=None, + qblock_size=None, + ref_dtype=torch.float32, + tol=0.01, + use_triton=True, + bias=None, + activation=""): + torch_output = __class__._ref_forward(A, B, ref_dtype=ref_dtype, bias=bias, activation=activation) + triton_output = __class__.forward(A, B, use_triton=use_triton, bias=bias, activation=activation) + assert torch.allclose(triton_output.cpu().type(torch_output.dtype), torch_output.cpu(), rtol=tol) + print(f"{__class__.__name__}: PASSed the parity check") + return triton_output, torch_output + + @staticmethod + def _read_autotune_table(): + TritonMatmul._read_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel) + TritonMatmul._read_autotune_table(__class__.__name__ + "_4d_kernel", __class__._4d_kernel) + + @staticmethod + def _write_autotune_table(): + TritonMatmul._write_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel) + TritonMatmul._write_autotune_table(__class__.__name__ + "_4d_kernel", __class__._4d_kernel) + + @staticmethod + def _update_autotune_table(): + TritonMatmul._update_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel) + TritonMatmul._update_autotune_table(__class__.__name__ + "_4d_kernel", __class__._4d_kernel) + + +# ----------------------------------------------------------------------------- +# mapping +if deepspeed.HAS_TRITON: + fp16_matmul = Fp16Matmul() + matmul = MatmulExt.forward + matmul_4d = fp16_matmul._matmul_4d + score_4d_matmul = fp16_matmul._score_4d_matmul + context_4d_matmul = fp16_matmul._context_4d_matmul +else: + fp16_matmul = None + matmul = None + matmul_4d = None + score_4d_matmul = None + context_4d_matmul = None + + +@atexit.register +def matmul_ext_update_autotune_table(): + if deepspeed.HAS_TRITON: + fp16_matmul._update_autotune_table() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/mlp.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..1708080b27efb0671d361b7cdeaff4b262cf0ce8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/mlp.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import math +import torch.nn as nn +from deepspeed.accelerator import get_accelerator +from deepspeed import comm as dist +from ..op_binding import MLPGemmOp, VectorMatMulOp, GELUGemmOp, ResidualAddOp + + +class TritonMLP(nn.Module): + + def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False): + super(TritonMLP, self).__init__() + + self.config = config + data_type = self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype + device = get_accelerator().current_device_name() + self.attn_nw = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + self.attn_nb = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + intm_size_per_partition = self.config.intermediate_size // self.config.mp_size + self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, + intm_size_per_partition, + dtype=data_type, + device=device), + requires_grad=False) + self.inter_b = nn.Parameter(torch.empty(intm_size_per_partition, dtype=data_type_fp, device=device), + requires_grad=False) + self.output_w = nn.Parameter(torch.empty(intm_size_per_partition, + self.config.hidden_size, + dtype=data_type, + device=device), + requires_grad=False) + self.output_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), + requires_grad=False) + + # used for quantization + self.q_scales = q_scales + self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups + self.merge_count = int(math.log2(merge_count)) + self.mp_group = mp_group + + self.mlp_gemm_func = MLPGemmOp(config) + self.vector_matmul_func = VectorMatMulOp(config) + self.fused_gemm_gelu = GELUGemmOp(config) + self.residual_add_func = ResidualAddOp(config) + + def forward(self, input, residual, residual_norm, bias): + residual_add = None + if self.attn_nw is None: + output = self.fused_gemm_gelu(input=residual_norm, + weight=self.inter_w, + bias=self.inter_b, + weight_out=self.output_w) + else: + output, residual_add = self.mlp_gemm_func(input=input, + residual=residual, + input_bias=bias, + weight_interm=self.inter_w, + weight_out=self.output_w, + bias=self.inter_b, + gamma=self.attn_nw, + beta=self.attn_nb) + residual = self.residual_add_func(hidden_state=output, + residual=residual, + attention_output=input, + attention_bias=bias if bias is not None else self.output_b, + final_bias=self.output_b, + add_bias=bias is not None, + residual_add=residual_add) + + if self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: + dist.all_reduce(residual, group=self.mp_group) + + return residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..dbed453137806a66d13fbef62292297e1c30c3f4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/ops.py @@ -0,0 +1,123 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import deepspeed.ops.transformer.inference.triton.matmul_ext as matmul_ext +from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp +from deepspeed.ops.transformer.inference.triton.layer_norm import layer_norm, layer_norm_residual +from deepspeed.utils.types import ActivationFuncType + + +def vector_matmul_func(input, weight, async_op, q_scale, q_int8, transposed_mode): + assert not transposed_mode and not async_op and not q_int8 + return matmul_ext.matmul(input, weight, bias=None, activation="", use_triton=True) + + +def fused_gemm_gelu(input, + weight, + weight_scale, + bias, + weight_out, + weight_out_scale, + epsilon, + pre_layer_norm, + q_int8, + async_op, + transposed_mode, + use_triton_ln=True): + assert not transposed_mode + + # activation + activation = "gelu" + + # intermediate fc in FF + intm_out = matmul_ext.matmul(input, weight, bias=bias, activation=activation, use_triton=True) + + # output fc in FF + ff_out = matmul_ext.matmul( + intm_out, + weight_out, + bias=None, + activation="", # bias added layer with residual_add + bias + layerNorm layer + use_triton=True) + return ff_out + + +def linear_func(input, weight, bias, add_bias, do_flash_attn, num_heads, transposed_mode=False): + assert not transposed_mode and not do_flash_attn + qkv_out = matmul_ext.matmul(input, weight, bias=(bias if add_bias else None), activation="", use_triton=True) + + return qkv_out + + +def mlp_gemm_func(input, + residual, + input_bias, + weight_interm, + weight_out, + bias, + gamma, + beta, + epsilon, + pre_layer_norm, + mlp_after_attn, + weight_interm_scale, + weight_out_scale, + q_int8, + mlp_act_func_type, + transposed_mode, + use_triton_ln=True): + assert not transposed_mode + + # residual add and layerNorm after attention + if use_triton_ln: + mlp_input = layer_norm_residual(input, input_bias, residual, gamma, beta, epsilon) + else: + mlp_input = LayerNormOp.layer_norm_residual(input, input_bias, residual, gamma, beta, epsilon) + + # activation + if ActivationFuncType(mlp_act_func_type) == ActivationFuncType.GELU: + activation = "gelu" + elif ActivationFuncType(mlp_act_func_type) == ActivationFuncType.ReLU: + activation = "relu" + else: + activation = "" + + # intermediate fc in FF + intm_out = matmul_ext.matmul(mlp_input, weight_interm, bias=bias, activation=activation, use_triton=True) + # output fc in FF + ff_out = matmul_ext.matmul( + intm_out, + weight_out, + bias=None, + activation="", # bias added layer with residual_add + bias + layerNorm layer + use_triton=True) + + return ff_out, mlp_input + + +def qkv_gemm_func( + input, + weight, + q_scale, + bias, + gamma, + beta, + epsilon, + add_bias, + q_int8, + transposed_mode=False, + use_triton_ln=True, +): + + assert not transposed_mode + # residual add and layerNorm after attention + if use_triton_ln: + qkv_input = layer_norm(input, gamma, beta, epsilon) + else: + qkv_input = LayerNormOp()(input, gamma, beta, epsilon) + + qkv_out = matmul_ext.matmul(qkv_input, weight, bias=(bias if add_bias else None), activation="", use_triton=True) + + return qkv_out, qkv_input diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/residual_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/residual_add.py new file mode 100644 index 0000000000000000000000000000000000000000..063e7a7e4a2d9740ad641b3116ea2cec37b472f4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/residual_add.py @@ -0,0 +1,88 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import triton +import triton.language as tl +from deepspeed.accelerator import get_accelerator + + +@triton.jit +def residual_add_bias_kernel( + hidden_state_ptr, + residual_ptr, + attn_output_ptr, + hidden_state_size, + attn_bias_ptr, + final_bias_ptr, + bias_size, + output_ptr, + mp_size: tl.constexpr, + mlp_after_attn: tl.constexpr, + pre_attn_norm: tl.constexpr, + add_attn_bias: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + block_start = pid * BLOCK_SIZE + + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_state_size + + bias_offsets = offsets % bias_size + bias_mask = bias_offsets < bias_size + + tl_hidden_state = tl.load(hidden_state_ptr + offsets, mask=mask) + tl_residual = tl.load(residual_ptr + offsets, mask=mask) + tl_attn_output = tl.load(attn_output_ptr + offsets, mask=mask) + tl_attn_bias = tl.load(attn_bias_ptr + bias_offsets, mask=bias_mask) + tl_final_bias = tl.load(final_bias_ptr + bias_offsets, mask=bias_mask) + + if mlp_after_attn: + if pre_attn_norm: + output = tl_hidden_state + (tl_residual + tl_final_bias + tl_attn_output + tl_attn_bias) / mp_size + else: + output = tl_hidden_state + tl_residual + tl_final_bias + else: + output = tl_hidden_state + tl_attn_output + (tl_residual + tl_final_bias) / mp_size + if add_attn_bias: + output += tl_attn_bias / mp_size + + tl.store(output_ptr + offsets, output, mask=mask) + + +def residual_add_bias(hidden_state: torch.Tensor, residual: torch.Tensor, attn_output: torch.Tensor, + attn_bias: torch.Tensor, final_bias: torch.Tensor, mp_size: int, mlp_after_attn: bool, + add_attn_bias: bool, pre_attn_norm: bool): + # check that all tensors are on the same device + assert get_accelerator().on_accelerator(hidden_state) \ + and get_accelerator().on_accelerator(residual) \ + and get_accelerator().on_accelerator(attn_output) \ + and get_accelerator().on_accelerator(attn_bias) \ + and get_accelerator().on_accelerator(final_bias) + + # check that all tensors have the same dtype + assert hidden_state.dtype == residual.dtype == attn_output.dtype \ + == attn_bias.dtype == final_bias.dtype + + # check that all tensors have the right shape + assert hidden_state.shape == residual.shape == attn_output.shape + assert attn_bias.shape == final_bias.shape + assert attn_bias.shape[0] == hidden_state.shape[2] + + output = torch.empty_like(hidden_state) + + hidden_state_size = output.numel() + bias_size = attn_bias.numel() + + grid = lambda meta: (triton.cdiv(hidden_state_size, meta['BLOCK_SIZE']), ) + + residual_add_bias_kernel[grid](hidden_state, residual, attn_output, hidden_state_size,\ + attn_bias, final_bias, bias_size, output, mp_size, mlp_after_attn, pre_attn_norm, \ + add_attn_bias, \ + BLOCK_SIZE=1024) + + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/softmax.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/softmax.py new file mode 100644 index 0000000000000000000000000000000000000000..1ee10d63e6cf8bfa6723856b53b7ca9ec30d3fdd --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/softmax.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import triton +import triton.language as tl +''' +softmax +modified the triton kernel in +https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/tutorials/02-fused-softmax.py +''' + + +@triton.jit +def softmax_kernel(output_ptr, input_ptr, stride, n_cols, BLOCK_SIZE: tl.constexpr): + row_idx = tl.program_id(0) + row_start_ptr = input_ptr + row_idx * stride + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf')).to(tl.float32) + row_minus_max = row - tl.max(row, axis=0) + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + output_row_start_ptr = output_ptr + row_idx * stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols) + + +@triton.jit +def masked_softmax_kernel(output_ptr, input_ptr, stride, mask_ptr, mask_stride, n_cols, BLOCK_SIZE: tl.constexpr): + row_idx = tl.program_id(0) + row_start_ptr = input_ptr + row_idx * stride + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + mask_ptrs = mask_ptr + col_offsets + row_idx * mask_stride # mask_stride is 0 for 1d mask + row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf')).to(tl.float32) + mask = tl.load(mask_ptrs, mask=col_offsets < n_cols, other=0).to(tl.float32) + row_minus_max = row - tl.max(row, axis=0) + row_minus_max = row_minus_max + mask + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + output_row_start_ptr = output_ptr + row_idx * stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols) + + +def softmax(input: torch.Tensor, mask: torch.Tensor = None, dim=-1) -> torch.Tensor: + assert input.is_contiguous() + assert (dim == -1) or (dim == len(input.shape) - 1), "Only dim=-1 is supported" + + use_mask = False if mask is None else True + input_arg = input.view(-1, input.shape[-1]) + n_rows, n_cols = input_arg.shape + BLOCK_SIZE = max(triton.next_power_of_2(n_cols), 2) + num_warps = 4 + if BLOCK_SIZE >= 2048: + num_warps = 8 + if BLOCK_SIZE >= 4096: + num_warps = 16 + # Allocate output + output = torch.empty_like(input) + if use_mask: + assert mask.is_contiguous() + mask = mask.view(-1, mask.shape[-1]) + mask_stride = mask.shape[-1] if mask.shape[-2] > 1 else 0 + masked_softmax_kernel[(n_rows, )]( + output, + input, + input_arg.stride(0), + mask, + mask_stride, + n_cols, + num_warps=num_warps, + BLOCK_SIZE=BLOCK_SIZE, + ) + else: + softmax_kernel[(n_rows, )]( + output, + input, + input_arg.stride(0), + n_cols, + num_warps=num_warps, + BLOCK_SIZE=BLOCK_SIZE, + ) + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..e2128e046df049ddbd846131b3dc6001083e991b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py @@ -0,0 +1,398 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import triton +import triton.language as tl +from .gelu import gelu_functor +import torch + +AUTOTUNE_TOP_K = 10 +SKIP_AUTOTUNE = False + + +def _triton_ops_matmul_early_config_prune(configs, named_args): + device = torch.cuda.current_device() #ignore-cuda + capability = torch.cuda.get_device_capability() #ignore-cuda + # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages + dtsize = named_args['A'].element_size() + dtype = named_args['A'].dtype + + # 1. make sure we have enough smem + pruned_configs = [] + for config in configs: + kw = config.kwargs + BLOCK_M, BLOCK_N, BLOCK_K, num_stages = \ + kw['BLOCK_M'], kw['BLOCK_N'], kw['BLOCK_K'], config.num_stages + + max_shared_memory = triton.runtime.driver.utils.get_device_properties(device)["max_shared_mem"] + required_shared_memory = (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize + if required_shared_memory <= max_shared_memory: + pruned_configs.append(config) + + return pruned_configs + + +def _fp16_matmul_prune_config(configs, named_args, skip_autotune=SKIP_AUTOTUNE): + if skip_autotune: + configs = [configs[0]] + else: + configs = _triton_ops_matmul_early_config_prune(configs, named_args) + return configs + + +""" +fp16 matmul implementation is adapted from triton matmul: +https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/triton/ops/matmul.py +""" + + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': 128, + 'BLOCK_N': 256, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=3, num_warps=8), + triton.Config({ + 'BLOCK_M': 256, + 'BLOCK_N': 128, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=3, num_warps=8), + triton.Config({ + 'BLOCK_M': 256, + 'BLOCK_N': 64, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=4, num_warps=4), + triton.Config({ + 'BLOCK_M': 64, + 'BLOCK_N': 256, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=4, num_warps=4), + triton.Config({ + 'BLOCK_M': 128, + 'BLOCK_N': 128, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=4, num_warps=4), + triton.Config({ + 'BLOCK_M': 128, + 'BLOCK_N': 64, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=4, num_warps=4), + triton.Config({ + 'BLOCK_M': 64, + 'BLOCK_N': 128, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=4, num_warps=4), + triton.Config({ + 'BLOCK_M': 128, + 'BLOCK_N': 32, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=4, num_warps=4), + triton.Config({ + 'BLOCK_M': 64, + 'BLOCK_N': 32, + 'BLOCK_K': 32, + 'SPLIT_K': 1 + }, num_stages=5, num_warps=2), + ], + key=['CACHE_M', 'CACHE_N', 'CACHE_K'], + prune_configs_by={ + 'early_config_prune': _fp16_matmul_prune_config, + 'perf_model': None, + 'top_k': AUTOTUNE_TOP_K + }, +) +@triton.heuristics({ + 'EVEN_K': lambda args: args['K'] % (args['BLOCK_K'] * args['SPLIT_K']) == 0, +}) +@triton.jit +def _fp_matmul( + A, + B, + C, + M, + N, + K, + bias, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + CACHE_M, + CACHE_N, + CACHE_K, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + SPLIT_K: tl.constexpr, + EVEN_K: tl.constexpr, + ACC_TYPE: tl.constexpr, + BIAS_ADD: tl.constexpr, + ACTIVATION: tl.constexpr, +): + # matrix multiplication + pid = tl.program_id(0) + pid_z = tl.program_id(1) + grid_m = (M + BLOCK_M - 1) // BLOCK_M + grid_n = (N + BLOCK_N - 1) // BLOCK_N + # re-order program ID for better L2 performance + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // (group_size) + # do matrix multiplication + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) + # pointers + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + for k in range(K, 0, -BLOCK_K * SPLIT_K): + if EVEN_K: + a = tl.load(A) + b = tl.load(B) + else: + a = tl.load(A, mask=rk[None, :] < k, other=0.) + b = tl.load(B, mask=rk[:, None] < k, other=0.) + acc += tl.dot(a, b) + A += BLOCK_K * SPLIT_K * stride_ak + B += BLOCK_K * SPLIT_K * stride_bk + # bias addition + if BIAS_ADD: + bias_offset = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + bias_ptr = bias + bias_offset + b = tl.load(bias_ptr, mask=bias_offset < N) + acc = acc + b[None, :] + # activation + if ACTIVATION == "relu": + acc = tl.where(acc >= 0, acc, 0) + elif ACTIVATION == "leaky_relu": + acc = tl.where(acc >= 0, acc, 0.01 * acc) + elif ACTIVATION == "gelu": + #acc = tl.sigmoid(1.702 * acc) * acc + acc = gelu_functor(acc) + elif ACTIVATION == "sigmoid": + acc = tl.sigmoid(acc) # sigmoid + acc = acc.to(C.dtype.element_ty) + # rematerialize rm and rn to save registers + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn) + mask = (rm < M)[:, None] & (rn < N)[None, :] + # handles write-back with reduction-splitting + if SPLIT_K == 1: + tl.store(C, acc, mask=mask) + else: + tl.atomic_add(C, acc, mask=mask) + + +def matmul_4d_prune_config(configs, named_args, skip_autotune=SKIP_AUTOTUNE): + if skip_autotune: + configs = [configs[0]] + else: + device = torch.cuda.current_device() #ignore-cuda + capability = torch.cuda.get_device_capability() #ignore-cuda + # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages + dtsize = named_args['a_ptr'].element_size() + dtype = named_args['a_ptr'].dtype + + # make sure we have enough smem + pruned_configs = [] + for config in configs: + kw = config.kwargs + BLOCK_M, BLOCK_N, BLOCK_K, num_stages = \ + kw['BLOCK_SIZE_M'], kw['BLOCK_SIZE_N'], kw['BLOCK_SIZE_K'], config.num_stages + + max_shared_memory = triton.runtime.driver.utils.get_device_properties(device)["max_shared_mem"] + required_shared_memory = (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize + if required_shared_memory <= max_shared_memory: + pruned_configs.append(config) + configs = pruned_configs + return configs + + +@triton.autotune( + configs=[ + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8 + }, + num_stages=1, # this is mainly for unit test, to minimize the share memory usage + num_warps=8), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=5, + num_warps=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=5, + num_warps=2, + ), + ], + key=['CACHE_M', 'CACHE_N', 'CACHE_K'], + prune_configs_by={ + 'early_config_prune': matmul_4d_prune_config, + 'perf_model': None, + 'top_k': AUTOTUNE_TOP_K + }, +) +@triton.jit +def matmul_4d_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + # Matrix dimensions + M, + N, + K, + CACHE_M, + CACHE_N, + CACHE_K, + stride_ab, + stride_ah, + stride_am, + stride_ak, + stride_bb, + stride_bh, + stride_bk, + stride_bn, + stride_cb, + stride_ch, + stride_cm, + stride_cn, + scale, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MASK: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + pid = tl.program_id(axis=0) + head = tl.program_id(axis=1) + batch = tl.program_id(axis=2) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + if MASK: + if (pid_m + 1) * BLOCK_SIZE_M - 1 < pid_n * BLOCK_SIZE_N: + c = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=c_ptr.dtype.element_ty) - float("inf") + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = (c_ptr + batch * stride_cb + head * stride_ch + stride_cm * offs_cm[:, None] + + stride_cn * offs_cn[None, :]) + tl.store(c_ptrs, c) + return + + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = (a_ptr + batch * stride_ab + head * stride_ah + + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)) + b_ptrs = (b_ptr + batch * stride_bb + head * stride_bh + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, K, BLOCK_SIZE_K): + a_mask = (offs_am[:, None] < M) & (offs_k[None, :] + k < K) + b_mask = (offs_k[:, None] + k < K) & (offs_bn[None, :] < N) + a = tl.load(a_ptrs, mask=a_mask, other=0.) + b = tl.load(b_ptrs, mask=b_mask, other=0.) + accumulator += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + c = accumulator.to(c_ptr.dtype.element_ty) + if scale > 0: + c = c * scale.to(c_ptr.dtype.element_ty) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if MASK: + c += tl.where(offs_cm[:, None] >= offs_cn[None, :], 0, float("-inf")) + c_ptrs = (c_ptr + batch * stride_cb + head * stride_ch + stride_cm * offs_cm[:, None] + + stride_cn * offs_cn[None, :]) + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..f98f45ef638eb875454c9f9a40e9ad1bbd347320 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton_ops.py @@ -0,0 +1,166 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Inspired by original Triton implementation: +https://github.com/openai/triton/blob/release/2.1.x/python/tutorials/06-fused-attention.py +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fwd_kernel( + Q, + K, + V, + sm_scale, + Out, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vk, + stride_vn, + stride_oz, + stride_oh, + stride_om, + stride_on, + Z, + H, + N_CTX, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + start_m = tl.program_id(0) + off_hz = tl.program_id(1) + qvk_offset = off_hz * stride_qh + Q_block_ptr = tl.make_block_ptr(base=Q + qvk_offset, + shape=(N_CTX, BLOCK_DMODEL), + strides=(stride_qm, stride_qk), + offsets=(start_m * BLOCK_M, 0), + block_shape=(BLOCK_M, BLOCK_DMODEL), + order=(1, 0)) + K_block_ptr = tl.make_block_ptr(base=K + qvk_offset, + shape=(BLOCK_DMODEL, N_CTX), + strides=(stride_kk, stride_kn), + offsets=(0, 0), + block_shape=(BLOCK_DMODEL, BLOCK_N), + order=(0, 1)) + V_block_ptr = tl.make_block_ptr(base=V + qvk_offset, + shape=(N_CTX, BLOCK_DMODEL), + strides=(stride_vk, stride_vn), + offsets=(0, 0), + block_shape=(BLOCK_N, BLOCK_DMODEL), + order=(1, 0)) + # initialize offsets + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + # scale sm_scale by log_2(e) and use + # 2^x instead of exp in the loop because CSE and LICM + # don't work as expected with `exp` in the loop + qk_scale = sm_scale * 1.44269504 + # load q: it will stay in SRAM throughout + q = tl.load(Q_block_ptr) + q = (q * qk_scale).to(tl.float16) + # loop over k, v and update accumulator + lo = 0 + #hi = (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX + hi = N_CTX + #hi = (start_m + 1) * BLOCK_M + for start_n in range(lo, hi, BLOCK_N): + # -- load k, v -- + k = tl.load(K_block_ptr) + v = tl.load(V_block_ptr) + # -- compute qk --- + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + #if IS_CAUSAL: + #qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) + qk += tl.dot(q, k) + # -- compute scaling constant --- + m_i_new = tl.maximum(m_i, tl.max(qk, 1)) + alpha = tl.math.exp2(m_i - m_i_new) + p = tl.math.exp2(qk - m_i_new[:, None]) + # -- scale and update acc -- + acc_scale = l_i * 0 + alpha # workaround some compiler bug + acc *= acc_scale[:, None] + acc += tl.dot(p.to(tl.float16), v) + # -- update m_i and l_i -- + l_i = l_i * alpha + tl.sum(p, 1) + m_i = m_i_new + # update pointers + K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + # write back l and m + acc = acc / l_i[:, None] + #l_ptrs = L + off_hz * N_CTX + offs_m + #tl.store(l_ptrs, m_i + tl.math.log2(l_i)) + # write back O + O_block_ptr = tl.make_block_ptr(base=Out + qvk_offset, + shape=(N_CTX, BLOCK_DMODEL), + strides=(stride_om, stride_on), + offsets=(start_m * BLOCK_M, 0), + block_shape=(BLOCK_M, BLOCK_DMODEL), + order=(1, 0)) + tl.store(O_block_ptr, acc.to(tl.float16)) + + +class triton_flash_attn(torch.nn.Module): + + def __init__(self, ): + super(triton_flash_attn, self).__init__() + + def forward(self, q, k, v, sm_scale, block_128=True): + BLOCK = 128 if block_128 else 64 + # shape constraints + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] + o = torch.empty_like(q) + grid = (triton.cdiv(q.shape[2], BLOCK), q.shape[0] * q.shape[1]) + num_warps = 4 if Lk <= 64 else 8 + + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + o, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + o.stride(0), + o.stride(1), + o.stride(2), + o.stride(3), + k.shape[0], + k.shape[1], + k.shape[2], + BLOCK_M=BLOCK, + BLOCK_N=BLOCK, + BLOCK_DMODEL=Lk, + num_warps=num_warps, + num_stages=1, + ) + return o diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..bfd4d60dcb1ceafff2d87b62a78a8cf480ab9448 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/transformer.py @@ -0,0 +1,412 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import json +import math +import torch +from torch import nn +from torch.autograd import Function +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import TransformerBuilder, StochasticTransformerBuilder + +# Cuda modules will be imported if needed +transformer_cuda_module = None +stochastic_transformer_cuda_module = None + + +class TransformerConfig(): + + def __init__(self, batch_size, hidden_size, intermediate_size, heads, attn_dropout_ratio, hidden_dropout_ratio, + num_hidden_layers, initializer_range): + self.layer_id = -1 + self.batch_size = batch_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.heads = heads + self.attn_dropout_ratio = attn_dropout_ratio + self.hidden_dropout_ratio = hidden_dropout_ratio + self.num_hidden_layers = num_hidden_layers + self.initializer_range = initializer_range + + +class DeepSpeedTransformerConfig(TransformerConfig): + """Initialize the DeepSpeed Transformer Config. + + Arguments: + batch_size: The maximum batch size used for running the kernel on each GPU + + hidden_size: The hidden size of the transformer layer + + intermediate_size: The intermediate size of the feed-forward part of transformer layer + + heads: The number of heads in the self-attention of the transformer layer + + attn_dropout_ratio: The ratio of dropout for the attention's output + + hidden_dropout_ratio: The ratio of dropout for the transformer's output + + num_hidden_layers: The number of transformer layers + + initializer_range: BERT model's initializer range for initializing parameter data + + local_rank: Optional: The rank of GPU running the transformer kernel, it is not required + to use if the model already set the current device, otherwise need to set it + so that the transformer kernel can work on the right device + + seed: The random seed for the dropout layers + + fp16: Enable half-precision computation + + pre_layer_norm: Select between Pre-LN or Post-LN transformer architecture + + normalize_invertible: Optional: Enable invertible LayerNorm execution (dropping the input activation), + default is False + + gelu_checkpoint: Optional: Enable checkpointing of Gelu activation output to save memory, + default is False + + adjust_init_range: Optional: Set as True (default) if the model adjusts the weight initial values of + its self-attention output and layer output, False keeps the initializer_range no change. + See the adjustment below: + output_std = self.config.initializer_range / math.sqrt(2.0 * num_layers) + + attn_dropout_checkpoint: Optional: Enable checkpointing of attention dropout to save memory, + default is False + + stochastic_mode: Enable for high performance, please note that this flag has some level of + non-determinism and can produce different results on different runs. However, we have seen + that by enabling it, the pretraining tasks such as BERT are not affected and can obtain + a high accuracy level. On the other hand, for the downstream tasks, such as fine-tuning, we recommend + to turn it off in order to be able to reproduce the same result through the regular kernel execution. + + return_tuple: Enable if using the return_tuple interface style for sending out the forward results. + + training: Enable for training rather than inference. + """ + + def __init__(self, + batch_size=-1, + hidden_size=-1, + intermediate_size=-1, + heads=-1, + attn_dropout_ratio=-1, + hidden_dropout_ratio=-1, + num_hidden_layers=-1, + initializer_range=-1, + layer_norm_eps=1e-12, + local_rank=-1, + seed=-1, + fp16=False, + pre_layer_norm=True, + normalize_invertible=False, + gelu_checkpoint=False, + adjust_init_range=True, + attn_dropout_checkpoint=False, + stochastic_mode=False, + return_tuple=False, + training=True): + super(DeepSpeedTransformerConfig, + self).__init__(batch_size, hidden_size, + (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, + attn_dropout_ratio, hidden_dropout_ratio, num_hidden_layers, initializer_range) + self.fp16 = fp16 + self.pre_layer_norm = pre_layer_norm + self.local_rank = local_rank + self.seed = seed + self.normalize_invertible = normalize_invertible + self.gelu_checkpoint = gelu_checkpoint # True: if higher batch size is required + self.adjust_init_range = adjust_init_range + self.test_gemm = False + self.layer_norm_eps = layer_norm_eps + self.training = training + self.is_grad_enabled = True + self.attn_dropout_checkpoint = attn_dropout_checkpoint + self.stochastic_mode = stochastic_mode + self.return_tuple = return_tuple + + @classmethod + def from_dict(cls, json_object): + config = DeepSpeedTransformerConfig() + for key, value in json_object.items(): + config.__dict__[key] = value + return config + + @classmethod + def from_json_file(cls, json_file): + with open(json_file, "r", encoding='utf-16') as reader: + text = reader.read() + return cls.from_dict(json.loads(text)) + + +class DeepSpeedTransformerFunction(Function): + + @staticmethod + def forward(ctx, input, input_mask, self, grads, layer_id, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, + attn_nb, inter_w, inter_b, output_w, output_b, norm_w, norm_b, config): + + cuda_module = stochastic_transformer_cuda_module if config.stochastic_mode else transformer_cuda_module + forward_func = cuda_module.forward_fp16 if config.fp16 else cuda_module.forward_fp32 + + inp_size = input.size() + if inp_size[1] % 16 != 0: + input = torch.cat( + (input, + torch.randn( + (inp_size[0], (16 - (inp_size[1] % 16)), inp_size[2]), device=input.device, dtype=input.dtype)), + 1) + input_mask = torch.cat((input_mask, torch.ones((inp_size[0], input_mask.shape[1], input_mask.shape[2], \ + (16 - (inp_size[1] % 16))), device=input_mask.device, dtype=input_mask.dtype) * -10000), 3) + + (output, inp_norm, qkv_tf, soft_inp, ctx_bufB, attn_o_inp, add_res, ff1_inp, gelu_inp, ff2_inp, + attn_prob_dropout_mask, attn_output_dropout_mask, layer_output_dropout_mask, attn_layer_norm_var, + attn_layer_norm_mean, layer_norm_var, layer_norm_mean) = forward_func( + config.layer_id, input, input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, + inter_b, output_w, output_b, norm_w, norm_b, config.training and config.is_grad_enabled, + config.pre_layer_norm, config.attn_dropout_checkpoint, config.normalize_invertible, + config.gelu_checkpoint) + + # For testing only. + if grads is not None: + for i in [2]: + attn_qkvw.register_hook(lambda x, i=i, self=self: grads.append([ + x[i * attn_ow.size(0):(i + 1) * attn_ow.size(0)], ("Q_W" if i == 0 else "K_W" if i == 1 else "V_W") + ])) + for i in [2]: + attn_qkvb.register_hook(lambda x, i=i, self=self: grads.append([ + x[i * attn_ow.size(0):(i + 1) * attn_ow.size(0)], ("Q_B" if i == 0 else "K_B" if i == 1 else "V_B") + ])) + + attn_ow.register_hook(lambda x, self=self: grads.append([x, "O_W"])) + attn_ob.register_hook(lambda x, self=self: grads.append([x, "O_B"])) + attn_nw.register_hook(lambda x, self=self: grads.append([x, "N2_W"])) + attn_nb.register_hook(lambda x, self=self: grads.append([x, "N2_B"])) + inter_w.register_hook(lambda x, self=self: grads.append([x, "int_W"])) + inter_b.register_hook(lambda x, self=self: grads.append([x, "int_B"])) + output_w.register_hook(lambda x, self=self: grads.append([x, "out_W"])) + output_b.register_hook(lambda x, self=self: grads.append([x, "out_B"])) + norm_w.register_hook(lambda x, self=self: grads.append([x, "norm_W"])) + norm_b.register_hook(lambda x, self=self: grads.append([x, "norm_B"])) + + if config.is_grad_enabled and config.training: + if (config.pre_layer_norm and config.normalize_invertible): + ctx.save_for_backward(input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, + inter_b, output_w, output_b, norm_w, norm_b) + else: + ctx.save_for_backward(output, input, input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, + attn_nb, inter_w, inter_b, output_w, output_b, norm_w, norm_b) + + ctx.config = config + if (config.pre_layer_norm or not config.normalize_invertible): + ctx.inp_norm = inp_norm + + ctx.qkv_tf = qkv_tf + ctx.soft_inp = soft_inp + if not config.attn_dropout_checkpoint: + ctx.ctx_bufB = ctx_bufB + + ctx.attn_o_inp = attn_o_inp + if not config.normalize_invertible: + ctx.add_res = add_res + + ctx.attn_layer_norm_mean = attn_layer_norm_mean + ctx.layer_norm_mean = layer_norm_mean + + ctx.ff1_inp = ff1_inp + if not config.gelu_checkpoint: + ctx.gelu_inp = gelu_inp + + ctx.ff2_inp = ff2_inp + ctx.attn_prob_dropout_mask = attn_prob_dropout_mask + ctx.attn_output_dropout_mask = attn_output_dropout_mask + ctx.layer_output_dropout_mask = layer_output_dropout_mask + ctx.attn_layer_norm_var = attn_layer_norm_var + ctx.layer_norm_var = layer_norm_var + + if inp_size[1] % 16 != 0: + output = torch.narrow(output, 1, 0, inp_size[1]) + + if config.return_tuple: + return (output, ) # outputs -> (output) : outputs[0] = output + else: + return output + + @staticmethod + def backward(ctx, grad_output): + bsz = grad_output.shape[0] + grad_output_shape = grad_output.size() + if grad_output_shape[1] % 16 != 0: + grad_output = torch.cat((grad_output, torch.zeros((bsz, (16 - (grad_output_shape[1] % 16)), \ + grad_output_shape[2]), device=grad_output.device, dtype=grad_output.dtype)), 1) + + assert ctx.config.training + + if (ctx.config.pre_layer_norm and ctx.config.normalize_invertible): + (input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, inter_b, output_w, + output_b, norm_w, norm_b) = ctx.saved_tensors + else: + (output, input, input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, inter_b, + output_w, output_b, norm_w, norm_b) = ctx.saved_tensors + + cuda_module = stochastic_transformer_cuda_module if ctx.config.stochastic_mode else transformer_cuda_module + backward_func = cuda_module.backward_fp16 if ctx.config.fp16 else cuda_module.backward_fp32 + + (grad_input, grad_attn_qkvw, grad_attn_qkvb, grad_attn_ow, grad_attn_ob, grad_attn_nw, grad_attn_nb, + grad_inter_w, grad_inter_b, grad_output_w, grad_output_b, grad_norm_w, grad_norm_b) = backward_func( + ctx.config.layer_id, grad_output, + (ctx.inp_norm if (ctx.config.pre_layer_norm and ctx.config.normalize_invertible) else output), + (ctx.inp_norm if (ctx.config.pre_layer_norm or not ctx.config.normalize_invertible) else input), + ctx.qkv_tf, ctx.soft_inp, (ctx.soft_inp if ctx.config.attn_dropout_checkpoint else ctx.ctx_bufB), + ctx.attn_o_inp, (ctx.ff1_inp if ctx.config.normalize_invertible else ctx.add_res), ctx.ff1_inp, + (ctx.ff2_inp if ctx.config.gelu_checkpoint else ctx.gelu_inp), ctx.ff2_inp, ctx.attn_prob_dropout_mask, + ctx.attn_output_dropout_mask, ctx.layer_output_dropout_mask, ctx.attn_layer_norm_var, + ctx.attn_layer_norm_mean, ctx.layer_norm_var, ctx.layer_norm_mean, + (ctx.inp_norm if + (ctx.config.pre_layer_norm and ctx.config.normalize_invertible) else input), input_mask, attn_qkvw, + attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, inter_b, output_w, output_b, norm_w, norm_b) + + # This appears to be an effective way to release context memory + ctx.qkv_tf = None + ctx.soft_inp = None + ctx.ctx_bufB = None + ctx.gelu_inp = None + ctx.ff2_inp = None + ctx.attn_o_inp = None + ctx.ff1_inp = None + ctx.add_res = None + ctx.inp_norm = None + ctx.config = None + ctx.attn_layer_norm_mean = None + ctx.layer_norm_mean = None + ctx.attn_prob_dropout_mask = None + ctx.attn_output_dropout_mask = None + ctx.layer_output_dropout_mask = None + ctx.attn_layer_norm_var = None + ctx.layer_norm_var = None + + if grad_output_shape[1] % 16 != 0: + grad_input = torch.narrow(grad_input, 1, 0, grad_output_shape[1]) + + return (grad_input, None, None, None, None, grad_attn_qkvw, grad_attn_qkvb, grad_attn_ow, grad_attn_ob, + grad_attn_nw, grad_attn_nb, grad_inter_w, grad_inter_b, grad_output_w, grad_output_b, grad_norm_w, + grad_norm_b, None) + + +class DeepSpeedTransformerLayer(nn.Module): + """Initialize the DeepSpeed Transformer Layer. + + Static variable: + layer_id: The layer-index counter starting from 0 and incrementing by 1 every time a layer object is instantiated, + e.g. if a model has 24 transformer layers, layer_id goes from 0 to 23. + Arguments: + config: An object of DeepSpeedTransformerConfig + + initial_weights: Optional: Only used for unit test + + initial_biases: Optional: Only used for unit test + """ + layer_id = 0 + + def __init__(self, config, initial_weights=None, initial_biases=None): + super(DeepSpeedTransformerLayer, self).__init__() + + self.config = config + self.config.layer_id = DeepSpeedTransformerLayer.layer_id + DeepSpeedTransformerLayer.layer_id = DeepSpeedTransformerLayer.layer_id + 1 + + print("DeepSpeed Transformer config is ", self.config.__dict__) + + if self.config.local_rank >= 0: + get_accelerator().set_device(self.config.local_rank) + + if initial_weights is None and initial_biases is None: + self.attn_qkvw = nn.Parameter(torch.Tensor(self.config.hidden_size * 3, self.config.hidden_size)) + self.attn_qkvb = nn.Parameter(torch.Tensor(self.config.hidden_size * 3)) + self.attn_ow = nn.Parameter(torch.Tensor(self.config.hidden_size, self.config.hidden_size)) + self.attn_ob = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.attn_nw = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.attn_nb = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.inter_w = nn.Parameter(torch.Tensor(self.config.intermediate_size, self.config.hidden_size)) + self.inter_b = nn.Parameter(torch.Tensor(self.config.intermediate_size)) + self.output_w = nn.Parameter(torch.Tensor(self.config.hidden_size, self.config.intermediate_size)) + self.output_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.norm_w = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.norm_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) + self.init_transformer_weights(self.config.adjust_init_range) + else: + # For testing only. + q = initial_weights[0].data + k = initial_weights[1].data + v = initial_weights[2].data + + self.attn_qkvw = nn.Parameter(torch.cat((q, k, v))) + #self.attn_qkvw[i * self.config.hidden_size:(i + 1) * self.config.hidden_size] = \ + # initial_weights[i].clone() + #torch.empty_like(initial_weights[i]).data.copy_(initial_weights[i].data) + self.attn_qkvb = nn.Parameter(torch.Tensor(self.config.hidden_size * 3)) + self.attn_qkvb.data.zero_() + self.attn_ow = initial_weights[3] + self.attn_ob = initial_biases[3] + self.attn_nw = initial_weights[4] + self.attn_nb = initial_biases[4] + self.inter_w = initial_weights[5] + self.inter_b = initial_biases[5] + self.output_w = initial_weights[6] + self.output_b = initial_biases[6] + self.norm_w = initial_weights[7] + self.norm_b = initial_biases[7] + + # Load cuda modules if needed + global transformer_cuda_module, stochastic_transformer_cuda_module + if transformer_cuda_module is None and not self.config.stochastic_mode: + transformer_cuda_module = TransformerBuilder().load() + if stochastic_transformer_cuda_module is None and self.config.stochastic_mode: + stochastic_transformer_cuda_module = StochasticTransformerBuilder().load() + + # create the layer in cuda kernels. + cuda_module = stochastic_transformer_cuda_module if self.config.stochastic_mode else transformer_cuda_module + create_layer_func = cuda_module.create_transformer_layer_fp16 if self.config.fp16 else cuda_module.create_transformer_layer_fp32 + + create_layer_func(self.config.layer_id, self.config.batch_size, self.config.hidden_size, self.config.heads, + self.config.intermediate_size, self.config.attn_dropout_ratio, + self.config.hidden_dropout_ratio, self.config.layer_norm_eps, self.config.seed, + self.config.pre_layer_norm, self.config.test_gemm, self.config.attn_dropout_checkpoint, + self.config.normalize_invertible, self.config.gelu_checkpoint, self.config.stochastic_mode) + + def init_transformer_weights(self, adjust_init_range=False): + num_layers = self.config.num_hidden_layers + output_std = self.config.initializer_range + if adjust_init_range and self.config.local_rank == 0: + print("Accounting for accumulation on the residual path") + output_std = self.config.initializer_range / math.sqrt(2.0 * num_layers) + + self.attn_qkvw.data.normal_(mean=0.0, std=self.config.initializer_range) + self.attn_qkvb.data.zero_() + self.attn_ow.data.normal_(mean=0.0, std=output_std) + self.attn_ob.data.zero_() + self.attn_nw.data.fill_(1.0) + self.attn_nb.data.zero_() + self.inter_w.data.normal_(mean=0.0, std=self.config.initializer_range) + self.inter_b.data.zero_() + self.output_w.data.normal_(mean=0.0, std=output_std) + self.output_b.data.zero_() + self.norm_w.data.fill_(1.0) + self.norm_b.data.zero_() + + def forward(self, + hidden_states, + attention_mask=None, + head_mask=None, + layer_head_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_value=None, + output_attentions=False, + grads=None): + self.config.is_grad_enabled = torch.is_grad_enabled() + self.config.training = self.training + return DeepSpeedTransformerFunction.apply(hidden_states, attention_mask, self, grads, self.config.layer_id, + self.attn_qkvw, self.attn_qkvb, self.attn_ow, self.attn_ob, + self.attn_nw, self.attn_nb, self.inter_w, self.inter_b, + self.output_w, self.output_b, self.norm_w, self.norm_b, self.config) diff --git a/venv/lib/python3.10/site-packages/deepspeed/pipe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/pipe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f25471713b4568b43216fe16786a46904575ae5e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/pipe/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..runtime.pipe import PipelineModule, LayerSpec, TiedLayerSpec diff --git a/venv/lib/python3.10/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..587abed6296d262c3b55789e267ffd47109f8762 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/profiling/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06b12d70150cbf207db9598fad677733f1ab4ae9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c933b8b9b89694fa6d41125bd3d24327d1562b00 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c0459e6637a1ea570b7fb87c6d5a21fedd37fce Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/config.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/config.py new file mode 100644 index 0000000000000000000000000000000000000000..e4f06630ea6f6801883046efb5f61df3e7132f63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/profiling/config.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject +from deepspeed.profiling.constants import * + + +class DeepSpeedFlopsProfilerConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedFlopsProfilerConfig, self).__init__() + + self.enabled = None + self.recompute_fwd_factor = None + self.profile_step = None + self.module_depth = None + self.top_modules = None + + if FLOPS_PROFILER in param_dict.keys(): + flops_profiler_dict = param_dict[FLOPS_PROFILER] + else: + flops_profiler_dict = {} + + self._initialize(flops_profiler_dict) + + def _initialize(self, flops_profiler_dict): + self.enabled = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_ENABLED, FLOPS_PROFILER_ENABLED_DEFAULT) + + self.recompute_fwd_factor = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR, + FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT) + + self.profile_step = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_PROFILE_STEP, + FLOPS_PROFILER_PROFILE_STEP_DEFAULT) + + self.module_depth = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_MODULE_DEPTH, + FLOPS_PROFILER_MODULE_DEPTH_DEFAULT) + + self.top_modules = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_TOP_MODULES, + FLOPS_PROFILER_TOP_MODULES_DEFAULT) + + self.detailed = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_DETAILED, FLOPS_PROFILER_DETAILED_DEFAULT) + + self.output_file = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_OUTPUT_FILE, + FLOPS_PROFILER_OUTPUT_FILE_DEFAULT) diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/constants.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..0374303d7d969c76cd6771b2b18bf46f94df8d04 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/profiling/constants.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######################################### +# flops profiler +######################################### +# Flops profiler. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +FLOPS_PROFILER_FORMAT = ''' +flops profiler should be enabled as: +"session_params": { + "flops_profiler": { + "enabled": true, + "recompute_fwd_factor": 0.0, + "profile_step": 1, + "module_depth": -1, + "top_modules": 3, + "detailed": true, + "output_file": null + } +} +''' + +FLOPS_PROFILER = "flops_profiler" + +FLOPS_PROFILER_ENABLED = "enabled" +FLOPS_PROFILER_ENABLED_DEFAULT = False + +FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR = "recompute_fwd_factor" +FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT = 0.0 + +FLOPS_PROFILER_PROFILE_STEP = "profile_step" +FLOPS_PROFILER_PROFILE_STEP_DEFAULT = 1 + +FLOPS_PROFILER_MODULE_DEPTH = "module_depth" +FLOPS_PROFILER_MODULE_DEPTH_DEFAULT = -1 + +FLOPS_PROFILER_TOP_MODULES = "top_modules" +FLOPS_PROFILER_TOP_MODULES_DEFAULT = 1 + +FLOPS_PROFILER_DETAILED = "detailed" +FLOPS_PROFILER_DETAILED_DEFAULT = True + +FLOPS_PROFILER_OUTPUT_FILE = "output_file" +FLOPS_PROFILER_OUTPUT_FILE_DEFAULT = None diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..980d8a5e92fc4c3491f822479d05c5e2776233df --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .profiler import * diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1bdc72be53afd5521f231e53f84455751f809815 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7143ac310b1a352fe164b889b0f13bad0ae3a52c Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/profiler.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..1f051077c36c24b272f031d629ac5ee84f88aa3f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/profiler.py @@ -0,0 +1,1297 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import time +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial +from typing import List, Optional +from collections import OrderedDict +import numpy as np +from deepspeed.accelerator import get_accelerator +from deepspeed.utils import logger +from deepspeed.moe.layer import MoE +from deepspeed.utils.timer import FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, STEP_GLOBAL_TIMER +from deepspeed.utils.torch import required_torch_version +import einops + +Tensor = torch.Tensor + +module_flop_count = [] +module_mac_count = [] +old_functions = {} + +DEFAULT_PRECISION = 2 + + +class FlopsProfiler(object): + """Measures the latency, number of estimated floating-point operations and parameters of each module in a PyTorch model. + + The flops-profiler profiles the forward pass of a PyTorch model and prints the model graph with the measured profile attached to each module. It shows how latency, flops and parameters are spent in the model and which modules or layers could be the bottleneck. It also outputs the names of the top k modules in terms of aggregated latency, flops, and parameters at depth l with k and l specified by the user. The output profile is computed for each batch of input. + The DeepSpeed flops profiler can be used with the DeepSpeed runtime or as a standalone package. + When using DeepSpeed for model training, the flops profiler can be configured in the deepspeed_config file and no user code change is required. + + If using the profiler as a standalone package, one imports the flops_profiler package and use the APIs. + + Here is an example for usage in a typical training workflow: + + .. code-block:: python + + model = Model() + prof = FlopsProfiler(model) + + for step, batch in enumerate(data_loader): + if step == profile_step: + prof.start_profile() + + loss = model(batch) + + if step == profile_step: + flops = prof.get_total_flops(as_string=True) + params = prof.get_total_params(as_string=True) + prof.print_model_profile(profile_step=profile_step) + prof.end_profile() + + loss.backward() + optimizer.step() + + To profile a trained model in inference, use the `get_model_profile` API. + + Args: + object (torch.nn.Module): The PyTorch model to profile. + """ + + def __init__(self, model, ds_engine=None, recompute_fwd_factor=0.0): + self.model = model + self.ds_engine = ds_engine + self.recompute_fwd_factor = recompute_fwd_factor + self.started = False + self.func_patched = False + + def start_profile(self, ignore_list=None): + """Starts profiling. + + Extra attributes are added recursively to all the modules and the profiled torch.nn.functionals are monkey patched. + + Args: + ignore_list (list, optional): the list of modules to ignore while profiling. Defaults to None. + """ + logger.info("Flops profiler started") + self.reset_profile() + _patch_functionals() + _patch_tensor_methods() + _patch_miscellaneous_operations() + + def register_module_hooks(module, ignore_list): + if ignore_list and type(module) in ignore_list: + return + + # if computing the flops of a module directly + if type(module) in MODULE_HOOK_MAPPING: + if not hasattr(module, "__flops_handle__"): + module.__flops_handle__ = module.register_forward_hook(MODULE_HOOK_MAPPING[type(module)]) + return + + # if computing the flops of the functionals in a module + def pre_hook(module, input): + module_flop_count.append([]) + module_mac_count.append([]) + + if not hasattr(module, "__pre_hook_handle__"): + module.__pre_hook_handle__ = module.register_forward_pre_hook(pre_hook) + + def post_hook(module, input, output): + if module_flop_count: + module.__flops__ += sum([elem[1] for elem in module_flop_count[-1]]) + module_flop_count.pop() + module.__macs__ += sum([elem[1] for elem in module_mac_count[-1]]) + module_mac_count.pop() + + if not hasattr(module, "__post_hook_handle__"): + module.__post_hook_handle__ = module.register_forward_hook(post_hook) + + def start_time_hook(module, input): + get_accelerator().synchronize() + module.__start_time__ = time.time() + + if not hasattr(module, "__start_time_hook_handle__"): + module.__start_time_hook_handle__ = module.register_forward_pre_hook(start_time_hook) + + def end_time_hook(module, input, output): + get_accelerator().synchronize() + module.__duration__ += time.time() - module.__start_time__ + + if not hasattr(module, "__end_time_hook_handle__"): + module.__end_time_hook_handle__ = module.register_forward_hook(end_time_hook) + + self.model.apply(partial(register_module_hooks, ignore_list=ignore_list)) + self.started = True + self.func_patched = True + + def stop_profile(self): + """Stop profiling. + + All torch.nn.functionals are restored to their originals. + """ + if self.started and self.func_patched: + _reload_functionals() + _reload_tensor_methods() + _reload_miscellaneous_operations() + self.func_patched = False + + def remove_profile_attrs(module): + if hasattr(module, "__pre_hook_handle__"): + module.__pre_hook_handle__.remove() + del module.__pre_hook_handle__ + if hasattr(module, "__post_hook_handle__"): + module.__post_hook_handle__.remove() + del module.__post_hook_handle__ + if hasattr(module, "__flops_handle__"): + module.__flops_handle__.remove() + del module.__flops_handle__ + if hasattr(module, "__start_time_hook_handle__"): + module.__start_time_hook_handle__.remove() + del module.__start_time_hook_handle__ + if hasattr(module, "__end_time_hook_handle__"): + module.__end_time_hook_handle__.remove() + del module.__end_time_hook_handle__ + + self.model.apply(remove_profile_attrs) + + def reset_profile(self): + """Resets the profiling. + + Adds or resets the extra attributes. + """ + + def get_param_count_and_ep(param): + """ + Return the number of parameters in the layer, whether the layer is an MoE layer, + and its expert parallelism size if so + """ + prefix = 'ep_size_' + offset = len(prefix) + expert_parallelism = 0 + if getattr(param, "group_name", "").startswith(prefix): + try: + expert_parallelism = int(param.group_name[offset:]) + except ValueError: + pass + return param.numel(), expert_parallelism, param.element_size() + + def add_or_reset_attrs(module): + module.__flops__ = 0 + module.__macs__ = 0 + module.__params__ = module.__expert_params__ = module.__model_expert_params__ = 0 + parameters = (get_param_count_and_ep(p) for p in module.parameters()) + for num_params, expert_parallelism, per_param_size in parameters: + params = num_params if not expert_parallelism else 0 + expert_params = num_params if expert_parallelism else 0 + # number of expert parameters taking into account other expert parallel groups + model_expert_params = num_params * expert_parallelism + module.__params__ += params + module.__expert_params__ += expert_params + module.__model_expert_params__ += model_expert_params + module.__start_time__ = 0 + module.__duration__ = 0 + + self.model.apply(add_or_reset_attrs) + + def end_profile(self): + """Ends profiling. + + The added attributes and handles are removed recursively on all the modules. + """ + if not self.started: + return + self.stop_profile() + self.started = False + + def remove_profile_attrs(module): + if hasattr(module, "__flops__"): + del module.__flops__ + if hasattr(module, "__macs__"): + del module.__macs__ + if hasattr(module, "__params__"): + del module.__params__ + if hasattr(module, "__expert_params__"): + del module.__expert_params__ + if hasattr(module, "__model_expert_params__"): + del module.__model_expert_params__ + if hasattr(module, "__start_time__"): + del module.__start_time__ + if hasattr(module, "__duration__"): + del module.__duration__ + + self.model.apply(remove_profile_attrs) + logger.info("Flops profiler finished") + + def get_total_flops(self, as_string=False): + """Returns the total flops of the model. + + Args: + as_string (bool, optional): whether to output the flops as string. Defaults to False. + + Returns: + The number of multiply-accumulate operations of the model forward pass. + """ + total_flops = get_module_flops(self.model) + return number_to_string(total_flops) if as_string else total_flops + + def get_total_macs(self, as_string=False): + """Returns the total MACs of the model. + + Args: + as_string (bool, optional): whether to output the flops as string. Defaults to False. + + Returns: + The number of multiply-accumulate operations of the model forward pass. + """ + total_macs = get_module_macs(self.model) + return macs_to_string(total_macs) if as_string else total_macs + + def get_total_duration(self, as_string=False): + """Returns the total duration of the model forward pass. + + Args: + as_string (bool, optional): whether to output the duration as string. Defaults to False. + + Returns: + The latency of the model forward pass. + """ + total_duration = get_module_duration(self.model) + return duration_to_string(total_duration) if as_string else total_duration + + def get_total_params(self, as_string=False): + """Returns the total number of parameters stored per rank. + + Args: + as_string (bool, optional): whether to output the parameters as string. Defaults to False. + + Returns: + The total number of parameters stored per rank. + """ + total_params = self.model.__expert_params__ + self.model.__params__ + return params_to_string(total_params) if as_string else total_params + + def is_expert_tensor_parallelism_enabled(self): + for _, module in self.model.named_modules(): + if isinstance(module, MoE) and hasattr(module, 'enable_expert_tensor_parallelism'): + return module.enable_expert_tensor_parallelism + return False + + def print_model_profile(self, profile_step=1, module_depth=-1, top_modules=1, detailed=True, output_file=None): + """Prints the model graph with the measured profile attached to each module. + + Args: + profile_step (int, optional): The global training step at which to profile. Note that warm up steps are needed for accurate time measurement. + module_depth (int, optional): The depth of the model to which to print the aggregated module information. When set to -1, it prints information from the top to the innermost modules (the maximum depth). + top_modules (int, optional): Limits the aggregated profile output to the number of top modules specified. + detailed (bool, optional): Whether to print the detailed model profile. + output_file (str, optional): Path to the output file. If None, the profiler prints to stdout. + """ + if not self.started: + return + import sys + import os.path + original_stdout = None + f = None + if output_file and output_file != "": + dir_path = os.path.dirname(os.path.abspath(output_file)) + if not os.path.exists(dir_path): + os.makedirs(dir_path) + original_stdout = sys.stdout + f = open(output_file, "w") + sys.stdout = f + + total_flops = self.get_total_flops() + total_macs = self.get_total_macs() + total_duration = self.get_total_duration() + total_params = self.get_total_params() + expert_tensor_parallelism = None # silence the linters + total_model_expert_params = total_model_nonexpert_params = 0 + if self.ds_engine: + total_model_nonexpert_params = self.model.__params__ * self.ds_engine.mp_world_size + if self.ds_engine.has_moe_layers: + expert_tensor_parallelism = self.ds_engine.mp_world_size if self.is_expert_tensor_parallelism_enabled( + ) else 1 + total_model_expert_params = self.model.__model_expert_params__ * expert_tensor_parallelism + + self.flops = total_flops + self.macs = total_macs + self.params = total_params + + print("\n-------------------------- DeepSpeed Flops Profiler --------------------------") + print(f'Profile Summary at step {profile_step}:') + print("Notations:\n" + "data parallel size (dp_size), model parallel size(mp_size),\n" + "number of parameters (params), number of multiply-accumulate operations(MACs),\n" + "number of floating-point operations (flops), floating-point operations per second (FLOPS),\n" + "fwd latency (forward propagation latency), bwd latency (backward propagation latency),\n" + "step (weights update latency), iter latency (sum of fwd, bwd and step latency)\n") + line_fmt = '{:<70} {:<8}' + if self.ds_engine: + print(line_fmt.format('world size: ', self.ds_engine.world_size)) + print(line_fmt.format('data parallel size: ', self.ds_engine.dp_world_size)) + print(line_fmt.format('model parallel size: ', self.ds_engine.mp_world_size)) + print(line_fmt.format('batch size per GPU: ', self.ds_engine.train_micro_batch_size_per_gpu())) + if self.ds_engine.has_moe_layers: + print(line_fmt.format('expert tensor parallelism enabled: ', expert_tensor_parallelism > 1)) + + print(line_fmt.format('params per GPU: ', params_to_string(total_params))) + if total_model_expert_params > 0: + print( + line_fmt.format('params of model: ', + params_to_string(total_model_nonexpert_params + total_model_expert_params))) + print(line_fmt.format(' non-expert params of model: ', params_to_string(total_model_nonexpert_params))) + print(line_fmt.format(' expert params of model: ', params_to_string(total_model_expert_params))) + else: + print( + line_fmt.format('params of model = params per GPU * mp_size: ', + params_to_string(total_model_nonexpert_params))) + + print(line_fmt.format('fwd MACs per GPU: ', macs_to_string(total_macs))) + + print(line_fmt.format('fwd flops per GPU: ', number_to_string(total_flops))) + + print( + line_fmt.format('fwd flops of model = fwd flops per GPU * mp_size: ', + number_to_string(total_flops * (self.ds_engine.mp_world_size if self.ds_engine else 1)))) + + fwd_latency = self.get_total_duration() + if self.ds_engine and self.ds_engine.wall_clock_breakdown(): + fwd_latency = self.ds_engine.timers(FORWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 + print(line_fmt.format('fwd latency: ', duration_to_string(fwd_latency))) + print( + line_fmt.format('fwd FLOPS per GPU = fwd flops per GPU / fwd latency: ', + flops_to_string(total_flops / fwd_latency))) + + if self.ds_engine and self.ds_engine.wall_clock_breakdown(): + bwd_factor = 2 + self.recompute_fwd_factor + bwd_latency = self.ds_engine.timers(BACKWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 + step_latency = self.ds_engine.timers(STEP_GLOBAL_TIMER).elapsed(False) / 1000.0 + print(line_fmt.format('bwd latency: ', duration_to_string(bwd_latency))) + print( + line_fmt.format(f'bwd FLOPS per GPU = {bwd_factor:g} * fwd flops per GPU / bwd latency: ', + flops_to_string(bwd_factor * total_flops / bwd_latency))) + print( + line_fmt.format( + f'fwd+bwd FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / (fwd+bwd latency): ', + flops_to_string((bwd_factor + 1) * total_flops / (fwd_latency + bwd_latency)))) + + print(line_fmt.format('step latency: ', duration_to_string(step_latency))) + + iter_latency = fwd_latency + bwd_latency + step_latency + print(line_fmt.format('iter latency: ', duration_to_string(iter_latency))) + print( + line_fmt.format(f'FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / iter latency: ', + flops_to_string((bwd_factor + 1) * total_flops / iter_latency))) + + samples_per_iter = self.ds_engine.train_micro_batch_size_per_gpu() * self.ds_engine.world_size + print(line_fmt.format('samples/second: ', round(samples_per_iter / iter_latency, DEFAULT_PRECISION))) + + def flops_repr(module): + params = module.__params__ + module.__expert_params__ + flops = get_module_flops(module) + macs = get_module_macs(module) + duration = get_module_duration(module) + items = [ + "{} = {:g}% Params".format( + params_to_string(params), + round(100 * params / total_params, DEFAULT_PRECISION) if total_params else 0), + "{} = {:g}% MACs".format(macs_to_string(macs), + round(100 * macs / total_macs, DEFAULT_PRECISION) if total_macs else 0), + "{} = {:g}% latency".format( + duration_to_string(duration), + round(100 * duration / total_duration, DEFAULT_PRECISION) if total_duration else 0), + flops_to_string(round(flops / duration, DEFAULT_PRECISION) if duration else 0), + ] + original_extra_repr = module.original_extra_repr() + if original_extra_repr: + items.append(original_extra_repr) + return ", ".join(items) + + def add_extra_repr(module): + flops_extra_repr = flops_repr.__get__(module) + if module.extra_repr != flops_extra_repr: + module.original_extra_repr = module.extra_repr + module.extra_repr = flops_extra_repr + assert module.extra_repr != module.original_extra_repr + + def del_extra_repr(module): + if hasattr(module, "original_extra_repr"): + module.extra_repr = module.original_extra_repr + del module.original_extra_repr + + self.model.apply(add_extra_repr) + + print("\n----------------------------- Aggregated Profile per GPU -----------------------------") + self.print_model_aggregated_profile(module_depth=module_depth, top_modules=top_modules) + + if detailed: + print("\n------------------------------ Detailed Profile per GPU ------------------------------") + print( + "Each module profile is listed after its name in the following order: \nparams, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS" + ) + print( + "\nNote: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). They are not counted as submodules, thus not to be printed out. However they make up the difference between a parent's MACs (or latency) and the sum of its submodules'.\n2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.\n3. The fwd latency listed in the top module's profile is directly captured at the module forward function in PyTorch, thus it's less than the fwd latency shown above which is captured in DeepSpeed.\n" + ) + print(self.model) + + self.model.apply(del_extra_repr) + + print("------------------------------------------------------------------------------") + + if output_file: + sys.stdout = original_stdout + f.close() + + def print_model_aggregated_profile(self, module_depth=-1, top_modules=1): + """Prints the names of the top top_modules modules in terms of aggregated time, flops, and parameters at depth module_depth. + + Args: + module_depth (int, optional): the depth of the modules to show. Defaults to -1 (the innermost modules). + top_modules (int, optional): the number of top modules to show. Defaults to 1. + """ + info = {} + if not hasattr(self.model, "__flops__"): + print("no __flops__ attribute in the model, call this function after start_profile and before end_profile") + return + + def walk_module(module, curr_depth, info): + if curr_depth not in info: + info[curr_depth] = {} + if module.__class__.__name__ not in info[curr_depth]: + info[curr_depth][module.__class__.__name__] = [ + 0, + 0, + 0, + ] # macs, params, time + info[curr_depth][module.__class__.__name__][0] += get_module_macs(module) + info[curr_depth][module.__class__.__name__][1] += module.__params__ + module.__expert_params__ + info[curr_depth][module.__class__.__name__][2] += get_module_duration(module) + has_children = len(module._modules.items()) != 0 + if has_children: + for child in module.children(): + walk_module(child, curr_depth + 1, info) + + walk_module(self.model, 0, info) + + depth = module_depth + if module_depth == -1: + depth = len(info) - 1 + + print(f'Top {top_modules} modules in terms of params, MACs or fwd latency at different model depths:') + + for d in range(depth): + num_items = min(top_modules, len(info[d])) + + sort_macs = { + k: macs_to_string(v[0]) + for k, v in sorted(info[d].items(), key=lambda item: item[1][0], reverse=True)[:num_items] + } + sort_params = { + k: params_to_string(v[1]) + for k, v in sorted(info[d].items(), key=lambda item: item[1][1], reverse=True)[:num_items] + } + sort_time = { + k: duration_to_string(v[2]) + for k, v in sorted(info[d].items(), key=lambda item: item[1][2], reverse=True)[:num_items] + } + + print(f"depth {d}:") + print(f" params - {sort_params}") + print(f" MACs - {sort_macs}") + print(f" fwd latency - {sort_time}") + + +def _prod(dims): + p = 1 + for v in dims: + p *= v + return p + + +def _linear_flops_compute(input, weight, bias=None): + out_features = weight.shape[0] + macs = input.numel() * out_features + return 2 * macs, macs + + +def _relu_flops_compute(input, inplace=False): + return input.numel(), 0 + + +def _prelu_flops_compute(input: Tensor, weight: Tensor): + return input.numel(), 0 + + +def _elu_flops_compute(input: Tensor, alpha: float = 1.0, inplace: bool = False): + return input.numel(), 0 + + +def _leaky_relu_flops_compute(input: Tensor, negative_slope: float = 0.01, inplace: bool = False): + return input.numel(), 0 + + +def _relu6_flops_compute(input: Tensor, inplace: bool = False): + return input.numel(), 0 + + +def _silu_flops_compute(input: Tensor, inplace: bool = False): + return input.numel(), 0 + + +def _gelu_flops_compute(input, **kwargs): + return input.numel(), 0 + + +def _pool_flops_compute(input, + kernel_size, + stride=None, + padding=0, + dilation=None, + ceil_mode=False, + count_include_pad=True, + divisor_override=None, + return_indices=None): + return input.numel(), 0 + + +def _conv_flops_compute(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): + assert weight.shape[1] * groups == input.shape[1] + + batch_size = input.shape[0] + in_channels = input.shape[1] + out_channels = weight.shape[0] + kernel_dims = list(weight.shape[2:]) + input_dims = list(input.shape[2:]) + + length = len(input_dims) + + strides = stride if type(stride) is tuple else (stride, ) * length + dilations = dilation if type(dilation) is tuple else (dilation, ) * length + if isinstance(padding, str): + if padding == 'valid': + paddings = (0, ) * length + elif padding == 'same': + paddings = () + for d, k in zip(dilations, kernel_dims): + total_padding = d * (k - 1) + paddings += (total_padding // 2, ) + elif isinstance(padding, tuple): + paddings = padding + else: + paddings = (padding, ) * length + + output_dims = [] + for idx, input_dim in enumerate(input_dims): + output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * + (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1 + output_dims.append(output_dim) + + filters_per_channel = out_channels // groups + conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel + active_elements_count = batch_size * int(_prod(output_dims)) + overall_conv_macs = conv_per_position_macs * active_elements_count + overall_conv_flops = 2 * overall_conv_macs + + bias_flops = 0 + if bias is not None: + bias_flops = out_channels * active_elements_count + + return int(overall_conv_flops + bias_flops), int(overall_conv_macs) + + +def _conv_trans_flops_compute( + input, + weight, + bias=None, + stride=1, + padding=0, + output_padding=0, + groups=1, + dilation=1, +): + batch_size = input.shape[0] + in_channels = input.shape[1] + out_channels = weight.shape[1] + kernel_dims = list(weight.shape[2:]) + input_dims = list(input.shape[2:]) + + length = len(input_dims) + + paddings = padding if type(padding) is tuple else (padding, ) * length + strides = stride if type(stride) is tuple else (stride, ) * length + dilations = dilation if type(dilation) is tuple else (dilation, ) * length + + output_dims = [] + for idx, input_dim in enumerate(input_dims): + + output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * + (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1 + output_dims.append(output_dim) + + paddings = padding if type(padding) is tuple else (padding, padding) + strides = stride if type(stride) is tuple else (stride, stride) + dilations = dilation if type(dilation) is tuple else (dilation, dilation) + + filters_per_channel = out_channels // groups + conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel + active_elements_count = batch_size * int(_prod(input_dims)) + overall_conv_macs = conv_per_position_macs * active_elements_count + overall_conv_flops = 2 * overall_conv_macs + + bias_flops = 0 + if bias is not None: + bias_flops = out_channels * batch_size * int(_prod(output_dims)) + + return int(overall_conv_flops + bias_flops), int(overall_conv_macs) + + +def _batch_norm_flops_compute( + input, + running_mean, + running_var, + weight=None, + bias=None, + training=False, + momentum=0.1, + eps=1e-05, +): + has_affine = weight is not None + if training: + # estimation + return input.numel() * (5 if has_affine else 4), 0 + flops = input.numel() * (2 if has_affine else 1) + return flops, 0 + + +def _layer_norm_flops_compute( + input: Tensor, + normalized_shape: List[int], + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + eps: float = 1e-5, +): + has_affine = weight is not None + # estimation + return input.numel() * (5 if has_affine else 4), 0 + + +def _group_norm_flops_compute(input: Tensor, + num_groups: int, + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + eps: float = 1e-5): + has_affine = weight is not None + # estimation + return input.numel() * (5 if has_affine else 4), 0 + + +def _instance_norm_flops_compute( + input: Tensor, + running_mean: Optional[Tensor] = None, + running_var: Optional[Tensor] = None, + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + use_input_stats: bool = True, + momentum: float = 0.1, + eps: float = 1e-5, +): + has_affine = weight is not None + # estimation + return input.numel() * (5 if has_affine else 4), 0 + + +def _upsample_flops_compute(*args, **kwargs): + input = args[0] + size = kwargs.get('size', None) + if size is None and len(args) > 1: + size = args[1] + + if size is not None: + if isinstance(size, tuple) or isinstance(size, list): + return int(_prod(size)), 0 + else: + return int(size), 0 + + scale_factor = kwargs.get('scale_factor', None) + if scale_factor is None and len(args) > 2: + scale_factor = args[2] + assert scale_factor is not None, "either size or scale_factor should be defined" + + flops = input.numel() + if isinstance(scale_factor, (list, tuple)): + # see documentation of `F.interpolate` + # the spatial dims are defined as the last `n-2` dims of the tensor + assert len(scale_factor) == input.ndim - 2 + flops *= _prod(scale_factor) + else: + flops *= scale_factor**(input.ndim - 2) + + return int(flops), 0 + + +def _softmax_flops_compute(input, dim=None, _stacklevel=3, dtype=None): + return input.numel(), 0 + + +def _embedding_flops_compute( + input, + weight, + padding_idx=None, + max_norm=None, + norm_type=2.0, + scale_grad_by_freq=False, + sparse=False, +): + return 0, 0 + + +def _dropout_flops_compute(input, p=0.5, training=True, inplace=False): + return 0, 0 + + +def _matmul_flops_compute(input, other, *, out=None): + """ + Count flops for the matmul operation. + """ + macs = _prod(input.shape) * other.shape[-1] + return 2 * macs, macs + + +def _addmm_flops_compute(input, mat1, mat2, *, beta=1, alpha=1, out=None): + """ + Count flops for the addmm operation. + """ + macs = _prod(mat1.shape) * mat2.shape[-1] + return 2 * macs + _prod(input.shape), macs + + +def _einsum_flops_compute(equation, *operands): + """ + Count flops for the einsum operation. + """ + equation = equation.replace(" ", "") + input_shapes = [o.shape for o in operands] + + # Re-map equation so that same equation with different alphabet + # representations will look the same. + letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys() + mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)} + equation = equation.translate(mapping) + + np_arrs = [np.zeros(s) for s in input_shapes] + optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1] + for line in optim.split("\n"): + if "optimized flop" in line.lower(): + flop = int(float(line.split(":")[-1])) + return flop, 0 + raise NotImplementedError("Unsupported einsum operation.") + + +def _einops_einsum_flops_compute(*args): + """ + Count flops for the einops.einsum operation. + """ + *operands, equation = args + input_shapes = [o.shape for o in operands] + + # Re-map equation so that same equation with different alphabet + # representations will look the same. + letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys() + mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)} + equation = equation.translate(mapping) + + np_arrs = [np.zeros(s) for s in input_shapes] + optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1] + for line in optim.split("\n"): + if "optimized flop" in line.lower(): + flop = int(float(line.split(":")[-1])) + return flop, 0 + + raise NotImplementedError("Unsupported einops.einsum operation.") + + +def _tensor_addmm_flops_compute(self, mat1, mat2, *, beta=1, alpha=1, out=None): + """ + Count flops for the tensor addmm operation. + """ + macs = _prod(mat1.shape) * mat2.shape[-1] + return 2 * macs + _prod(self.shape), macs + + +def _mul_flops_compute(input, other, *, out=None): + return _elementwise_flops_compute(input, other) + + +def _add_flops_compute(input, other, *, alpha=1, out=None): + return _elementwise_flops_compute(input, other) + + +def _elementwise_flops_compute(input, other): + if not torch.is_tensor(input): + if torch.is_tensor(other): + return _prod(other.shape), 0 + else: + return 1, 0 + elif not torch.is_tensor(other): + return _prod(input.shape), 0 + else: + dim_input = len(input.shape) + dim_other = len(other.shape) + max_dim = max(dim_input, dim_other) + + final_shape = [] + for i in range(max_dim): + in_i = input.shape[i] if i < dim_input else 1 + ot_i = other.shape[i] if i < dim_other else 1 + if in_i > ot_i: + final_shape.append(in_i) + else: + final_shape.append(ot_i) + flops = _prod(final_shape) + return flops, 0 + + +def _attn_flops_compute(q, k, v, *args, **kwargs): + """ + Count flops for the scaled_dot_product_attention operation. + """ + macs = _prod(q.shape) * k.shape[-2] + macs += _prod(q.shape[:-1]) * k.shape[-2] * v.shape[-1] + return 2 * macs, macs + + +def wrapFunc(func, funcFlopCompute): + oldFunc = func + name = func.__str__ + old_functions[name] = oldFunc + + def newFunc(*args, **kwds): + flops, macs = funcFlopCompute(*args, **kwds) + if module_flop_count: + module_flop_count[-1].append((name, flops)) + if module_mac_count and macs: + module_mac_count[-1].append((name, macs)) + return oldFunc(*args, **kwds) + + newFunc.__str__ = func.__str__ + + return newFunc + + +def _patch_functionals(): + # FC + F.linear = wrapFunc(F.linear, _linear_flops_compute) + + # convolutions + F.conv1d = wrapFunc(F.conv1d, _conv_flops_compute) + F.conv2d = wrapFunc(F.conv2d, _conv_flops_compute) + F.conv3d = wrapFunc(F.conv3d, _conv_flops_compute) + + # conv transposed + F.conv_transpose1d = wrapFunc(F.conv_transpose1d, _conv_trans_flops_compute) + F.conv_transpose2d = wrapFunc(F.conv_transpose2d, _conv_trans_flops_compute) + F.conv_transpose3d = wrapFunc(F.conv_transpose3d, _conv_trans_flops_compute) + + # activations + F.relu = wrapFunc(F.relu, _relu_flops_compute) + F.prelu = wrapFunc(F.prelu, _prelu_flops_compute) + F.elu = wrapFunc(F.elu, _elu_flops_compute) + F.leaky_relu = wrapFunc(F.leaky_relu, _leaky_relu_flops_compute) + F.relu6 = wrapFunc(F.relu6, _relu6_flops_compute) + if hasattr(F, "silu"): + F.silu = wrapFunc(F.silu, _silu_flops_compute) + F.gelu = wrapFunc(F.gelu, _gelu_flops_compute) + + # Normalizations + F.batch_norm = wrapFunc(F.batch_norm, _batch_norm_flops_compute) + F.layer_norm = wrapFunc(F.layer_norm, _layer_norm_flops_compute) + F.instance_norm = wrapFunc(F.instance_norm, _instance_norm_flops_compute) + F.group_norm = wrapFunc(F.group_norm, _group_norm_flops_compute) + + # poolings + F.avg_pool1d = wrapFunc(F.avg_pool1d, _pool_flops_compute) + F.avg_pool2d = wrapFunc(F.avg_pool2d, _pool_flops_compute) + F.avg_pool3d = wrapFunc(F.avg_pool3d, _pool_flops_compute) + F.max_pool1d = wrapFunc(F.max_pool1d, _pool_flops_compute) + F.max_pool2d = wrapFunc(F.max_pool2d, _pool_flops_compute) + F.max_pool3d = wrapFunc(F.max_pool3d, _pool_flops_compute) + F.adaptive_avg_pool1d = wrapFunc(F.adaptive_avg_pool1d, _pool_flops_compute) + F.adaptive_avg_pool2d = wrapFunc(F.adaptive_avg_pool2d, _pool_flops_compute) + F.adaptive_avg_pool3d = wrapFunc(F.adaptive_avg_pool3d, _pool_flops_compute) + F.adaptive_max_pool1d = wrapFunc(F.adaptive_max_pool1d, _pool_flops_compute) + F.adaptive_max_pool2d = wrapFunc(F.adaptive_max_pool2d, _pool_flops_compute) + F.adaptive_max_pool3d = wrapFunc(F.adaptive_max_pool3d, _pool_flops_compute) + + # upsample + F.upsample = wrapFunc(F.upsample, _upsample_flops_compute) + F.interpolate = wrapFunc(F.interpolate, _upsample_flops_compute) + + # softmax + F.softmax = wrapFunc(F.softmax, _softmax_flops_compute) + + # embedding + F.embedding = wrapFunc(F.embedding, _embedding_flops_compute) + + # attn - scaled_dot_product_attention added in torch 2.0+ + if required_torch_version(min_version=2.0): + F.scaled_dot_product_attention = wrapFunc(F.scaled_dot_product_attention, _attn_flops_compute) + + +def _patch_tensor_methods(): + torch.matmul = wrapFunc(torch.matmul, _matmul_flops_compute) + torch.Tensor.matmul = wrapFunc(torch.Tensor.matmul, _matmul_flops_compute) + torch.Tensor.__matmul__ = wrapFunc(torch.Tensor.__matmul__, _matmul_flops_compute) + torch.mm = wrapFunc(torch.mm, _matmul_flops_compute) + torch.Tensor.mm = wrapFunc(torch.Tensor.mm, _matmul_flops_compute) + torch.bmm = wrapFunc(torch.bmm, _matmul_flops_compute) + torch.Tensor.bmm = wrapFunc(torch.Tensor.bmm, _matmul_flops_compute) + + torch.addmm = wrapFunc(torch.addmm, _addmm_flops_compute) + torch.Tensor.addmm = wrapFunc(torch.Tensor.addmm, _tensor_addmm_flops_compute) + + torch.mul = wrapFunc(torch.mul, _mul_flops_compute) + torch.Tensor.mul = wrapFunc(torch.Tensor.mul, _mul_flops_compute) + + torch.add = wrapFunc(torch.add, _add_flops_compute) + torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute) + + torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) + + torch.baddbmm = wrapFunc(torch.baddbmm, _tensor_addmm_flops_compute) + + +def _patch_miscellaneous_operations(): + einops.einsum = wrapFunc(einops.einsum, _einops_einsum_flops_compute) + + +def _reload_functionals(): + # torch.nn.functional does not support importlib.reload() + F.linear = old_functions[F.linear.__str__] + F.conv1d = old_functions[F.conv1d.__str__] + F.conv2d = old_functions[F.conv2d.__str__] + F.conv3d = old_functions[F.conv3d.__str__] + F.conv_transpose1d = old_functions[F.conv_transpose1d.__str__] + F.conv_transpose2d = old_functions[F.conv_transpose2d.__str__] + F.conv_transpose3d = old_functions[F.conv_transpose3d.__str__] + F.relu = old_functions[F.relu.__str__] + F.prelu = old_functions[F.prelu.__str__] + F.elu = old_functions[F.elu.__str__] + F.leaky_relu = old_functions[F.leaky_relu.__str__] + F.relu6 = old_functions[F.relu6.__str__] + if hasattr(F, "silu"): + F.silu = old_functions[F.silu.__str__] + F.gelu = old_functions[F.gelu.__str__] + F.batch_norm = old_functions[F.batch_norm.__str__] + F.layer_norm = old_functions[F.layer_norm.__str__] + F.instance_norm = old_functions[F.instance_norm.__str__] + F.group_norm = old_functions[F.group_norm.__str__] + F.avg_pool1d = old_functions[F.avg_pool1d.__str__] + F.avg_pool2d = old_functions[F.avg_pool2d.__str__] + F.avg_pool3d = old_functions[F.avg_pool3d.__str__] + F.max_pool1d = old_functions[F.max_pool1d.__str__] + F.max_pool2d = old_functions[F.max_pool2d.__str__] + F.max_pool3d = old_functions[F.max_pool3d.__str__] + F.adaptive_avg_pool1d = old_functions[F.adaptive_avg_pool1d.__str__] + F.adaptive_avg_pool2d = old_functions[F.adaptive_avg_pool2d.__str__] + F.adaptive_avg_pool3d = old_functions[F.adaptive_avg_pool3d.__str__] + F.adaptive_max_pool1d = old_functions[F.adaptive_max_pool1d.__str__] + F.adaptive_max_pool2d = old_functions[F.adaptive_max_pool2d.__str__] + F.adaptive_max_pool3d = old_functions[F.adaptive_max_pool3d.__str__] + F.upsample = old_functions[F.upsample.__str__] + F.interpolate = old_functions[F.interpolate.__str__] + F.softmax = old_functions[F.softmax.__str__] + F.embedding = old_functions[F.embedding.__str__] + + +def _reload_tensor_methods(): + torch.matmul = old_functions[torch.matmul.__str__] + torch.Tensor.matmul = old_functions[torch.Tensor.matmul.__str__] + torch.mm = old_functions[torch.mm.__str__] + torch.Tensor.mm = old_functions[torch.Tensor.mm.__str__] + torch.bmm = old_functions[torch.matmul.__str__] + torch.Tensor.bmm = old_functions[torch.Tensor.bmm.__str__] + torch.addmm = old_functions[torch.addmm.__str__] + torch.Tensor.addmm = old_functions[torch.Tensor.addmm.__str__] + torch.mul = old_functions[torch.mul.__str__] + torch.Tensor.mul = old_functions[torch.Tensor.mul.__str__] + torch.add = old_functions[torch.add.__str__] + torch.Tensor.add = old_functions[torch.Tensor.add.__str__] + + torch.einsum = old_functions[torch.einsum.__str__] + + torch.baddbmm = old_functions[torch.baddbmm.__str__] + + +def _reload_miscellaneous_operations(): + einops.einsum = old_functions[einops.einsum.__str__] + + +def _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size): + gates_size = w_ih.shape[0] + # matrix matrix mult ih state and internal state + flops += 2 * w_ih.shape[0] * w_ih.shape[1] - gates_size + # matrix matrix mult hh state and internal state + flops += 2 * w_hh.shape[0] * w_hh.shape[1] - gates_size + if isinstance(rnn_module, (nn.RNN, nn.RNNCell)): + # add both operations + flops += rnn_module.hidden_size + elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)): + # hadamard of r + flops += rnn_module.hidden_size + # adding operations from both states + flops += rnn_module.hidden_size * 3 + # last two hadamard _product and add + flops += rnn_module.hidden_size * 3 + elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)): + # adding operations from both states + flops += rnn_module.hidden_size * 4 + # two hadamard _product and add for C state + flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size + # final hadamard + flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size + return flops + + +def _rnn_forward_hook(rnn_module, input, output): + flops = 0 + # input is a tuple containing a sequence to process and (optionally) hidden state + inp = input[0] + batch_size = inp.shape[0] + seq_length = inp.shape[1] + num_layers = rnn_module.num_layers + + for i in range(num_layers): + w_ih = rnn_module.__getattr__("weight_ih_l" + str(i)) + w_hh = rnn_module.__getattr__("weight_hh_l" + str(i)) + if i == 0: + input_size = rnn_module.input_size + else: + input_size = rnn_module.hidden_size + flops = _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size) + if rnn_module.bias: + b_ih = rnn_module.__getattr__("bias_ih_l" + str(i)) + b_hh = rnn_module.__getattr__("bias_hh_l" + str(i)) + flops += b_ih.shape[0] + b_hh.shape[0] + + flops *= batch_size + flops *= seq_length + if rnn_module.bidirectional: + flops *= 2 + rnn_module.__flops__ += int(flops) + + +def _rnn_cell_forward_hook(rnn_cell_module, input, output): + flops = 0 + inp = input[0] + batch_size = inp.shape[0] + w_ih = rnn_cell_module.__getattr__("weight_ih") + w_hh = rnn_cell_module.__getattr__("weight_hh") + input_size = inp.shape[1] + flops = _rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size) + if rnn_cell_module.bias: + b_ih = rnn_cell_module.__getattr__("bias_ih") + b_hh = rnn_cell_module.__getattr__("bias_hh") + flops += b_ih.shape[0] + b_hh.shape[0] + + flops *= batch_size + rnn_cell_module.__flops__ += int(flops) + + +MODULE_HOOK_MAPPING = { + # RNN + nn.RNN: _rnn_forward_hook, + nn.GRU: _rnn_forward_hook, + nn.LSTM: _rnn_forward_hook, + nn.RNNCell: _rnn_cell_forward_hook, + nn.LSTMCell: _rnn_cell_forward_hook, + nn.GRUCell: _rnn_cell_forward_hook, +} + + +def macs_to_string(macs, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(macs, units=units, precision=precision)}MACs" + + +def number_to_string(num, units=None, precision=DEFAULT_PRECISION): + if units is None: + if num >= 1e12: + magnitude, units = 1e12, "T" + elif num >= 1e9: + magnitude, units = 1e9, "G" + elif num >= 1e6: + magnitude, units = 1e6, "M" + elif num >= 1e3: + magnitude, units = 1e3, "K" + elif num >= 1 or num == 0: + magnitude, units = 1, "" + elif num >= 1e-3: + magnitude, units = 1e-3, "m" + else: + magnitude, units = 1e-6, "u" + else: + if units == "T": + magnitude = 1e12 + elif units == "G": + magnitude = 1e9 + elif units == "M": + magnitude = 1e6 + elif units == "K": + magnitude = 1e3 + elif units == "m": + magnitude = 1e-3 + elif units == "u": + magnitude = 1e-6 + else: + magnitude = 1 + return f"{round(num / magnitude, precision):g} {units}" + + +def flops_to_string(flops, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(flops, units=units, precision=precision)}FLOPS" + + +def bytes_to_string(b, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(b, units=units, precision=precision)}B" + + +def params_to_string(params_num, units=None, precision=DEFAULT_PRECISION): + units = units.replace("B", "G") if units else units + return number_to_string(params_num, units=units, precision=precision).replace("G", "B").strip() + + +def duration_to_string(duration, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(duration, units=units, precision=precision)}s" + + + # can not iterate over all submodules using self.model.modules() + # since modules() returns duplicate modules only once +def get_module_flops(module): + sum = module.__flops__ + # iterate over immediate children modules + for child in module.children(): + sum += get_module_flops(child) + return sum + + +def get_module_macs(module): + sum = module.__macs__ + # iterate over immediate children modules + for child in module.children(): + sum += get_module_macs(child) + return sum + + +def get_module_duration(module): + duration = module.__duration__ + if duration == 0: # e.g. ModuleList + for m in module.children(): + duration += get_module_duration(m) + return duration + + +def get_model_profile(model, + input_shape=None, + args=[], + kwargs={}, + print_profile=True, + detailed=True, + module_depth=-1, + top_modules=1, + warm_up=1, + as_string=True, + output_file=None, + ignore_modules=None, + mode='forward'): + """Returns the total floating-point operations, MACs, and parameters of a model. + + Example: + + .. code-block:: python + + model = torchvision.models.alexnet() + batch_size = 256 + flops, macs, params = get_model_profile(model=model, input_shape=(batch_size, 3, 224, 224))) + + Args: + model ([torch.nn.Module]): the PyTorch model to be profiled. + input_shape (tuple): input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument. + args (list): list of positional arguments to the model. + kwargs (dict): dictionary of keyword arguments to the model. + print_profile (bool, optional): whether to print the model profile. Defaults to True. + detailed (bool, optional): whether to print the detailed model profile. Defaults to True. + module_depth (int, optional): the depth into the nested modules. Defaults to -1 (the inner most modules). + top_modules (int, optional): the number of top modules to print in the aggregated profile. Defaults to 3. + warm_up (int, optional): the number of warm-up steps before measuring the latency of each module. Defaults to 1. + as_string (bool, optional): whether to print the output as string. Defaults to True. + output_file (str, optional): path to the output file. If None, the profiler prints to stdout. + ignore_modules ([type], optional): the list of modules to ignore during profiling. Defaults to None. + + Returns: + The number of floating-point operations, multiply-accumulate operations (MACs), and parameters in the model. + """ + assert isinstance(model, nn.Module), "model must be a PyTorch module" + prof = FlopsProfiler(model) + model.eval() + + if input_shape is not None: + assert type(input_shape) is tuple, "input_shape must be a tuple" + assert len(input_shape) >= 1, "input_shape must have at least one element" + try: + input = torch.ones(()).new_empty( + (*input_shape, ), + dtype=next(model.parameters()).dtype, + device=next(model.parameters()).device, + ) + except StopIteration: + input = torch.ones(()).new_empty((*input_shape, )) + + args = [input] + assert (len(args) > 0) or (len(kwargs) > 0), "args and/or kwargs must be specified if input_shape is None" + + logger.info("Flops profiler warming-up...") + for _ in range(warm_up): + if kwargs: + if mode == 'forward': + _ = model(*args, **kwargs) + if mode == 'generate': + _ = model.generate(*args, **kwargs) + else: + if mode == 'forward': + _ = model(*args) + if mode == 'generate': + _ = model.generate(*args) + prof.start_profile(ignore_list=ignore_modules) + + if kwargs: + if mode == 'forward': + _ = model(*args, **kwargs) + if mode == 'generate': + _ = model.generate(*args, **kwargs) + else: + if mode == 'forward': + _ = model(*args) + if mode == 'generate': + _ = model.generate(*args) + + flops = prof.get_total_flops() + macs = prof.get_total_macs() + params = prof.get_total_params() + if print_profile: + prof.print_model_profile(profile_step=warm_up, + module_depth=module_depth, + top_modules=top_modules, + detailed=detailed, + output_file=output_file) + + prof.end_profile() + if as_string: + return number_to_string(flops), macs_to_string(macs), params_to_string(params) + + return flops, macs, params diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..244cce2d968ae43ae2af3d3e3091e59629dac594 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..935099cf58376956c4aa48e43b33c3d4f109c8cb Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f1f4fb58297c10dcb4ac05a20a64295acac0ef8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74e10a2913bb155d71495a6c4e2fa292bbf7a27d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3633baf1b5995e9e6f144c2758d0b83e9f02203e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a15ee659fb9175c6a2b66ea12e8698ce7706d0d0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acd297ba1d81ed63d6bd7f60ebfaf911aae96626 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f095f2cfb62d26de447fb110ea276e72b53ec6fd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa87a925d865a118c2ec2ae6300850c8a2830340 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ba1be5fc39297b0c2a298bb19c346494aa3cc85 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a60d37ed4603ba2d09236d2a5c99fcaf4862224b2404215258104c6ddd878d60 +size 130592 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..752e43150e33a01fd5b233a65c648d353dcb6c58 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c95d76f6e8d82ef1b1f985db8a44345f9ca580a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2cd14efc968707f89e0ce054edbbdf6d5a5a504 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..506a6d77e0e439d0edaf5a16b49a5b52b8279fc8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dd0305714d7a9dfce33f01105993f5cde1d4bd0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da793330a8a8fa49753fb6ba7dd68f26285610b0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8bea0b44c741a7e77a8cfc4c4083ec6bf590428b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c9bdfa072a9d533652f584925dd48ab1a06c40d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14b5b7effa380bd687a9df0679284657fbd661b8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d46e8d5a8323e12a90609aa6d6082024387a143 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16728f1309a65ffd1b0550081df0e1c9f3da3847 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c186426c2d6eb7d2a48d385b2b2addd1bcc0e82f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58d4cccec19e7c640516e23a01894c316e8519df Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..08c4b81937f918425d6da976b9c8391b30eee741 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -0,0 +1,1142 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Use to partition the activations stored for backward propagation +Therefore reduces the memory consumption +Also implements CPU checkpointing and contiguous memory checkpointing +Reduces memory consumption and memory fragmentation + +Code for rng checkpointing taken from NVIDIA Megatron-LM mpu/random.py +b886b7bb972afe72bac0f5de4f42a4a7bae8ebef +""" + +# Parts of the code here are adapted from PyTorch +# repo: https://github.com/pytorch/pytorch +import copy +import torch +import contextlib +from deepspeed import comm as dist +import weakref + +import mmap +from torch import _C + +from deepspeed.runtime.config import DeepSpeedConfig +from deepspeed.utils import logger +from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage +from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers, FORWARD_GLOBAL_TIMER +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime import compiler + +# DeepSpeed Checkpointing Enabled or Disabled +deepspeed_checkpointing_enabled = False + +# MP parameters +mpu = None + +#set default values +mp_rank = 0 +mp_size = 1 +mp_group = None + +# Model Parameters +num_layers = None + +# Checkpointing buffers +contiguous_data_buffers = [] +data_offsets = [] + +contiguous_size_buffers = [] +size_offsets = [] + +timers = None + +# optimization flags +PARTITION_ACTIVATIONS = False +CPU_CHECKPOINT = False +CONTIGUOUS_CHECKPOINTING = False +SYNCHRONIZE = False +PROFILE_TIME = False + +# Default name for the model parallel rng tracker. +_MODEL_PARALLEL_RNG_TRACKER_NAME = 'model-parallel-rng' + + +def detach_variable(inputs, device=None): + if isinstance(inputs, tuple): + out = [] + for inp in inputs: + if not isinstance(inp, torch.Tensor): + out.append(inp) + continue + + requires_grad = inp.requires_grad + + if device is not None: + x = inp.to(device=device) + else: + x = inp + + x = x.detach() + x.requires_grad = requires_grad + out.append(x) + return tuple(out) + else: + raise RuntimeError("Only tuple of tensors is supported. Got Unsupported input type: ", type(inputs).__name__) + + +def _set_cuda_rng_state(new_state, device=-1): + """Sets the random number generator state of the current GPU. + + Arguments: + new_state (torch.ByteTensor): The desired state + This function is adapted from PyTorch repo (torch.cuda.set_rng_state) #ignore-cuda + with a single change: the input state is not cloned. Cloning caused + major performance issues for +4 GPU cases. + """ + if hasattr(_C, '_cuda_setRNGState') and callable(_C._cuda_setRNGState): + # older PyTorch + def cb(): + with get_accelerator().device(device): + _C._cuda_setRNGState(new_state) + else: + # newer PyTorch + if device == -1: + device = torch.device(get_accelerator().device_name()) + elif isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device(get_accelerator().device_name(), device) + + def cb(): + idx = device.index + if idx is None: + idx = get_accelerator().current_device() + default_generator = get_accelerator().default_generator(idx) + default_generator.set_state(new_state) + + get_accelerator().lazy_call(cb) + + +class CudaRNGStatesTracker: + """Tracker for the cuda RNG states. + + Using the `add` method, a cuda rng state is initialized based on + the input `seed` and is assigned to `name`. Later, by forking the + rng state, we can perform operations and return to our starting + cuda state. + """ + + def __init__(self): + # Map from a string name to the cuda rng state. + self.states_ = {} + # Seeds are just for book keeping and ensure no seed is set twice. + self.seeds_ = set() + + def reset(self): + """Set to the initial state (no tracker).""" + self.states_ = {} + self.seeds_ = set() + + def get_states(self): + """Get rng states. Copy the dictionary so we have direct + pointers to the states, not just a pointer to the dictionary.""" + return copy.copy(self.states_) + + def set_states(self, states): + """Set the rng states. For efficiency purposes, we do not check + the size of seed for compatibility.""" + self.states_ = states + + def add(self, name, seed): + """Track the rng state.""" + # Check seed is not already used. + if seed in self.seeds_: + raise Exception('seed {} already exists'.format(seed)) + self.seeds_.add(seed) + # Check that state is not already defined. + if name in self.states_: + raise Exception('cuda rng state {} already exists'.format(name)) + # Get the current rng state. + orig_rng_state = get_accelerator().get_rng_state() + # Set the new state and store it. + get_accelerator().manual_seed(seed) + self.states_[name] = get_accelerator().get_rng_state() + # Reset rng state to what it was. + _set_cuda_rng_state(orig_rng_state) + + @contextlib.contextmanager + def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME): + """Fork the cuda rng state, perform operations, and exit with + the original state.""" + # Check if we have added the state + if name not in self.states_: + raise Exception('cuda rng state {} is not added'.format(name)) + # Store current rng state. + orig_cuda_rng_state = get_accelerator().get_rng_state() + # Set rng state to the desired one + _set_cuda_rng_state(self.states_[name]) + # Do the stuff we wanted to do. + try: + yield + finally: + # Update the current rng state for later use. + self.states_[name] = get_accelerator().get_rng_state() + # And set the state to the original state we started with. + _set_cuda_rng_state(orig_cuda_rng_state) + + +# RNG tracker object. +_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker() + + +def get_cuda_rng_tracker(): + """Get cuda rng tracker.""" + return _CUDA_RNG_STATE_TRACKER + + +def model_parallel_cuda_manual_seed(seed): + """Initialize model parallel cuda seed. + + This function should be called after the model parallel is + initialized. Also, no get_accelerator().manual_seed should be called + after this function. Basically, this is replacement for that + function. + Two set of RNG states are tracked: + default state: This is for data parallelism and is the same among a + set of model parallel GPUs but different across + different model parallel groups. This is used for + example for dropout in the non-model-parallel regions. + model-parallel state: This state is different among a set of model + parallel GPUs, but the same across data parallel + groups. This is used for example for dropout in + model parallel regions. + """ + global mpu + + tp_rank = bwc_tensor_model_parallel_rank(mpu) + + # 2718 is just for fun and any POSITIVE value will work. + offset = seed + 2718 + model_parallel_seed = offset + tp_rank + # Data parallel gets the original seed. + data_parallel_seed = seed + + if dist.get_rank() == 0: + logger.info( + '> initializing model parallel cuda seeds on global rank {}, ' + 'model parallel rank {}, and data parallel rank {} with ' + 'model parallel seed: {} and data parallel seed: {}'.format(dist.get_rank(), tp_rank, + mpu.get_data_parallel_rank(), + model_parallel_seed, data_parallel_seed), ) + _CUDA_RNG_STATE_TRACKER.reset() + # Set the default state. + get_accelerator().manual_seed(data_parallel_seed) + # and model parallel state. + _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, model_parallel_seed) + + +def model_parallel_reconfigure_tp_seed(seed): + global mpu + tp_rank = bwc_tensor_model_parallel_rank(mpu) + model_parallel_seed = seed + 2718 + tp_rank + with _CUDA_RNG_STATE_TRACKER.fork(): + get_accelerator().manual_seed(model_parallel_seed) + + +def get_partition_start(item): + global mp_rank, mp_size, mp_group + size = item.numel() + partition_size = size / mp_size + start = partition_size * mp_rank + return int(start) + + +def get_partition_size(item): + global mp_rank, mp_size, mp_group + size = item.numel() + assert size % mp_size == 0, "Doesn't handle if partition activation if item is not divisible by mp size" + partition_size = size / mp_size + return int(partition_size) + + +def gather_partitioned_activations(tensors, device=None): + global mp_rank, mp_size, mp_group + assert len(tensors) % 2 == 0, f'Expected even count of tensors, instead got {len(tensors)}' + inputs = [] + num_args = int(len(tensors) / 2) + for i in range(num_args): + + item = tensors[2 * i] + size = tensors[2 * i + 1] + + if not is_activation_to_checkpoint(item): + inputs.append(item) + continue + + # don't need to do all_gather if model parallel is not enabled + if mp_group is None or mp_size == 1: + item = item.view(list(size.numpy())) + if device is not None: + item = item.to(device) + inputs.append(item) + continue + + partition_size = item.numel() + tensor_size = partition_size * mp_size + if device is not None: + flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=device) + else: + flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=item.device) + part = flat_tensor.narrow(0, partition_size * mp_rank, partition_size) + part.copy_(item) + dist.all_gather_into_tensor(flat_tensor, part, group=mp_group) + input_tensor = flat_tensor.view(list(size.numpy())) + item.data = input_tensor.data + + inputs.append(item) + + return tuple(inputs) + + +def extract_tensors(all_objects): + """ + Separate objects in list/tuple into tensors and non-tensors and create a mapping to enable re-aggregation. + The order of tensors and non-tensors is preserved in their respective output groups. + + Parameters: + all_objects (list/tuple): Objects containing tensors and non-tensors to be split. + + Returns: + tuple: Containing tensors, non-tensors, and bools of whether each position in original list/tuple was a tensor. + + """ + tensor_objects = [v for v in all_objects if torch.is_tensor(v)] + non_tensor_objects = [v for v in all_objects if not torch.is_tensor(v)] + tensor_flags = [torch.is_tensor(v) for v in all_objects] + if type(all_objects) is tuple: + return tuple(tensor_objects), tuple(non_tensor_objects), tuple(tensor_flags) + return tensor_objects, non_tensor_objects, tensor_flags + + +def merge_tensors(tensor_objects, non_tensor_objects, tensor_flags): + """ + Merge two lists (or tuples) of tensors and non-tensors using a mapping of positions in merged list (or tuple). + + Parameters: + tensor_objects (list/tuple): Tensors to merge. + non_tensor_objects (list/tuple): Non-tensors to merge. + tensor_flags (list/tuple): Indicates whether each position in output is a tensor. + + Returns: + tuple: Merge of tensors and non-tensors + """ + merged_objects = [] + tensor_idx = 0 + non_tensor_idx = 0 + + real_tensor_flags = None + + # remove the flags that are assigned to the size of the flattened tensors + if PARTITION_ACTIVATIONS: + real_tensor_flags = [] + previous_flag = False + for flag in tensor_flags: + if previous_flag: + previous_flag = False + continue + previous_flag = flag + real_tensor_flags.append(flag) + else: + real_tensor_flags = tensor_flags + + for is_tensor in real_tensor_flags: + if is_tensor: + merged_objects.append(tensor_objects[tensor_idx]) + tensor_idx += 1 + else: + merged_objects.append(non_tensor_objects[non_tensor_idx]) + non_tensor_idx += 1 + + return tuple(merged_objects) + + +def is_activation_to_checkpoint(item): + """ + Is an activation to be checkpointed + """ + global mp_size + extra_flag = (not hasattr(item, 'no_checkpointing')) or (hasattr(item, 'no_checkpointing') + and item.no_checkpointing == False) + return torch.is_tensor(item) and item.is_floating_point() and item.numel() >= mp_size and extra_flag + + +def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): + global contiguous_data_buffers, data_offsets + + inputs = [] + num_non_fp_tensors = 0 + + for arg_index, item in enumerate(args): + if not is_activation_to_checkpoint(item): + inputs.append(item) + num_non_fp_tensors += 1 + continue + + i = arg_index - num_non_fp_tensors + partition_size = get_partition_size(item) + partition = item.detach().contiguous().view(-1).narrow(0, get_partition_start(item), partition_size).clone() + + buffer_device = torch.device('cpu') if cpu_checkpoint else partition.device + + if contiguous_checkpoint: + if i >= len(contiguous_data_buffers): + tensor_list = [ + torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) + for _ in range(num_layers) + ] + contiguous_data_buffers.append(tensor_list) + data_offsets.append(0) + elif contiguous_data_buffers[i] is None: + tensor_list = [ + torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) + for _ in range(num_layers) + ] + contiguous_data_buffers[i] = tensor_list + data_offsets[i] = 0 + + # Because the 'new_empty' returns uninitialized pages, + # the pages need to be populated during the cudaMemcpy time + # which increases the data copy time. To avoid this, we + # pre-populate these pages by simply writing 0 ahead of + # the actual cudaMemcpy operation time. Due to the + # previously launched GPU kernels, there is a small + # window of time here for CPUs to populate pages asynchronously. + contiguous_data_buffers[i][data_offsets[i]].data[range( + 0, contiguous_data_buffers[i][data_offsets[i]].data.shape[0], + int(mmap.PAGESIZE / contiguous_data_buffers[i][data_offsets[i]].data.element_size()))] = 0 + + contiguous_partition = contiguous_data_buffers[i][data_offsets[i]].data.copy_(partition.data) + data_offsets[i] = data_offsets[i] + 1 + inputs.append(contiguous_partition) + else: + partition = partition.cpu() if CPU_CHECKPOINT else partition + inputs.append(partition) + + return inputs + + +def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint): + global contiguous_size_buffers, size_offsets + + new_args = [] + num_non_fp_tensors = 0 + + for arg_index, (arg, inp) in enumerate(zip(args, inputs)): + size = torch.tensor(arg.size()) if torch.is_tensor(arg) else None + if not is_activation_to_checkpoint(arg): + new_args.append(arg) + new_args.append(size) + num_non_fp_tensors += 1 + continue + + arg.data = torch.empty([], device=arg.device).data + arg.saved_data = inp.data + + new_args.append(arg) + i = arg_index - num_non_fp_tensors + + if contiguous_checkpoint: + numel = size.numel() + if i >= len(contiguous_size_buffers): + tmp = torch.tensor(()) + contiguous_size_buffers.append( + tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device)) + size_offsets.append(0) + elif contiguous_size_buffers[i] is None: + tmp = torch.tensor(()) + contiguous_size_buffers[i] = tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device) + size_offsets[i] = 0 + + contiguous_size = contiguous_size_buffers[i].narrow(0, size_offsets[i], numel).data.copy_(size.data) + contiguous_size = contiguous_size.view_as(size) + size_offsets[i] = size_offsets[i] + numel + new_args.append(contiguous_size) + else: + new_args.append(size) + + return new_args + + +def get_cpu_activations_for_backward(args, inputs): + new_args = [] + for i, (arg, inp) in enumerate(zip(args, inputs)): + if not is_activation_to_checkpoint(arg): + new_args.append(arg) + continue + + arg.data = torch.empty([], device=arg.device).data + arg.saved_data = inp.data + new_args.append(arg) + + return new_args + + +class CheckpointFunction(torch.autograd.Function): + """This function is adapted from torch.utils.checkpoint with + two main changes: + 1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state` #ignore-cuda + 2) the states in the model parallel tracker are also properly + tracked/set/reset. + 3) Performance activation partitioning, contiguous memory optimization + 4) CPU Checkpointing + 5) Profile forward and backward functions + """ + + @staticmethod + def forward(ctx, run_function, all_outputs, *args): + global mpu, timers, SYNCHRONIZE, PROFILE_TIME + + def save_args_for_backward(*all_args): + tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args) + ctx.deepspeed_saved_tensors = tensor_args + ctx.non_tensor_args = non_tensor_args + ctx.tensor_flags = tensor_flags + + if SYNCHRONIZE: + get_accelerator().synchronize() + + if timers is None and PROFILE_TIME: + timers = Timers() + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).start() + + ctx.run_function = run_function + global num_layers + global mp_rank, mp_size, mp_group + global contiguous_data_buffers, contiguous_size_buffers + global data_offsets, size_offsets + global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset + + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + + if PARTITION_ACTIVATIONS: + inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING) + elif CPU_CHECKPOINT: + inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint) + + # just in case something funky is happening such as reuse of inputs + inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint) + + # Copy the rng states. + ctx.fwd_cpu_rng_state = torch.get_rng_state() + ctx.fwd_cuda_rng_state = get_accelerator().get_rng_state() + ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + see_memory_usage("Before running forward on the layer", force=False) + # ctx.save_for_backward(*args) + with torch.no_grad(): + outputs = run_function(*inputs_cuda) + + see_memory_usage("After running forward on the layer", force=False) + del inputs_cuda + + if PARTITION_ACTIVATIONS: + new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING) + assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}' + save_args_for_backward(*new_args) + elif CPU_CHECKPOINT: + new_args = get_cpu_activations_for_backward(args, inputs) + save_args_for_backward(*new_args) + else: + save_args_for_backward(*args) + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).stop() + timers.log([FORWARD_GLOBAL_TIMER]) + if SYNCHRONIZE: + get_accelerator().synchronize() + + # Tensors returned from forward() may not be differentiable. + if torch.is_tensor(outputs): + non_grad_outputs = [outputs] if not outputs.is_floating_point() else [] + else: + non_grad_outputs = [o for o in outputs if torch.is_tensor(o) and not o.is_floating_point()] + ctx.mark_non_differentiable(*non_grad_outputs) + + if torch.is_tensor(outputs): + all_outputs += [outputs] + return outputs + else: + all_outputs += outputs + outputs, _, _ = extract_tensors(all_objects=outputs) + return tuple(outputs) + + @staticmethod + def backward(ctx, *grads): + global timers + see_memory_usage("In backward", force=False) + # removing pointers to the contiguous buffer memory + # so that they can be garbage collected once the checkpoints + # have been used + if SYNCHRONIZE: + get_accelerator().synchronize() + if PROFILE_TIME: + timers('backward').start() + + if CONTIGUOUS_CHECKPOINTING: + global data_offsets, size_offsets + global contiguous_data_buffers, contiguous_size_buffers + + for buffers in contiguous_data_buffers: + buffers = [] + + # frees up all the pointers to the checkpoints except for the ones + # stored by save for backward + contiguous_data_buffers = [] + contiguous_size_buffers = [] + data_offsets = [] + size_offsets = [] + + see_memory_usage("In backward checkpointing code", force=False) + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError("Checkpointing is not compatible with .grad(), " + "please use .backward() if possible") + + global PARTITION_ACTIVATIONS + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + # Rebuild deepspeed_saved_tensors + for t in ctx.deepspeed_saved_tensors: + if t is not None and hasattr(t, 'saved_data') and t.saved_data is not None: + t.data = t.saved_data.to(t.device) + t.saved_data = None + + if PARTITION_ACTIVATIONS: + # with get_accelerator().stream(transport_stream): + inputs = gather_partitioned_activations(ctx.deepspeed_saved_tensors, + device=cuda_device if CPU_CHECKPOINT else None) + detached_inputs = detach_variable(inputs) + elif CPU_CHECKPOINT: + inputs = move_to_device(ctx.deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint) + detached_inputs = detach_variable(inputs) + else: + inputs = ctx.deepspeed_saved_tensors + detached_inputs = detach_variable(inputs) + + # Add non tensor input args + detached_inputs = merge_tensors(tensor_objects=detached_inputs, + non_tensor_objects=ctx.non_tensor_args, + tensor_flags=ctx.tensor_flags) + + # Store the current states. + bwd_cpu_rng_state = torch.get_rng_state() + bwd_cuda_rng_state = get_accelerator().get_rng_state() + bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + # Set the states to what it used to be before the forward pass. + torch.set_rng_state(ctx.fwd_cpu_rng_state) + _set_cuda_rng_state(ctx.fwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker) + + # if PARTITION_ACTIVATIONS: + # current_stream=get_accelerator().current_stream() + # current_stream.wait_stream(transport_stream) + + see_memory_usage("In backward checkpointing code before forward", force=False) + + with torch.enable_grad(): + outputs = ctx.run_function(*detached_inputs) + + see_memory_usage("In backward checkpointing code after forward", force=False) + # Set the states back to what it was at the start of this function. + torch.set_rng_state(bwd_cpu_rng_state) + _set_cuda_rng_state(bwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker) + + if isinstance(outputs, torch.Tensor): + outputs = (outputs, ) + + # Filter out non tensor outputs + outputs, _, _ = extract_tensors(all_objects=outputs) + + # Construct arguments to autograd.backward(). + # This is usually just outputs and grads, but forward() can return tensors that + # are not differentiable. + output_tensors = [] + grad_tensors = [] + for out, grad in zip(outputs, grads): + if out.requires_grad: + output_tensors.append(out) + grad_tensors.append(grad) + + see_memory_usage("In backward checkpointing code before backward", force=False) + + torch.autograd.backward(output_tensors, grad_tensors) + + # Force clear our stashed tensors to prevent a memory leak in certain scenarios + ctx.deepspeed_saved_tensors = None + ctx.non_tensor_args = None + ctx.tensor_flags = None + + see_memory_usage("After backward checkpointing code after backward", force=False) + + if PROFILE_TIME: + timers('backward').stop() + timers.log(['backward']) + if SYNCHRONIZE: + get_accelerator().synchronize() + ret_list = [None, None] # first None for ctx + for inp in detached_inputs: + if torch.is_tensor(inp): + ret_list.append(inp.grad) + else: + ret_list.append(None) + + return tuple(ret_list) + + +def non_reentrant_checkpoint(function, *args): + """This function is union of `torch.utils.checkpoint._checkpoint_without_reentrant` and `CheckpointFunction` in this module + + This function is aim to solve the back probagation error raised from all input requires no grad. + * has already been implemented in pytorch for a while, the solution is stable at most time except for jit module mode. + * can help to solve the issue which is hacked by `deepspeed.runtime.pipe.module.PipelineModule._is_checkpointable` + + Main modifications compared to the implementation of torch: + 1. adapt to the signature of `checkpoint` function in this module + 2. solve the non-deterministic by random state management consistent with deepspeed `CheckpointFunction` + 3. when there is partition or cpu checkpointing, gather them in the unpack_hook during back probagation + 4. make all after backward blocks in the hook which will executed after all leaf nodes backward execution. + 5. above 4. is inspired by `torch.autograd.graph.register_multi_grad_hook`, which is only implemented after 2.0.0 + """ + global mpu, timers, SYNCHRONIZE, PROFILE_TIME + + deepspeed_saved_tensors = None + non_tensor_args = None + tensor_flags = None + + def save_args_for_backward(*all_args): + """keep this function to reduce the modification from original implementation""" + nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags + tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args) + deepspeed_saved_tensors = tensor_args + non_tensor_args = non_tensor_args + tensor_flags = tensor_flags + + if SYNCHRONIZE: + get_accelerator().synchronize() + + if timers is None and PROFILE_TIME: + timers = Timers() + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).start() + + global num_layers + global mp_rank, mp_size, mp_group + global contiguous_data_buffers, contiguous_size_buffers + global data_offsets, size_offsets + global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset + + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + + if PARTITION_ACTIVATIONS: + inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING) + elif CPU_CHECKPOINT: + inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint) + + # just in case something funky is happening such as reuse of inputs + inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint) + + # Copy the rng states. + fwd_cpu_rng_state = torch.get_rng_state() + fwd_cuda_rng_state = get_accelerator().get_rng_state() + fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + if PARTITION_ACTIVATIONS: + new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING) + assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}' + save_args_for_backward(*new_args) + elif CPU_CHECKPOINT: + new_args = get_cpu_activations_for_backward(args, inputs) + save_args_for_backward(*new_args) + else: + save_args_for_backward(*args) + + class Holder(): + """the place holder object used as activations to save memory""" + pass + + # weakref seems utilized to discover the tensor deletion before a whole + # forward backward pair loop finished + storage: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + weak_holder_list = [] + leaf_tensors = [] + backward_visited_leaf_nodes = 0 + + def checkpoint_pack(tensor_from_forward): + """used to record the activation order in the `weak_holder_list` + + the activation order in holder list is consistent between the first forward and recomputing forward. + * the jit compiled forward will break the order consistency * + """ + res = Holder() + weak_holder_list.append(weakref.ref(res)) + + # if this is a leaf tensor, save it for backward progression trace + # leaf tensor used to be input or parameters, which is not activations and + # has no memory overhead + if tensor_from_forward.requires_grad and tensor_from_forward.is_leaf: + leaf_tensors.append(tensor_from_forward) + return res + + def checkpoint_unpack(holder_from_backward): + """retrieve the activations from recompute""" + nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags + + # if this is the first step of backward probagation, recompute the graph and save + # all the activations with the same order as `checkpoint_pack` does + if len(storage) == 0: + unpack_counter = 0 + + def replay_pack(tensor_from_replay): + """save recompute activations""" + nonlocal unpack_counter + unpack_counter += 1 + + if weak_holder_list[unpack_counter - 1]() is None: + return + + detached_activations = tensor_from_replay.detach() + storage[weak_holder_list[unpack_counter - 1]()] = detached_activations + + return + + def replay_unpack(none_value): + """recompute graph need not to backward""" + raise RuntimeError("You are calling backwards on a tensor that is never exposed.") + + global timers + see_memory_usage("In backward", force=False) + # removing pointers to the contiguous buffer memory + # so that they can be garbage collected once the checkpoints + # have been used + if SYNCHRONIZE: + get_accelerator().synchronize() + if PROFILE_TIME: + timers('backward').start() + + if CONTIGUOUS_CHECKPOINTING: + global data_offsets, size_offsets + global contiguous_data_buffers, contiguous_size_buffers + + for buffers in contiguous_data_buffers: + buffers = [] + + # frees up all the pointers to the checkpoints except for the ones + # stored by save for backward + contiguous_data_buffers = [] + contiguous_size_buffers = [] + data_offsets = [] + size_offsets = [] + + see_memory_usage("In backward checkpointing code", force=False) + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError("Checkpointing is not compatible with .grad(), " + "please use .backward() if possible") + + global PARTITION_ACTIVATIONS + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + + # gather inputs which is partitioned or checkpointed before first forward + if PARTITION_ACTIVATIONS: + # with get_accelerator().stream(transport_stream): + inputs = gather_partitioned_activations(deepspeed_saved_tensors, + device=cuda_device if CPU_CHECKPOINT else None) + detached_inputs = detach_variable(inputs) + elif CPU_CHECKPOINT: + inputs = move_to_device(deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint) + detached_inputs = detach_variable(inputs) + else: + inputs = deepspeed_saved_tensors + detached_inputs = detach_variable(inputs) + + # Add non tensor input args + detached_inputs = merge_tensors(tensor_objects=detached_inputs, + non_tensor_objects=non_tensor_args, + tensor_flags=tensor_flags) + + # Store the current states. + bwd_cpu_rng_state = torch.get_rng_state() + bwd_cuda_rng_state = get_accelerator().get_rng_state() + bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + # Set the states to what it used to be before the forward pass. + torch.set_rng_state(fwd_cpu_rng_state) + _set_cuda_rng_state(fwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(fwd_cuda_rng_state_tracker) + + see_memory_usage("In backward checkpointing code before forward", force=False) + with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(replay_pack, replay_unpack): + _unused = function(*detached_inputs) + + see_memory_usage("In backward checkpointing code after forward", force=False) + # Set the states back to what it was at the start of this function. + torch.set_rng_state(bwd_cpu_rng_state) + _set_cuda_rng_state(bwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker) + + deepspeed_saved_tensors = None + non_tensor_args = None + tensor_flags = None + + if holder_from_backward not in storage: + raise RuntimeError("Attempt to retrieve a tensor saved by autograd multiple times without checkpoint" + " recomputation being triggered in between, this is not currently supported.") + + return storage[holder_from_backward] + + def after_backward_hook(_nonuse_grads): + """the hook registered to all leaf tensors""" + nonlocal leaf_tensors, backward_visited_leaf_nodes + backward_visited_leaf_nodes += 1 + + if backward_visited_leaf_nodes == len(leaf_tensors): + see_memory_usage("After backward checkpointing code after backward", force=False) + + if PROFILE_TIME: + timers('backward').stop() + timers.log(['backward']) + if SYNCHRONIZE: + get_accelerator().synchronize() + + with torch.autograd.graph.saved_tensors_hooks(checkpoint_pack, checkpoint_unpack): + outputs = function(*inputs_cuda) + if PROFILE_TIME or SYNCHRONIZE: + for leaf_tensor in leaf_tensors: + leaf_tensor.register_hook(after_backward_hook) + + see_memory_usage("After running forward on the layer", force=False) + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).stop() + timers.log([FORWARD_GLOBAL_TIMER]) + if SYNCHRONIZE: + get_accelerator().synchronize() + + all_outputs = [] + if torch.is_tensor(outputs): + all_outputs += [outputs] + else: + all_outputs += outputs + + if len(all_outputs) == 1: + return all_outputs[0] + else: + return tuple(all_outputs) + + +@compiler.disable # WA from Pytorch repo for compile + zero 3 accuracy issue +def checkpoint(function, *args): + """Checkpoint a model or part of the model. + This has been directly copied from torch.utils.checkpoint. """ + + all_outputs = [] + CheckpointFunction.apply(function, all_outputs, *args) + if len(all_outputs) == 1: + return all_outputs[0] + else: + return tuple(all_outputs) + + +def partition_activations_in_checkpoint(partition_activation): + global PARTITION_ACTIVATIONS + PARTITION_ACTIVATIONS = partition_activation + if dist.get_rank() == 0: + logger.info(f"**************Partition Activations {PARTITION_ACTIVATIONS}************") + + +def set_num_layers(nlayers): + global num_layers + num_layers = nlayers + + +def reset(): + """Resets memory buffers related to contiguous memory optimizations. + Should be called during eval when multiple forward propagations are + computed without any backward propagation that usually clears these + buffers. + Arguments: + None + + Return: + None + """ + if CONTIGUOUS_CHECKPOINTING: + global data_offsets, size_offsets + global contiguous_data_buffers, contiguous_size_buffers + + for buffers in contiguous_data_buffers: + buffers = [] + + # frees up all the pointers to the checkpoints except for the ones + # stored by save for backward + contiguous_data_buffers = [] + contiguous_size_buffers = [] + data_offsets = [] + size_offsets = [] + + +def _configure_using_config_file(config, mpu=None): + global num_layers, PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ + CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME + + config = DeepSpeedConfig(config, mpu=mpu).activation_checkpointing_config + if dist.get_rank() == 0: + logger.info(config.repr()) + PARTITION_ACTIVATIONS = config.partition_activations + CONTIGUOUS_CHECKPOINTING = config.contiguous_memory_optimization + num_layers = config.number_checkpoints + CPU_CHECKPOINT = config.cpu_checkpointing + SYNCHRONIZE = config.synchronize_checkpoint_boundary + PROFILE_TIME = config.profile + + +def _configure_defaults(): + + global mpu, num_layers, deepspeed_checkpointing_enabled + + global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ + CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME + + PARTITION_ACTIVATIONS = False + CONTIGUOUS_CHECKPOINTING = False + num_layers = False + CPU_CHECKPOINT = False + SYNCHRONIZE = False + PROFILE_TIME = False + deepspeed_checkpointing_enabled = True + + +def configure( + mpu_, + deepspeed_config=None, + partition_activations=None, + contiguous_checkpointing=None, + num_checkpoints=None, + checkpoint_in_cpu=None, + synchronize=None, + profile=None, +): + """Configure DeepSpeed Activation Checkpointing. + + Arguments: + mpu_: Optional: An object that implements the following methods + get_model_parallel_rank/group/world_size, and get_data_parallel_rank/group/world_size + + deepspeed_config: Optional: DeepSpeed Config json file when provided will be used to + configure DeepSpeed Activation Checkpointing + + partition_activations: Optional: Partitions activation checkpoint across model parallel + GPUs when enabled. By default False. Will overwrite deepspeed_config if provided + + contiguous_checkpointing: Optional: Copies activation checkpoints to a contiguous memory + buffer. Works only with homogeneous checkpoints when partition_activations is enabled. + Must provide num_checkpoints. By default False. Will overwrite deepspeed_config if + provided + + num_checkpoints: Optional: Number of activation checkpoints stored during the forward + propagation of the model. Used to calculate the buffer size for contiguous_checkpointing + Will overwrite deepspeed_config if provided + + checkpoint_in_cpu: Optional: Moves the activation checkpoint to CPU. Only works with + partition_activation. Default is false. Will overwrite deepspeed_config if provided + + synchronize: Optional: Performs get_accelerator().synchronize() at the beginning and end of + each call to deepspeed.checkpointing.checkpoint for both forward and backward pass. + By default false. Will overwrite deepspeed_config if provided + + profile: Optional: Logs the forward and backward time for each + deepspeed.checkpointing.checkpoint invocation. Will overwrite deepspeed_config + if provided + + Returns: + None + """ + global mpu, num_layers, deepspeed_checkpointing_enabled + + global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ + CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME + + _configure_defaults() + + if mpu_ is not None: + mpu = mpu_ + + if deepspeed_config is not None: + _configure_using_config_file(deepspeed_config, mpu=mpu) + + if partition_activations is not None: + PARTITION_ACTIVATIONS = partition_activations + + if contiguous_checkpointing is not None: + CONTIGUOUS_CHECKPOINTING = contiguous_checkpointing + + if num_checkpoints is not None: + num_layers = num_checkpoints + + if checkpoint_in_cpu is not None: + CPU_CHECKPOINT = checkpoint_in_cpu + + if synchronize is not None: + SYNCHRONIZE = synchronize + + if profile is not None: + PROFILE_TIME = profile + + if CONTIGUOUS_CHECKPOINTING: + assert PARTITION_ACTIVATIONS, "Contiguous Checkpointing is only available with partitioned activations. Set partitioned activations to true in deepspeed config" + if CONTIGUOUS_CHECKPOINTING: + assert num_layers is not None, "Must specify the number of layers with contiguous memory checkpointing" + + global mp_rank, mp_size, mp_group + + if mpu is not None: + if hasattr(mpu, 'get_tensor_model_parallel_rank'): + mp_rank = mpu.get_tensor_model_parallel_rank() + mp_size = mpu.get_tensor_model_parallel_world_size() + mp_group = mpu.get_tensor_model_parallel_group() + else: + mp_rank = mpu.get_model_parallel_rank() + mp_size = mpu.get_model_parallel_world_size() + mp_group = mpu.get_model_parallel_group() + + #print configuration only once + see_memory_usage("After configuration", force=False) + if dist.get_rank() == 0: + logger.info(f"Activation Checkpointing Information") + logger.info(f"----Partition Activations {PARTITION_ACTIVATIONS}, CPU CHECKPOINTING {CPU_CHECKPOINT}") + logger.info(f"----contiguous Memory Checkpointing {CONTIGUOUS_CHECKPOINTING} with {num_layers} total layers") + logger.info(f"----Synchronization {SYNCHRONIZE}") + logger.info(f"----Profiling time in checkpointing {PROFILE_TIME}") + + +def is_configured(): + """True if deepspeed activation checkpointing has been configured + by calling deepspeed.checkpointing.configure, else returns false + + Arguments: + None + + Return: + True of configured, else False + """ + return deepspeed_checkpointing_enabled diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/config.py new file mode 100644 index 0000000000000000000000000000000000000000..dc07388a95da039b50bad87b4aa57b12f4e41f6f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/config.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject + +######################################### +# DeepSpeed Activation Checkpointing +######################################### +# Activation Checkpointing Allows to save memory by only keeping a select few +#activations for the backpropagation. +ACTIVATION_CHKPT_FORMAT = ''' +Activation Checkpointing should be configured as: +"session_params": { + "activation_checkpointing": { + "partitioned_activations": [true|false], + "number_checkpoints": 100, + "contiguous_memory_optimization": [true|false], + "cpu_checkpointing": [true|false], + "profile": [true|false], + "synchronize_checkpoint_boundary": [true|false], + } +} +''' + +ACT_CHKPT_PARTITION_ACTIVATIONS = 'partition_activations' +ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT = False + +ACT_CHKPT_NUMBER_CHECKPOINTS = 'number_checkpoints' +ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT = None + +ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION = 'contiguous_memory_optimization' +ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT = False + +ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY = 'synchronize_checkpoint_boundary' +ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT = False + +ACT_CHKPT_PROFILE = 'profile' +ACT_CHKPT_PROFILE_DEFAULT = False + +ACT_CHKPT_CPU_CHECKPOINTING = 'cpu_checkpointing' +ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT = False + +ACT_CHKPT = 'activation_checkpointing' + +ACT_CHKPT_DEFAULT = { + ACT_CHKPT_PARTITION_ACTIVATIONS: ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT, + ACT_CHKPT_NUMBER_CHECKPOINTS: ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT, + ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION: ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT, + ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY: ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT, + ACT_CHKPT_PROFILE: ACT_CHKPT_PROFILE_DEFAULT, + ACT_CHKPT_CPU_CHECKPOINTING: ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT +} + + +class DeepSpeedActivationCheckpointingConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedActivationCheckpointingConfig, self).__init__() + + self.partition_activations = None + self.contiguous_memory_optimization = None + self.cpu_checkpointing = None + self.number_checkpoints = None + self.synchronize_checkpoint_boundary = None + self.profile = None + + if ACT_CHKPT in param_dict.keys(): + act_chkpt_config_dict = param_dict[ACT_CHKPT] + else: + act_chkpt_config_dict = ACT_CHKPT_DEFAULT + + self._initialize(act_chkpt_config_dict) + + def _initialize(self, act_chkpt_config_dict): + self.partition_activations = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_PARTITION_ACTIVATIONS, + ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT) + + self.contiguous_memory_optimization = get_scalar_param(act_chkpt_config_dict, + ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION, + ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT) + + self.cpu_checkpointing = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_CPU_CHECKPOINTING, + ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT) + + self.number_checkpoints = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_NUMBER_CHECKPOINTS, + ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT) + + self.profile = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_PROFILE, ACT_CHKPT_PROFILE_DEFAULT) + + self.synchronize_checkpoint_boundary = get_scalar_param(act_chkpt_config_dict, + ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY, + ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/base_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/base_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..726124027131c4edf110ef0c080e07abcf5f19ff --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/base_optimizer.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch + +from deepspeed.utils import logger +from deepspeed.utils.tensor_fragment import map_to_flat_opt_states +from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, see_memory_usage +from deepspeed.runtime.torch_autocast import get_autocast_dtype, is_autocast_initialized + + +class DeepSpeedOptimizer(object): + pass + + +class ZeROOptimizer(DeepSpeedOptimizer): + + def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, checkpoint_dir: str) -> None: + checkpoint_dir = os.path.join(checkpoint_dir, "zero") + optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") + assert os.path.isfile( + optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' + optim_sd = torch.load(optim_state_path, weights_only=False) + + self._load_global_state(optim_sd) + + tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) + if self.mpu is None: + logger.warning("MPU is not provided, setting tp size to 1 in checkpoint loading.") + tp_world_size = 1 + else: + tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \ + else self.mpu.get_tensor_model_parallel_world_size() + + for i, (param_group, + loaded_param_group) in enumerate(zip(self.optimizer.param_groups, optim_sd['param_groups'])): + # We have an assumption that all params in the same param_group have the same keys + opt_keys = set() + steps = [] + + lp_groups = getattr(self, lp_groups_name) + for lp in lp_groups[i]: + if lp._hp_mapping is not None: + #print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}") + step = lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank, + tp_world_size) + for key in lp._hp_mapping.get_optim_state_keys(): + opt_keys.add(key) + steps.append(step) + + hp_param = param_group['params'][0] + assert all(step == steps[0] for step in steps), f"Steps {steps} are not equal" + if steps[0] is not None: + self.optimizer.state[hp_param]['step'] = steps[0] + + map_to_flat_opt_states(hp_param, lp_groups[i], self.optimizer.state, opt_keys) + + for key, value in loaded_param_group.items(): + if key == 'params': + continue + param_group[key] = value + + def report_ipg_memory_usage(self, tag, param_elems, dtype=None): + dtypes = self.ipg_buckets.keys() if dtype is None else [dtype] + + for dt in dtypes: + bucket = self.ipg_buckets[dt] + elem_count = bucket.elements + param_elems + percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size + see_memory_usage( + f"{tag}: elems in_bucket {dt} {bucket.elements} param {param_elems} max_percent {percent_of_bucket_size}" + ) + + def get_param_comm_dtype(self, param): + if is_autocast_initialized(): + return get_autocast_dtype(param) + else: + return self.communication_data_type diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..cd00e9e43b00d03c7b66c445afa68de6700579a1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py @@ -0,0 +1,566 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from collections import OrderedDict +import torch +import sys +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed import comm as dist +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.runtime.base_optimizer import ZeROOptimizer +from packaging import version as pkg_version +from deepspeed.git_version_info import version +from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim, + align_dense_tensors, all_gather_dp_groups, is_model_parallel_parameter, + see_memory_usage, graph_process, get_norm_with_moe_layers) +from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state, fragment_address, groups +from deepspeed.moe.utils import is_moe_param, is_moe_param_group +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank +from deepspeed.utils.torch import register_grad_hook +from deepspeed.checkpoint import enable_universal_checkpoint +from deepspeed.checkpoint.constants import (DS_VERSION, PARTITION_COUNT, BASE_OPTIMIZER_STATE, + SINGLE_PARTITION_OF_FP32_GROUPS, CLIP_GRAD, GROUP_PADDINGS, + PARAM_SLICE_MAPPINGS) + +setattr(sys.modules[__name__], 'fragment_address', fragment_address) + + +def print_rank_0(message, debug=False, force=False): + if dist.get_rank() == 0 and (debug or force): + print(message) + + +class BF16_Optimizer(ZeROOptimizer): + + def __init__(self, + init_optimizer, + param_names, + bfloat16_config, + mpu=None, + clip_grad=0.0, + norm_type=2, + allgather_bucket_size=5000000000, + dp_process_group=None, + timers=None, + grad_acc_dtype=None, + graph_harvesting=False, + has_moe_layers=False): + super().__init__() + see_memory_usage('begin bf16_optimizer', force=True) + self.timers = timers + self.optimizer = init_optimizer + self.param_names = param_names + self.using_real_optimizer = not isinstance(self.optimizer, DummyOptim) + + assert bfloat16_config.enabled, f"BF16Optimizer: requires bfloat16 to be enabled" + assert grad_acc_dtype in [torch.float32, torch.bfloat16 + ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}" + self.grad_acc_dtype = grad_acc_dtype + + self.immediate_grad_update = bfloat16_config.immediate_grad_update + + self.clip_grad = clip_grad + self.norm_type = norm_type + self.mpu = mpu + self.allgather_bucket_size = int(allgather_bucket_size) + self.dp_process_group = dp_process_group + self.dp_rank = dist.get_rank(group=self.dp_process_group) + self.has_moe_layers = has_moe_layers + self.non_expert_gradients = [] + self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] + if self.has_moe_layers: + self._configure_moe_settings() + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + + #align nccl all-gather send buffers to 4-bye boundary + self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 + + # Build BF16/FP32 groups + self.bf16_groups = [] + self.bf16_groups_flat = [] + self.bf16_partitioned_groups = [] + + self.fp32_groups_flat_partition = [] + + # Maintain different fp32 gradients views for convenience + self.fp32_groups_gradients = [] + self.fp32_groups_gradient_dict = {} + self.fp32_groups_gradients_flat = [] + self.fp32_groups_actual_gradients_flat = [] + self.fp32_groups_gradient_flat_partition = [] + self.fp32_groups_has_gradients = [] + + self.group_paddings = [] + self.graph_harvesting = graph_harvesting + if self.using_real_optimizer: + self._setup_for_real_optimizer() + + see_memory_usage('end bf16_ optimizer', force=True) + + def destroy(self): + for i, _ in enumerate(self.optimizer.param_groups): + for p in self.bf16_groups[i]: + if getattr(p, '_hp_mapping', None): + p._hp_mapping = None + for hook in self._grad_acc_hooks: + hook.remove() + print_rank_0("Removed grad acc hooks") + + def _configure_moe_settings(self): + assert any( + [is_moe_param_group(group) for group in self.optimizer.param_groups] + ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" + + for i, group in enumerate(self.optimizer.param_groups): + if is_moe_param_group(group): + assert all([is_moe_param(param) + for param in group['params']]), "All params in MoE group must be MoE params" + self.real_dp_process_group[i] = groups._get_expert_data_parallel_group(group['name']) + self.expert_gradients = {} + if self.has_moe_layers: + for key in groups._get_expert_data_parallel_group_dict().keys(): + self.expert_gradients[key] = [] + + def _setup_for_real_optimizer(self): + self.partition_count = [dist.get_world_size(group=pg) for pg in self.real_dp_process_group] + + for i, param_group in enumerate(self.optimizer.param_groups): + real_dp_world_size = dist.get_world_size(group=self.real_dp_process_group[i]) + see_memory_usage(f'before initializing group {i}', force=True) + + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + + # grab the original list + trainable_parameters = [param for param in param_group['params'] if param.requires_grad] + self.bf16_groups.append(trainable_parameters) + + # create flat bf16 params + self.bf16_groups_flat.append( + self._flatten_dense_tensors_aligned(self.bf16_groups[i], + self.nccl_start_alignment_factor * real_dp_world_size)) + # Make bf16 params point to flat tensor storage + self._update_storage_to_flattened_tensor(tensor_list=self.bf16_groups[i], + flat_tensor=self.bf16_groups_flat[i]) + + # divide flat weights into equal sized partitions + partition_size = self.bf16_groups_flat[i].numel() // real_dp_world_size + bf16_dp_partitions = [ + self.bf16_groups_flat[i].narrow(0, dp_index * partition_size, partition_size) + for dp_index in range(real_dp_world_size) + ] + self.bf16_partitioned_groups.append(bf16_dp_partitions) + + # create fp32 params partition + self.fp32_groups_flat_partition.append(bf16_dp_partitions[partition_id].clone().float().detach()) + self.fp32_groups_flat_partition[i].requires_grad = True + + num_elem_list = [t.numel() for t in self.bf16_groups[i]] + + # create fp32 gradients + fp32_flat_buffer = torch.zeros_like(self.bf16_groups_flat[i], dtype=self.grad_acc_dtype) + self.fp32_groups_gradients_flat.append(fp32_flat_buffer) + if self.has_moe_layers and is_moe_param_group(param_group): + self.expert_gradients[param_group['name']].append(fp32_flat_buffer) + else: + self.non_expert_gradients.append(fp32_flat_buffer) + + # track individual fp32 gradients for entire model + fp32_gradients = self._split_flat_tensor(flat_tensor=self.fp32_groups_gradients_flat[i], + num_elem_list=num_elem_list) + self.fp32_groups_gradients.append(fp32_gradients) + self.fp32_groups_gradient_dict[i] = fp32_gradients + + # flat tensor corresponding to actual fp32 gradients (i.e., minus alignment padding) + length_without_padding = sum(num_elem_list) + self.fp32_groups_actual_gradients_flat.append( + torch.narrow(self.fp32_groups_gradients_flat[i], 0, 0, length_without_padding)) + + # flat tensor corresponding to gradient partition + self.fp32_groups_gradient_flat_partition.append( + torch.narrow(self.fp32_groups_gradients_flat[i], 0, partition_id * partition_size, partition_size)) + + # track fp32 gradient updates + self.fp32_groups_has_gradients.append([False] * len(self.bf16_groups[i])) + + # Record padding required for alignment + if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: + padding = self.bf16_groups_flat[i].numel() - length_without_padding + else: + padding = 0 + + self.group_paddings.append(padding) + + # update optimizer param groups to reference fp32 params partition + param_group['params'] = [self.fp32_groups_flat_partition[i]] + + see_memory_usage(f'after initializing group {i}', force=True) + + self._grad_acc_hooks = [] + if self.immediate_grad_update: + self.create_grad_acc_hooks() + + # Need optimizer states initialized before linking lp to optimizer state + self._link_all_hp_params() + self._hp_optimizer_states_linked = False + self._enable_universal_checkpoint() + self._param_slice_mappings = self._create_param_mapping() + + def _enable_universal_checkpoint(self): + for lp_param_group in self.bf16_groups: + enable_universal_checkpoint(param_list=lp_param_group) + + def _create_param_mapping(self): + param_mapping = [] + for i, _ in enumerate(self.optimizer.param_groups): + param_mapping_per_group = OrderedDict() + for lp in self.bf16_groups[i]: + if lp._hp_mapping is not None: + lp_name = self.param_names[lp] + param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() + param_mapping.append(param_mapping_per_group) + + return param_mapping + + def _link_all_hp_params(self): + for i, _ in enumerate(self.optimizer.param_groups): + real_dp_world_size = dist.get_world_size(group=self.real_dp_process_group[i]) + + # Link bf16 and fp32 params in partition + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + partition_size = self.bf16_groups_flat[i].numel() // real_dp_world_size + flat_hp_partition = self.fp32_groups_flat_partition[i] + link_hp_params(lp_param_list=self.bf16_groups[i], + flat_hp_partition=flat_hp_partition, + gradient_dict=self.fp32_groups_gradient_dict, + offload_gradient_dict=None, + use_offload=False, + param_group_index=i, + partition_start=partition_id * partition_size, + partition_size=partition_size, + dp_group=self.real_dp_process_group[i]) + + def _lazy_init_hp_params_optimizer_state(self): + if not self._hp_optimizer_states_linked: + for i, _ in enumerate(self.optimizer.param_groups): + lazy_init_hp_params_optimizer_state(self.bf16_groups[i], self.fp32_groups_flat_partition[i], + self.optimizer.state) + self._hp_optimizer_states_linked = True + + def _split_flat_tensor(self, flat_tensor, num_elem_list): + assert sum(num_elem_list) <= flat_tensor.numel() + tensor_list = [] + offset = 0 + for num_elem in num_elem_list: + dense_tensor = torch.narrow(flat_tensor, 0, offset, num_elem) + tensor_list.append(dense_tensor) + offset += num_elem + + return tensor_list + + def _update_storage_to_flattened_tensor(self, tensor_list, flat_tensor): + updated_params = self.unflatten(flat_tensor, tensor_list) + for p, q in zip(tensor_list, updated_params): + p.data = q.data + + def _flatten_dense_tensors_aligned(self, tensor_list, alignment): + return self.flatten(align_dense_tensors(tensor_list, alignment)) + + @torch.no_grad() + def step(self, closure=None): + if closure is not None: + raise NotImplementedError(f'{self.__class__} does not support closure.') + + non_expert_grads_for_norm, expert_grads_for_norm = self.get_grads_for_norm() + non_expert_groups_norm = get_global_norm_of_tensors(input_tensors=non_expert_grads_for_norm, + mpu=self.mpu, + norm_type=self.norm_type, + use_graph=self.graph_harvesting) + all_groups_norm = non_expert_groups_norm + if self.has_moe_layers: + all_groups_norm = get_norm_with_moe_layers(non_expert_groups_norm, + mpu=self.mpu, + expert_tensors=expert_grads_for_norm, + norm_type=self.norm_type) + + self._global_grad_norm = all_groups_norm + + assert all_groups_norm > 0. + if self.clip_grad > 0.: + clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(for_clipping=True), + max_norm=self.clip_grad, + global_norm=all_groups_norm, + mpu=self.mpu, + use_graph=self.graph_harvesting) + + for param_partition, grad_partition in zip(self.fp32_groups_flat_partition, + self.fp32_groups_gradient_flat_partition): + # In case of grad acc dtype different than FP32, need to cast to high precision. + param_partition.grad = grad_partition.to( + param_partition.dtype) if grad_partition.dtype != param_partition.dtype else grad_partition + + self.optimizer.step() + + if self.grad_acc_dtype is not torch.float32: + for param_partition in self.fp32_groups_flat_partition: + param_partition.grad = None + + # We need to link optimizer state after the first step() call + self._lazy_init_hp_params_optimizer_state() + + self.update_lp_params() + + self.clear_hp_grads() + + def backward(self, loss, retain_graph=False, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs): + """Perform a backward pass and copy the low-precision gradients to the + high-precision copy. + + We copy/accumulate to the high-precision grads now to prevent accumulating in the + bf16 grads after successive backward() calls (i.e., grad accumulation steps > 1) + + The low-precision grads are deallocated during this procedure. + """ + self.clear_lp_grads() + loss.backward(retain_graph=retain_graph, **bwd_kwargs) + + if update_hp_grads: + self.update_hp_grads(clear_lp_grads=clear_lp_grads) + + @torch.no_grad() + def _update_hp_grad(self, lp, group_idx, param_idx, clear_lp_grads): + if lp.grad is None: + return + + hp_grad = self.fp32_groups_gradients[group_idx][param_idx] + assert hp_grad is not None, \ + f'high precision param has no gradient, lp param_id = {id(lp)} group_info = [{group_idx}][{param_idx}]' + + hp_grad.data.add_(lp.grad.data.to(hp_grad.dtype).view(hp_grad.shape)) + lp._hp_grad = hp_grad + self.fp32_groups_has_gradients[group_idx][param_idx] = True + + # clear gradients + if clear_lp_grads: + lp.grad.zero_() + + @torch.no_grad() + def _update_hp_grads_func(self, clear_lp_grads=False): + for i, group in enumerate(self.bf16_groups): + for j, lp in enumerate(group): + self._update_hp_grad(lp, i, j, clear_lp_grads) + + @torch.no_grad() + def update_hp_grads(self, clear_lp_grads=False): + if self.immediate_grad_update: + return + + if self.graph_harvesting: + graph_process(False, self._update_hp_grads_func, clear_lp_grads) + else: + self._update_hp_grads_func(clear_lp_grads) + #cpu op + for i, group in enumerate(self.bf16_groups): + for j, lp in enumerate(group): + if lp.grad is None: + continue + self.fp32_groups_has_gradients[i][j] = True + + @torch.no_grad() + def get_grads_for_reduction(self): + if self.has_moe_layers: + return self.non_expert_gradients, self.expert_gradients + return self.non_expert_gradients, {} + + @torch.no_grad() + def get_grads_for_norm(self, for_clipping=False): + """ + Returns: + tuple[list[Tensor], dict[ep_name, List[Tensor]] | list: + If for_clipping, return all gradients. + Otherwise, separate and return dict of expert_grad and list of non_expert_grad + """ + # (grads, expert_group_name) + expert_grads_for_norm = {} + + # grads + non_expert_grads_for_norm = [] + all_grads_for_clip = [] + + tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) + assert len(self.bf16_groups) == len(self.optimizer.param_groups) + for i, group in enumerate(self.bf16_groups): + for j, lp in enumerate(group): + if not for_clipping: + if hasattr(lp, PIPE_REPLICATED) and lp.ds_pipe_replicated: + continue + + # skip duplicated parameters. perform norm only on cards with tp_rank=0. + # non-duplicated parameters include: + # - Parameters with tp: Use allreducesum of mp_group. + # - Moe Parameters with ep: Use allreducesum of ep_group. + if not (tensor_mp_rank == 0 or is_model_parallel_parameter(lp) or is_moe_param(lp)): + continue + + if not self.fp32_groups_has_gradients[i][j]: + continue + if not for_clipping: + param_group = self.optimizer.param_groups[i] + if self.has_moe_layers and is_moe_param_group(param_group): + if param_group['name'] not in expert_grads_for_norm: + expert_grads_for_norm[param_group['name']] = [] + expert_grads_for_norm[param_group['name']].append(self.fp32_groups_gradients[i][j]) + else: + non_expert_grads_for_norm.append(self.fp32_groups_gradients[i][j]) + else: + all_grads_for_clip.append(self.fp32_groups_gradients[i][j]) + if not for_clipping: + return non_expert_grads_for_norm, expert_grads_for_norm + return all_grads_for_clip + + @torch.no_grad() + def update_lp_params(self): + for i, (bf16_partitions, + fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + bf16_partitions[partition_id].data.copy_(fp32_partition.data) + + all_gather_dp_groups(groups_flat=self.bf16_groups_flat, + partitioned_param_groups=self.bf16_partitioned_groups, + dp_process_group=self.real_dp_process_group, + start_alignment_factor=self.nccl_start_alignment_factor, + allgather_bucket_size=self.allgather_bucket_size) + + def clear_hp_grads(self): + for flat_gradients in self.fp32_groups_gradients_flat: + flat_gradients.zero_() + + for i, group in enumerate(self.fp32_groups_gradients): + self.fp32_groups_has_gradients[i] = [False] * len(group) + + def clear_lp_grads(self, set_to_none=False): + + # using zero_() fixed memory address for graph replay + if self.graph_harvesting: + assert not set_to_none, "graph harvesting is incompatible with setting lp grads to None" + + zero_grads_list = [] + for group in self.bf16_groups: + for param in group: + if set_to_none: + param.grad = None + elif param.grad is not None: + if param.grad.grad_fn is not None: + param.grad.detach_() + zero_grads_list.append(param.grad) + if not set_to_none and len(zero_grads_list) > 0: + torch._foreach_zero_(zero_grads_list) + + def zero_grad(self, set_to_none=True): + self.clear_lp_grads(set_to_none) + self.clear_hp_grads() + + def state_dict(self): + state_dict = {} + state_dict[CLIP_GRAD] = self.clip_grad + state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() + state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = self.fp32_groups_flat_partition + state_dict[GROUP_PADDINGS] = self.group_paddings + state_dict[PARTITION_COUNT] = self.partition_count + state_dict[DS_VERSION] = version + state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings + + return state_dict + + # Restore base optimizer fp32 weights bfloat16 weights + def _restore_from_bit16_weights(self): + for i, (bf16_partitions, + fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + fp32_partition.data.copy_(bf16_partitions[partition_id].data) + + def refresh_fp32_params(self): + self._restore_from_bit16_weights() + + def load_state_dict(self, + state_dict_list, + checkpoint_folder=None, + load_optimizer_states=True, + load_from_fp32_weights=False, + load_serial=None, + param_shapes=None): + if checkpoint_folder: + self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) + else: + self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) + + def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): + + dp_rank = dist.get_rank(group=self.dp_process_group) + current_rank_sd = state_dict_list[dp_rank] + + ckpt_version = current_rank_sd.get(DS_VERSION, False) + assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" + ckpt_version = pkg_version.parse(ckpt_version) + + self.clip_grad = current_rank_sd.get(CLIP_GRAD, self.clip_grad) + + if load_optimizer_states: + print(f"_load_legacy_checkpoint current_rank_sd[BASE_OPTIMIZER_STATE]") + self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) + + if load_from_fp32_weights: + for current, saved in zip(self.fp32_groups_flat_partition, + current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): + src_tensor = _get_padded_tensor(saved, current.numel()) + current.data.copy_(src_tensor.data) + + if load_optimizer_states: + self._link_all_hp_params() + + def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): + self.load_hp_checkpoint_state_from_checkpoint_dir("bf16_groups", checkpoint_folder) + + def _load_global_state(self, sd): + pass + + @property + def param_groups(self): + """Forward the wrapped optimizer's parameters.""" + return self.optimizer.param_groups + + @property + def state(self): + """Forward the wrapped optimizer's states.""" + return self.optimizer.state + + def accumulate_hp_grads_and_remove_lp(self, lp_param, group_idx, param_idx): + assert self.immediate_grad_update + self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=False) + + def create_grad_acc_hooks(self): + for i, param_group in enumerate(self.bf16_groups): + for j, param in enumerate(param_group): + if param.requires_grad: + + def wrapper(param, i, j): + + def accumulate_hp_grads_and_remove_lp(*notneeded): + self.accumulate_hp_grads_and_remove_lp(param, i, j) + + self._grad_acc_hooks.append(register_grad_hook(param, accumulate_hp_grads_and_remove_lp)) + + wrapper(param, i, j) + + +def _get_padded_tensor(src_tensor, size): + if src_tensor.numel() >= size: + return src_tensor + padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) + slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) + slice_tensor.data.copy_(src_tensor.data) + return padded_tensor diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a6484db58e0608eafd7ec33ba871c5982c939aed --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .fast_checkpoint_engine import FastCheckpointEngine +from .torch_checkpoint_engine import TorchCheckpointEngine +from .decoupled_checkpoint_engine import DecoupledCheckpointEngine +from .checkpoint_engine import CheckpointCommitInfo +from .utils import create_checkpoint_engine diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f7092ee2fc910edeb981d2effee7953a47f46ce Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbf77ba59e77a173d0c048ebf1aaca8122f75969 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41cc9c69b73767b55b29f4601d09a487a8ba5df6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..231f221a9ab25f40780795c4735fcd4e3f5e9ed7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42adbf4b7cea1423d2485ed9b047d73845037ee6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..786b24e04ee7abda9d7147d5b6ce8f8878779389 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff792fd1f460138608352c27dbacb8456de25b0b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..03254b935a2ac688ab0790e2ac0bc5ee88398855 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +import abc +from abc import ABC + +from dataclasses import dataclass + + +@dataclass +class CheckpointCommitInfo(object): + tag: str + save_dir: str + save_latest: bool + + +class CheckpointEngine(ABC): + # init checkpoint engine for save/load + def __init__(self, config_params=None): + self.name = None + + @abc.abstractmethod + def create(self, info: CheckpointCommitInfo): + # create checkpoint on give tag for save/load. + ... + + @abc.abstractmethod + def save(self, state_dict, path: str): + ... + + def makedirs(self, path, exist_ok=False): + os.makedirs(path, exist_ok=exist_ok) + + @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 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c0ea88836e376deb54a72f2157c88f907fc52f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py @@ -0,0 +1,160 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.multiprocessing as mp +from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ + CheckpointEngine, CheckpointCommitInfo +from deepspeed.runtime.checkpoint_engine.fast_checkpoint_engine import FastCheckpointEngine +from deepspeed import comm as dist +from deepspeed.runtime.utils import get_checkpoint_folder_size + +from enum import Enum + + +class DecoupledEvent(Enum): + SAVE_EVENT = 1 + COMMIT_EVENT = 2 + EXIT_EVENT = 3 + + +class CheckpointSize(object): + + def __init__(self): + self._pre = None + self._post = None + self._gigabytes = None + + def gb_size(self): + return self._gigabytes + + def set_pre_size(self, size): + self._pre = size + + def set_post_size(self, size): + self._post = size + self._gigabytes = (self._post - self._pre) / (1024**3) + + +def init_decoupled_checkpoint(config_params, dp_writer_config, save_event, save_queue, optimize_dp_state): + checkpoint_engine = FastCheckpointEngine(config_params, dp_writer_config, optimize_dp_state) + print(f'Created FastCheckpointEngine for Decoupled Checkpointing') + save_path_list = [] + while True: + (save_info, event_type) = save_queue.get() + if event_type == DecoupledEvent.SAVE_EVENT and save_info is not None: + state_dict, save_path = save_info + # print(f'Received decoupled checkpoint request for {save_path=}') + save_path_list.append(save_path) + checkpoint_engine.save(state_dict, save_path) + del state_dict + # print(f'Completed decoupled checkpoint request for {save_path=}') + + if event_type == DecoupledEvent.COMMIT_EVENT: + # print(f'Recieved commit request for {save_path_list=}') + save_path_list = [] + save_event.set() + + if event_type == DecoupledEvent.EXIT_EVENT: + # print(f'Received decoupled exit request') + break + + +ENGINE_NAME = "DecoupledCheckpointEngine" + + +class DecoupledCheckpointEngine(CheckpointEngine): + + def __init__(self, config_params, dp_writer_config, optimize_dp_state): + if mp.get_start_methods(allow_None=False) is None: + mp.set_start_method('spawn') + super().__init__(config_params) + self.name = ENGINE_NAME + self.dp_writer_config = dp_writer_config + self.commit_info = None + self.checkpoint_size = CheckpointSize() + self.global_rank = dist.get_rank() + self.optimize_dp_state = optimize_dp_state + if dp_writer_config is None: + self.save_event = None + self.save_queue = None + self.ckpt_process = None + self.local_rank = None + print( + f'[{ENGINE_NAME}]: No checkpoint process self.global_rank={self.global_rank} self.dp_writer_config={self.dp_writer_config}' + ) + else: + self.save_event = mp.Event() + self.save_queue = mp.SimpleQueue() + engine_args = (config_params, dp_writer_config, self.save_event, self.save_queue, self.optimize_dp_state) + self.ckpt_process = mp.Process(target=init_decoupled_checkpoint, args=engine_args) + self.ckpt_process.start() + self.local_rank = dp_writer_config.local_rank + print( + f'[{ENGINE_NAME}]: Create checkpoint process self.global_rank={self.global_rank} self.ckpt_process.pid={self.ckpt_process.pid} self.dp_writer_config={self.dp_writer_config}' + ) + + def __del__(self): + self.cleanup() + + def create(self, info: CheckpointCommitInfo): + self.commit_info = info + if self.checkpoint_size.gb_size() is None: + pre_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank) + self.checkpoint_size.set_pre_size(pre_size) + + def load(self, path: str, map_location=None): + sd = torch.load(path, map_location=map_location) + return sd + + def save(self, state_dict, path: str): + if self.ckpt_process is None: + return + save_info = (state_dict, path) + self.save_queue.put((save_info, DecoupledEvent.SAVE_EVENT)) + + def commit(self, info: CheckpointCommitInfo): + assert info == self.commit_info + if self.ckpt_process is not None: + self.save_queue.put((None, DecoupledEvent.COMMIT_EVENT)) + # print(f'[begin] wait for decoupled complete for {info.tag}') + self.save_event.wait() + # print(f'[end] wait for decoupled complete for {info.tag}') + self.save_event.clear() + self.commit_info = None + + if self.checkpoint_size.gb_size() is None: + dist.barrier() + post_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank) + self.checkpoint_size.set_post_size(post_size) + + if self.global_rank == 0: + print( + f'{self.name} self.global_rank={self.global_rank} created checkpoint of {round(self.checkpoint_size.gb_size(), 2)} GB' + ) + + return True + + def get_commit_info(self): + # print(f'getting commit info {self.commit_info=}') + return self.commit_info + + def is_decoupled(self): + return True + + def cleanup(self): + # print(f'Inside {self.name} cleanup') + + if self.get_commit_info() is not None: + self.commit(self.commit_info) + + if self.ckpt_process is not None: + self.save_queue.put((None, DecoupledEvent.EXIT_EVENT)) + self.ckpt_process.join() + self.ckpt_process = None + self.save_queue = None + + def is_data_parallel_writer(self, dp_rank): + return self.ckpt_process is not None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..4bfecf810dc7984fe1aa412763df915cd56fb8ca --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py @@ -0,0 +1,50 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ + CheckpointEngine, CheckpointCommitInfo +from deepspeed.runtime.model_checkpointing import ( + CHECKPOINT_WRITER, + CHECKPOINT_SERIALIZATION, + CheckpointWriterFactory, +) + + +class FastCheckpointEngine(CheckpointEngine): + + def __init__(self, config_params, dp_writer_config, optimize_dp_state): + super().__init__(config_params) + self.name = 'FastCheckpointEngine' + self.serialization_enabled = config_params.checkpoint_config[CHECKPOINT_SERIALIZATION] + self.optimize_dp_state = optimize_dp_state + if dp_writer_config is None: + self._writer = None + else: + self._writer = CheckpointWriterFactory(writer_config=config_params.checkpoint_config[CHECKPOINT_WRITER], + aio_config=config_params.aio_config, + dp_writer_config=dp_writer_config) + + def create(self, info: CheckpointCommitInfo): + pass + + def save(self, state_dict, path: str): + if self._writer is None: + return + + torch.save(obj=state_dict, + f=self._writer.create_writer(path, self.optimize_dp_state), + _use_new_zipfile_serialization=self.serialization_enabled) + self._writer.release_writer() + + def load(self, path: str, map_location=None): + sd = torch.load(path, map_location=map_location) + return sd + + def commit(self, info: CheckpointCommitInfo): + return True + + def is_data_parallel_writer(self, dp_rank): + return self._writer is not None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..485e21411fb8332a167458770de31d71680e678d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py @@ -0,0 +1,108 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +import torch_nebula + +from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ + CheckpointEngine, CheckpointCommitInfo +from deepspeed.utils import logger, log_dist +from deepspeed.nebula.constants import * + + +def _get_tag_from_path(path): + return os.path.basename(os.path.dirname(path)) + + +class NebulaCheckpointEngine(CheckpointEngine): + + def __init__(self, config_params=None): + super().__init__(config_params) + self.name = "NebulaCheckpointEngine" + self.checkpoint = None + self.tag_flag = None + self.enable_nebula_load = config_params.enable_nebula_load + self.nebula_load_path = config_params.load_path + if self.nebula_load_path is None: + self.nebula_load_path = config_params.persistent_storage_path + + nebula_config_params = { + NEBULA_PERSISTENT_STORAGE_PATH: config_params.persistent_storage_path, + NEBULA_PERSISTENT_TIME_INTERVAL: config_params.persistent_time_interval, + NEBULA_NUM_OF_VERSION_IN_RETENTION: config_params.num_of_version_in_retention, + } + torch_nebula.init(**nebula_config_params) + + def create(self, info: CheckpointCommitInfo): + log_dist(f"[Nebula] Start Checkpoint for tag:{info.tag}", ranks=[0]) + # -2 means: customer needs to explicitly tell nebula + # current checkpoint is complete by commit methond. + self.checkpoint = torch_nebula.Checkpoint(info.tag, -2) + + def save(self, state_dict, path: str): + log_dist(f"[Nebula] Create dummy files for loading.") + torch.save("", path) + + tag = _get_tag_from_path(path) + partititon_name = os.path.basename(path) + logger.info(f"[Nebula] Saving {partititon_name} under tag {tag}...") + self.checkpoint.save(partititon_name, state_dict) + logger.info(f"[Nebula] Saved {partititon_name} under tag {tag}.") + + def load(self, path: str, map_location=None): + tag = _get_tag_from_path(path) + first_load_flag = self.tag_flag is None or self.tag_flag == tag + if not self.enable_nebula_load and first_load_flag: + self.tag_flag = tag + logger.info(f"[Nebula] Disable nebula load. Loading checkpoint from {path} ...") + partition = torch.load(path, map_location=map_location, weights_only=False) + logger.info(f"[Nebula] Disable nebula load. Loaded checkpoint from {path} .") + return partition + + partition_name = os.path.basename(path) + logger.info(f"[Nebula] Loading {path} under tag {tag} from nebula path {self.nebula_load_path}...") + + checkpoint = None + if tag in (None, 'latest', 'latest_universal'): + # In some cases, there is the inconsistent tag between deepspeed metadata (latest file) + # and nebula metadata, will lead to the failure on loading with deepspeed tag. Then we + # will try to load the valid latest checkpoint from nebula(tier3 > tier1). So, in summary + # when met failure loading for given tag, the loading priority would be like: + # nebula tier3 latest > nebula tier1 latest. + checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path) + else: + checkpoint = torch_nebula.get_checkpoint(tag=tag, persist_path=self.nebula_load_path) + + if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''): + logger.info( + f"Unable to find valid checkpoint tag:{tag} from Nebula, try to get latest checkpoint again from nebula {self.nebula_load_path} path!" + ) + # nebula tier3 latest + checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path) + if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''): + logger.info( + f"Unable to find latest checkpoint from Nebula tier3, try to get latest checkpoint again from nebula tier1 path!" + ) + # nebula tier1 latest + checkpoint = torch_nebula.get_latest_checkpoint() + logger.warning(f"Unable to find valid checkpoint from Nebula under tag:{tag}.") + return None + + tag = checkpoint.tag + self.tag_flag = -1 + partition = checkpoint.load(partition_name, map_location=map_location) + logger.info(f"[Nebula] Loaded {path} under tag {tag} from {self.nebula_load_path}.") + return partition + + def commit(self, info: CheckpointCommitInfo): + tag = info.tag + # nebula commit will be call when all files under give tag are ready to be persisted in the async way. + logger.info(f"[Nebula] all files for {tag} are saved in tier1. It is ready to start persisting") + commit_rls = self.checkpoint.commit() + if not commit_rls: + logger.error(f"[Nebula] failed to commit the checkpoint, please check the log.") + return False + return commit_rls diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..9014677f1d5c5a50f4df5fc011a11489924b1e04 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils import log_dist +from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ + CheckpointEngine, CheckpointCommitInfo +from deepspeed.runtime.model_checkpointing import CHECKPOINT_SERIALIZATION + +ENGINE_NAME = "TorchCheckpointEngine" + + +class TorchCheckpointEngine(CheckpointEngine): + + def __init__(self, config_params=None): + super().__init__(config_params) + self.name = ENGINE_NAME + if config_params is None: + self.zipfile_serialization = False + else: + self.zipfile_serialization = config_params.checkpoint_config[CHECKPOINT_SERIALIZATION] + log_dist(f'[{ENGINE_NAME}] Initialized with serialization = {self.zipfile_serialization}', ranks=[0]) + + def create(self, info: CheckpointCommitInfo): + log_dist(f"[Torch] Checkpoint {info.tag} is begin to save!", ranks=[0]) + pass + + def save(self, state_dict, path: str): + # log_dist(f"[Torch] Saving [begin] {path}... {self.zipfile_serialization=}", ranks=[0]) + torch.save(state_dict, path, _use_new_zipfile_serialization=self.zipfile_serialization) + # log_dist(f"[Torch] Saving [end] {path}... {self.zipfile_serialization=}", ranks=[0]) + + def load(self, path: str, map_location=None): + log_dist(f"[Torch] Begin Load checkpoint from {path}...", ranks=[0]) + partition = torch.load(path, map_location=map_location, weights_only=False) + log_dist(f"[Torch] End Load checkpoint from {path}...", ranks=[0]) + return partition + + def commit(self, info: CheckpointCommitInfo): + #logger.info(f"[Torch] Checkpoint {tag} is ready now!") + return True diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8fa5e9725092338aee41b6951688c1cdad11b2e2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/utils.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.model_checkpointing.constants import * +from deepspeed.runtime.model_checkpointing.utils import create_data_parallel_writer_config +from deepspeed.utils import logger + +from .decoupled_checkpoint_engine import DecoupledCheckpointEngine +from .fast_checkpoint_engine import FastCheckpointEngine +from .torch_checkpoint_engine import TorchCheckpointEngine + + +def create_checkpoint_engine(config_params, groups, zero_stage, has_moe_layers, optimize_dp_state): + if config_params is not None: + if config_params.checkpoint_config[CHECKPOINT_WRITER] is not None: + writer_config = config_params.checkpoint_config[CHECKPOINT_WRITER] + dp_writer_config = create_data_parallel_writer_config( + groups=groups, + parallel_unit=writer_config[CHECKPOINT_DATA_PARALLEL], + zero_stage=zero_stage, + has_moe_layers=has_moe_layers) + if writer_config[CHECKPOINT_WRITER_DECOUPLED]: + return DecoupledCheckpointEngine(config_params, dp_writer_config, optimize_dp_state) + else: + return FastCheckpointEngine(config_params, dp_writer_config, optimize_dp_state) + + if config_params is not None and config_params.nebula_config.enabled: + try: + from .nebula_checkpoint_engine import NebulaCheckpointEngine + except ImportError as err: + logger.error(f"No torch_nebula was found! Will fall back to torch.save. Details: {err}") + return TorchCheckpointEngine(config_params) + else: + return NebulaCheckpointEngine(config_params=config_params.nebula_config) + + return TorchCheckpointEngine(config_params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64e70db95100a0e60ca0eccdde878fa49a673ff5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56b53e7618b3cfd994037be7ef05147124312ed5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54bd84ed04398619556000e2f9d47e742c10f085 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..427c567885a1b6b85127d06140a44209d7b02d9f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0136194c61e3fb4a841750ab42c41bae2c92d803 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d52e786fb516fd33de388a22abb1ff5fe7196f4a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/coalesced_collectives.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/coalesced_collectives.py new file mode 100644 index 0000000000000000000000000000000000000000..2fadce52222cda680600253ae840f84b89bda7ed --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/coalesced_collectives.py @@ -0,0 +1,218 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +batched collective operations for overhead amortization and better +bandwidth utilization +""" + +import math +from typing import List, Any +import torch +from torch import Tensor +from deepspeed import comm as dist +from deepspeed.comm import ProcessGroup, all_to_all_single +from deepspeed.accelerator import get_accelerator +from deepspeed.utils import instrument_w_nvtx +from deepspeed.ops import op_builder +from deepspeed.utils import logger + + +def _torch_reduce_scatter_fn(input_tensor: Tensor, output_tensor: Tensor, group=None, async_op=False, prof=False): + return instrument_w_nvtx(dist.reduce_scatter_fn)(output_tensor, input_tensor, group=group, async_op=False) + + +quantizer_module = None + + +@instrument_w_nvtx +@torch.no_grad() +def all_to_all_quant_reduce(tensors: List[Tensor], groups: {}) -> List[Tensor]: + global quantizer_module + if quantizer_module is None: + quantizer_module = op_builder.QuantizerBuilder().load() + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + this_rank = dist.get_rank() + intra_idx = int(this_rank / local_world_size) + inter_idx = this_rank % local_world_size + output_lst: List[Tensor] = [None] * len(tensors) + for idx, tensor in enumerate(tensors): + if tensor.dim() == 1: + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + elif tensor.numel() % (2 * global_world_size) != 0: + # Due to the constraint of 2-stage all-to-all, the input tensor must be divisible by 2 * global_world_size + # Otherwise, all-to-all cannot be performed because of shape mismatch. + # See more at https://github.com/deepspeedai/DeepSpeed/pull/5056 + logger.warning( + f"qgZ falls back to reduce_scatter because tensor size = {tensor.numel()} is not divisible by (2 * global_world_size) = {2 * global_world_size}. Please consider allocating a new world to enable qgZ" + ) + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + else: + intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size) + + inter_quant_group = intra_quant_group // local_world_size + intra_quant_int4, intra_q_scales = quantizer_module.swizzle_quant(tensor, intra_quant_group, 4, + quantizer_module.Symmetric, 1, num_nodes, + local_world_size) + local_output = torch.empty_like(intra_quant_int4) + scale_output = torch.empty_like(intra_q_scales) + all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']) + all_to_all_single(scale_output, intra_q_scales, group=groups[f'local_{intra_idx}']) + global_input_tensor, global_scales = quantizer_module.quantized_reduction( + local_output, scale_output, intra_quant_group, inter_quant_group, 4, quantizer_module.Symmetric, + local_world_size) + global_output = torch.empty_like(global_input_tensor) + global_scale_output = torch.empty_like(global_scales) + all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']) + all_to_all_single(global_scale_output, global_scales, group=groups[f'global_{inter_idx}']) + final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(), + 4, quantizer_module.Symmetric) + assert final_output.numel( + ) % num_nodes == 0, f"final_output.numel()={final_output.numel()} is not divisible by num_nodes={num_nodes}" + output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1) + return output_lst + + +@instrument_w_nvtx +@torch.no_grad() +def all_to_all_loco_quant_reduce( + params: List[Tensor], + groups: {}, + loco_param: Any = None, +) -> List[Tensor]: + global quantizer_module + global loco_idx + if quantizer_module is None: + quantizer_module = op_builder.QuantizerBuilder().load() + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + this_rank = dist.get_rank() + intra_idx = int(this_rank / local_world_size) + inter_idx = this_rank % local_world_size + output_lst: List[Tensor] = [None] * len(params) + for idx, p in enumerate(params): + tensor = p.grad + if tensor.dim() == 1: + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + elif tensor.numel() % (2 * global_world_size) != 0: + # Due to the constraint of 2-stage all-to-all, the input tensor must be divisible by 2 * global_world_size + # Otherwise, all-to-all cannot be performed because of shape mismatch. + # See more at https://github.com/deepspeedai/DeepSpeed/pull/5056 + logger.warning( + f"qgZ falls back to reduce_scatter because tensor size = {tensor.numel()} is not divisible by (2 * global_world_size) = {2 * global_world_size}. Please consider allocating a new world to enable qgZ" + ) + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + else: + err_beta = loco_param['err_beta'] + reset_T = loco_param['reset_T'] + if not hasattr(p, 'intra_ef_buf') or loco_idx > reset_T: + loco_idx = 0 + intra_err = torch.zeros_like(p.grad) + inter_err = torch.zeros(tensor.numel() // local_world_size, device=tensor.device, dtype=tensor.dtype) + else: + intra_err = quantizer_module.dequantize(p.intra_ef_buf[0], p.intra_ef_buf[1], + p.intra_ef_buf[1].numel(), 8, quantizer_module.Symmetric) + inter_err = quantizer_module.dequantize(p.inter_ef_buf[0], p.inter_ef_buf[1], + p.inter_ef_buf[1].numel(), 8, quantizer_module.Symmetric) + + intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size) + inter_quant_group = intra_quant_group // local_world_size + intra_quant_int4, intra_q_scales = quantizer_module.loco_swizzle_quant(tensor, intra_err, err_beta, + intra_quant_group, 4, + quantizer_module.Symmetric, 1, + num_nodes, local_world_size) + local_output = torch.empty_like(intra_quant_int4) + scale_output = torch.empty_like(intra_q_scales) + all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']) + all_to_all_single(scale_output, intra_q_scales, group=groups[f'local_{intra_idx}']) + + p.intra_ef_buf = quantizer_module.quantize(intra_err, intra_quant_group, 8, quantizer_module.Symmetric) + + global_input_tensor, global_scales = quantizer_module.loco_quantized_reduction( + local_output, scale_output, inter_err, err_beta, intra_quant_group, inter_quant_group, 4, + quantizer_module.Symmetric, local_world_size) + + global_output = torch.empty_like(global_input_tensor) + global_scale_output = torch.empty_like(global_scales) + all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']) + all_to_all_single(global_scale_output, global_scales, group=groups[f'global_{inter_idx}']) + + p.inter_ef_buf = quantizer_module.quantize(inter_err, inter_quant_group, 8, quantizer_module.Symmetric) + + final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(), + 4, quantizer_module.Symmetric) + assert final_output.numel( + ) % num_nodes == 0, f"final_output.numel()={final_output.numel()} is not divisible by num_nodes={num_nodes}" + output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1) + loco_idx += 1 + + return output_lst + + +@instrument_w_nvtx +@torch.no_grad() +def reduce_scatter_coalesced( + tensors: List[Tensor], + group: ProcessGroup = None, +) -> List[Tensor]: + """simultaneously reduce-scatter a list of tensors - this can be done more + efficiently than individual reduce scatter calls + TODO. see if PyTorch team wants a c++ version of this for ProcessGroupNCCL + """ + this_rank = dist.get_rank(group) + world_sz = dist.get_world_size(group) + + partition_lst_for_each_tensor = [None] * len(tensors) + for tensor_idx, tensor in enumerate(tensors): + flattened_tensor = tensor.view(-1) + chunk_sz = math.ceil(tensor.numel() / world_sz) + partition_lst_for_each_tensor[tensor_idx] = [ + flattened_tensor[rank * chunk_sz:rank * chunk_sz + chunk_sz] for rank in range(0, world_sz) + ] + + padded_partition_sz_for_each_tensor = tuple(math.ceil(t.numel() / world_sz) for t in tensors) + + if len(tensors) == 1 and tensors[0].numel() % world_sz == 0: + # if there's only one tensor being reduced and we don't need to pad + # we have an opportunity to avoid a memory allocation + tensor_partition_flat_buffer = tensors[0].view(-1) + else: + # interleave tensor partitions such that the correct reduced partitions of each tensor + # end up at each rank + tensor_partitions_lst_with_padding = [] + for rank in range(world_sz): + for tensor_idx in range(len(tensors)): + # add tensor content + tensor_chunk = partition_lst_for_each_tensor[tensor_idx][rank] + tensor_partitions_lst_with_padding.append(tensor_chunk) + + # add padding if necessary + padding_sz = padded_partition_sz_for_each_tensor[tensor_idx] - tensor_chunk.numel() + if padding_sz > 0: + tensor_partitions_lst_with_padding.append( + torch.empty(padding_sz, dtype=tensor_chunk.dtype, device=tensor_chunk.device)) + + tensor_partition_flat_buffer = instrument_w_nvtx(torch.cat)(tensor_partitions_lst_with_padding) + + tensor_partition_flat_buffer.div_(world_sz) # pre-divide + tensor_partition_buffer_for_each_rank: List[Tensor] = torch.chunk(tensor_partition_flat_buffer, world_sz) + + # batched reduce-scatter call + _torch_reduce_scatter_fn(tensor_partition_flat_buffer, + tensor_partition_buffer_for_each_rank[this_rank], + group=group) + + # reverse procedure of the interleaving done previously, done on the + # result of the batched reduce-scatter + output_lst: List[Tensor] = [None] * len(tensors) + offset = 0 + for tensor_idx in range(len(tensors)): + output_lst[tensor_idx] = tensor_partition_buffer_for_each_rank[this_rank].narrow( + 0, offset, partition_lst_for_each_tensor[tensor_idx][this_rank].numel()) + + offset += padded_partition_sz_for_each_tensor[tensor_idx] + return output_lst diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/compressed.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/compressed.py new file mode 100644 index 0000000000000000000000000000000000000000..2c5482eb1ad769bb3a0ecd87b4a3489d31ce28db --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/compressed.py @@ -0,0 +1,137 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +import torch +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import PackbitsBuilder + + +class CompressedBackend(object): + + def __init__(self, mpu=None): + if mpu is None: + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + else: + self.mpu = mpu + self.world_group = self.mpu.get_data_parallel_group() + self.size = dist.get_world_size(group=self.world_group) + self.rank = dist.get_rank(group=self.world_group) + self.packer = PackbitsBuilder().load() + + def my_igather(self, rank, size, group, sendbuf, recvbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) + else: + recvbuf[rank] = sendbuf + else: + req.append(dist.isend(sendbuf, group=group, dst=root)) + return req + + def my_gather(self, rank, size, group, sendbuf, recvbuf, root): + if rank == root: + for idx in range(size): + if idx != rank: + dist.recv(recvbuf[idx], src=idx, group=group) + else: + recvbuf[rank] = sendbuf + else: + dist.send(sendbuf, group=group, dst=root) + + def pack(self, buffer, size): + # pack float tensor into uint8 tensor + packed = self.packer.packbits(buffer.float(), buffer.numel(), self.rank) + return packed.reshape(size, -1) + + def unpack(self, buffer, size, dtype): + # unpack uint8 to float tensor + unpacked = self.packer.unpackbits(buffer, buffer.numel(), self.rank) + return unpacked.reshape(size, -1).to(dtype) + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + + # align size of original_buffer and error + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) + + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + sign_list_packed_tmp = self.pack(buffer_m, self.size).type(torch.int8) + + recvbuf_sign = torch.zeros([self.size, len(sign_list_packed_tmp[self.rank])], + dtype=sign_list_packed_tmp[0].dtype, + device=sign_list_packed_tmp.device) + + sign_list_packed = [sign_list_packed_tmp[idx] for idx in range(self.size)] + + recvbuf_scale = [ + torch.zeros(1, dtype=worker_scale.dtype, device=get_accelerator().current_device_name()) + for _ in range(self.size) + ] + + # communication phase 1 + # all to all for sign + dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) + # all gather for scale + dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) + + flattened_recvbuf_sign = recvbuf_sign.type(torch.uint8).flatten() + compensated_server_m = self.unpack(flattened_recvbuf_sign, self.size, torch.float32) \ + .mul_(torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) + + compensated_server_m.add_(server_error) + + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + server_sign_packed = self.pack(compensated_server_m, 1).type(torch.int8) + + # recvbuf_sign_server + recvbuf_sign_server_tmp = torch.zeros([self.size, len(server_sign_packed[0])], + dtype=recvbuf_sign.dtype, + device=server_sign_packed.device) + + recvbuf_sign_server = [recvbuf_sign_server_tmp[idx] for idx in range(self.size)] + + # recvbuf_scale_server + recvbuf_scale_server_tmp = torch.zeros([self.size, 1], + dtype=worker_scale.dtype, + device=server_sign_packed.device) + + recvbuf_scale_server = [recvbuf_scale_server_tmp[idx] for idx in range(self.size)] + + # communication Phase 2 + dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) + dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) + + recvbuf_sign_server = torch.stack(recvbuf_sign_server) + + flattened_recvbuf_sign_server = recvbuf_sign_server.type(torch.uint8).flatten() + + buffer_m.data.copy_( + self.unpack(flattened_recvbuf_sign_server, self.size, + torch.float32).mul_(recvbuf_scale_server_tmp).flatten().data) + + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/hccl.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/hccl.py new file mode 100644 index 0000000000000000000000000000000000000000..b8639c7da4c99327e4a6afe4e226d6c4a224886b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/hccl.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +import torch +import torch_npu +import deepspeed.comm as dist + + +class HcclBackend(object): + + def __init__(self, mpu=None): + if mpu is None: + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + else: + self.mpu = mpu + self.world_group = self.mpu.get_data_parallel_group() + self.size = dist.get_world_size(group=self.world_group) + self.rank = dist.get_rank(group=self.world_group) + + def my_igather(self, rank, size, group, sendbuf, recvbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) + else: + recvbuf[rank] = sendbuf + else: + req.append(dist.isend(sendbuf, group=group, dst=root)) + return req + + def my_gather(self, rank, size, group, sendbuf, recvbuf, root): + if rank == root: + for idx in range(size): + if idx != rank: + dist.recv(recvbuf[idx], src=idx, group=group) + else: + recvbuf[rank] = sendbuf + else: + dist.send(sendbuf, group=group, dst=root) + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + + # align size of original_buffer and error + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) + + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + sign_list_packed_tmp = torch_npu.npu_sign_bits_pack(buffer_m, self.size).type(torch.int8) + + recvbuf_sign = torch.zeros([self.size, len(sign_list_packed_tmp[self.rank])], + dtype=sign_list_packed_tmp[0].dtype, + device=sign_list_packed_tmp.device) + + sign_list_packed = [sign_list_packed_tmp[idx] for idx in range(self.size)] + + recvbuf_scale = [ + torch.zeros(1, dtype=worker_scale.dtype, device=torch.device(local_rank)) for _ in range(self.size) + ] + + # communication phase 1 + # all to all for sign + dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) + # all gather for scale + dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) + + flattened_recvbuf_sign = recvbuf_sign.type(torch.uint8).flatten() + compensated_server_m = torch_npu.npu_sign_bits_unpack(flattened_recvbuf_sign, self.size, torch.float32) \ + .mul_(torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) + + compensated_server_m.add_(server_error) + + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + server_sign_packed = torch_npu.npu_sign_bits_pack(compensated_server_m, 1).type(torch.int8) + + # recvbuf_sign_server + recvbuf_sign_server_tmp = torch.zeros([self.size, len(server_sign_packed[0])], + dtype=recvbuf_sign.dtype, + device=server_sign_packed.device) + + recvbuf_sign_server = [recvbuf_sign_server_tmp[idx] for idx in range(self.size)] + + # recvbuf_scale_server + recvbuf_scale_server_tmp = torch.zeros([self.size, 1], + dtype=worker_scale.dtype, + device=server_sign_packed.device) + + recvbuf_scale_server = [recvbuf_scale_server_tmp[idx] for idx in range(self.size)] + + # communication Phase 2 + dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) + dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) + + recvbuf_sign_server = torch.stack(recvbuf_sign_server) + + flattened_recvbuf_sign_server = recvbuf_sign_server.type(torch.uint8).flatten() + + buffer_m.data.copy_( + torch_npu.npu_sign_bits_unpack(flattened_recvbuf_sign_server, self.size, + torch.float32).mul_(recvbuf_scale_server_tmp).flatten().data) + + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/mpi.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/mpi.py new file mode 100644 index 0000000000000000000000000000000000000000..bc544787aa7a7de5181d83897e725991f572e5ce --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/mpi.py @@ -0,0 +1,215 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import cupy +import time +import numpy as np +from mpi4py import MPI + +from deepspeed.runtime.compression.cupy import CupyBackend + + +class MpiBackend(object): + + def __init__(self, cuda_aware): + self.comm = MPI.COMM_WORLD + self.rank = self.comm.Get_rank() + self.size = self.comm.Get_size() + self.cuda_aware = cuda_aware + self.compression_backend = CupyBackend() + + def my_igather(self, rank, size, comm, sendbuf, recbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(comm.Irecv(recbuf[idx], source=idx)) + else: + recbuf[rank] = sendbuf + else: + req.append(comm.Isend(sendbuf, dest=root)) + return req + + def gather_cuda(self, rank, world_size, comm, cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, + cupy_recvbuf_scale): + # We do in-place operations on cupy buffers so we do not return any buffers + requests = [] + for idx in range(world_size): + req_sign = self.my_igather(rank, world_size, comm, cupy_sign_list_packed[idx], cupy_recvbuf_sign, root=idx) + requests += req_sign + + for idx in range(world_size): + req_scale = self.my_igather(rank, world_size, comm, cupy_worker_scale, cupy_recvbuf_scale, root=idx) + requests += req_scale + + MPI.Request.Waitall(requests) + + def gather_host(self, rank, world_size, comm, cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, + cupy_recvbuf_scale): + + # In-place operations are not possible for newly created cupy arrays + # so we need to return the new buffers + numpy_recvbuf_sign = np.zeros([world_size, cupy_sign_list_packed[rank].size], + dtype=cupy_sign_list_packed[0].dtype) + numpy_recvbuf_scale = np.zeros([world_size, 1], dtype=cupy_worker_scale.dtype) + + # 1. convert from cupy to numpy + numpy_sign_list_packed = cupy_sign_list_packed + + for idx in range(world_size): + numpy_sign_list_packed[idx] = cupy.asnumpy(cupy_sign_list_packed[idx]) + + numpy_worker_scale = cupy.asnumpy(cupy_worker_scale) + numpy_recvbuf_scale = cupy.asnumpy(cupy_recvbuf_scale) + + cupy.cuda.get_current_stream().synchronize() + + # 2. use numpy buffers for communication + requests = [] + + for idx in range(world_size): + req_sign = self.my_igather(rank, + world_size, + comm, + numpy_sign_list_packed[idx], + numpy_recvbuf_sign, + root=idx) + requests += req_sign + + for idx in range(world_size): + req_scale = self.my_igather(rank, world_size, comm, numpy_worker_scale, numpy_recvbuf_scale, root=idx) + requests += req_scale + + MPI.Request.Waitall(requests) + + # 3. Convert back from numpy to cupy + cupy_recvbuf_sign = cupy.asarray(numpy_recvbuf_sign) + for idx in range(world_size): + cupy_sign_list_packed[idx] = cupy.asarray(numpy_sign_list_packed[idx]) + + cupy_worker_scale = cupy.asarray(numpy_worker_scale) + cupy_recvbuf_scale = cupy.asarray(numpy_recvbuf_scale) + cupy.cuda.get_current_stream().synchronize() + + return cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, cupy_recvbuf_scale + + def allgather_cuda(self, comm, cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server): + comm.Allgather(cupy_server_sign_packed, cupy_recvbuf_sign_server) + comm.Allgather(cupy_server_scale, cupy_recvbuf_scale_server) + + def allgather_host(self, comm, cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server): + + # 1. Convert cupy to numpy + numpy_recvbuf_sign_server = np.zeros([comm.Get_size(), cupy_server_sign_packed.size], + dtype=cupy_server_sign_packed.dtype) + numpy_recvbuf_scale_server = np.zeros([comm.Get_size(), 1], dtype=cupy_server_scale.dtype) + + numpy_server_sign_packed = cupy.asnumpy(cupy_server_sign_packed) + numpy_recvbuf_sign_server = cupy.asnumpy(cupy_recvbuf_sign_server) + numpy_server_scale = cupy.asnumpy(cupy_server_scale) + numpy_recvbuf_scale_server = cupy.asnumpy(cupy_recvbuf_scale_server) + cupy.cuda.get_current_stream().synchronize() + + # 2. Communicate numpy buffers + comm.Allgather(numpy_server_sign_packed, numpy_recvbuf_sign_server) + comm.Allgather(numpy_server_scale, numpy_recvbuf_scale_server) + comm.Barrier() + + # 3. Convert numpy back to cupy + cupy_server_sign_packed = cupy.asarray(numpy_server_sign_packed) + cupy_recvbuf_sign_server = cupy.asarray(numpy_recvbuf_sign_server) + cupy_server_scale = cupy.asarray(numpy_server_scale) + cupy_recvbuf_scale_server = cupy.asarray(numpy_recvbuf_scale_server) + cupy.cuda.get_current_stream().synchronize() + + return cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, cupy_recvbuf_scale_server + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + + all_start_time = time.time() + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + cupy.cuda.Device(local_rank).use() + + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + cupy_sign_list_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool()), self.size) + cupy_worker_scale = self.compression_backend.torch2cupy(worker_scale) + + cupy_recvbuf_sign = cupy.zeros([self.size, cupy_sign_list_packed[self.rank].size], + dtype=cupy_sign_list_packed[0].dtype) + cupy_recvbuf_scale = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) + + # Communication Phase 1 + gather_start = time.time() + if self.cuda_aware: + self.gather_cuda(self.rank, self.size, self.comm, cupy_sign_list_packed, cupy_recvbuf_sign, + cupy_worker_scale, cupy_recvbuf_scale) + else: + _, cupy_recvbuf_sign, _, cupy_recvbuf_scale = self.gather_host(self.rank, self.size, self.comm, + cupy_sign_list_packed, cupy_recvbuf_sign, + cupy_worker_scale, cupy_recvbuf_scale) + gather_end = time.time() + + # cupy_sign_list_packed, cupy_worker_scale, worker_scale = None, None, None + cupy_sign_list_packed = None + + compensated_server_m = self.compression_backend.cupy2torch( + (cupy.unpackbits(cupy_recvbuf_sign.flatten())).reshape(self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + self.compression_backend.cupy2torch(cupy_recvbuf_scale).mul_(1 / self.size)).sum(0) + compensated_server_m.add_(server_error) + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + cupy_server_scale = self.compression_backend.torch2cupy(server_scale) + + cupy_server_sign_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool()), 1) + compensated_server_m = None + + cupy_recvbuf_sign_server = cupy.zeros([self.size, cupy_server_sign_packed[0].size], + dtype=cupy_recvbuf_sign.dtype) + cupy_recvbuf_scale_server = cupy.zeros([self.size, 1], dtype=cupy_recvbuf_scale.dtype) + # cupy_recvbuf_sign, cupy_recvbuf_scale = None, None + cupy_recvbuf_sign = None + + # Communication Phase 2 + if self.cuda_aware: + self.allgather_cuda(self.comm, cupy_server_sign_packed[0], cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server) + else: + _, cupy_recvbuf_sign_server, _, cupy_recvbuf_scale_server = self.allgather_host( + self.comm, cupy_server_sign_packed[0], cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server) + + # cupy_server_sign_packed, cupy_server_scale, server_scale = None, None, None + cupy_server_sign_packed = None + + buffer_m.data.copy_( + self.compression_backend.cupy2torch((cupy.unpackbits(cupy_recvbuf_sign_server.flatten())).reshape( + self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + self.compression_backend.cupy2torch(cupy_recvbuf_scale_server)).flatten().data) + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + # cupy_recvbuf_sign_server, cupy_recvbuf_scale_server = None, None + + return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/nccl.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/nccl.py new file mode 100644 index 0000000000000000000000000000000000000000..a57b7519a295e584e77561843c52fa025bfaa66d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/nccl.py @@ -0,0 +1,166 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +import cupy +import numpy as np + +from deepspeed.runtime.compression.cupy import CupyBackend +from deepspeed.utils.torch import required_torch_version +from deepspeed.accelerator import get_accelerator + + +class NcclBackend(object): + + def __init__(self, mpu=None): + if mpu is None: + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + else: + self.mpu = mpu + self.world_group = self.mpu.get_data_parallel_group() + self.rank = dist.get_rank(group=self.world_group) + self.size = dist.get_world_size(group=self.world_group) + self.compression_backend = CupyBackend() + self.bool_not_supported = required_torch_version(min_version=1.10) + + def my_igather(self, rank, size, group, sendbuf, recvbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) + else: + recvbuf[rank] = sendbuf + else: + req.append(dist.isend(sendbuf, group=group, dst=root)) + return req + + def my_gather(self, rank, size, group, sendbuf, recvbuf, root): + if rank == root: + for idx in range(size): + if idx != rank: + dist.recv(recvbuf[idx], src=idx, group=group) + else: + recvbuf[rank] = sendbuf + else: + dist.send(sendbuf, group=group, dst=root) + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + + # all_start_time = time.time() + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + cupy.cuda.Device(local_rank).use() + + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(buffer_m.numel()) + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + if self.bool_not_supported: + cupy_sign_list_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool().to(dtype=torch.uint8)), self.size) + else: + cupy_sign_list_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool()), self.size) + cupy_worker_scale = self.compression_backend.torch2cupy(worker_scale) + + cupy_recvbuf_sign = cupy.zeros([self.size, cupy_sign_list_packed[self.rank].size], + dtype=cupy_sign_list_packed[0].dtype) + # cupy_recvbuf_scale = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) + + sign_list_packed = [ + self.compression_backend.cupy2torch(cupy_sign_list_packed[idx]) for idx in range(self.size) + ] + + # worker_scale = self.compression_backend.cupy2torch(cupy_worker_scale) + recvbuf_sign = self.compression_backend.cupy2torch(cupy_recvbuf_sign) + #recvbuf_scale = self.compression_backend.cupy2torch(cupy_recvbuf_scale) + recvbuf_scale = [ + torch.zeros(1, dtype=worker_scale.dtype, device=torch.device(get_accelerator().device_name(local_rank))) + for i in range(self.size) + ] + + # communication phase 1 + # gather_start = time.time() + # Alltoall for sign + dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) + # Allgather for scale + dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) + + # gather_end = time.time() + + # cupy_sign_list_packed, sign_list_packed, cupy_worker_scale, worker_scale = None, None, None, None + cupy_sign_list_packed = None + + cupy_recvbuf_sign = self.compression_backend.torch2cupy(recvbuf_sign) + #cupy_recvbuf_scale = self.compression_backend.torch2cupy(torch.stack(recvbuf_scale)) + + compensated_server_m = self.compression_backend.cupy2torch( + (cupy.unpackbits(cupy_recvbuf_sign.flatten())).reshape(self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) + compensated_server_m.add_(server_error) + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + # cupy_server_scale = self.compression_backend.torch2cupy(server_scale) + + if self.bool_not_supported: + cupy_server_sign_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool().to(dtype=torch.uint8)), + 1) + else: + cupy_server_sign_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool()), 1) + compensated_server_m = None + + cupy_recvbuf_sign_server = cupy.zeros([self.size, cupy_server_sign_packed[0].size], + dtype=cupy_recvbuf_sign.dtype) + # cupy_recvbuf_sign, recvbuf_sign = None, None + cupy_recvbuf_sign = None + + server_sign_packed = [self.compression_backend.cupy2torch(cupy_server_sign_packed[0])] + recvbuf_sign_server = [ + self.compression_backend.cupy2torch(cupy_recvbuf_sign_server[idx]) for idx in range(self.size) + ] + + # server_scale = self.compression_backend.cupy2torch(cupy_server_scale) + cupy_recvbuf_scale_server = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) + # cupy_recvbuf_scale, recvbuf_scale = None, None + + recvbuf_scale_server = [ + self.compression_backend.cupy2torch(cupy_recvbuf_scale_server[idx]) for idx in range(self.size) + ] + + # Communication Phase 2 + dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) + dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) + + cupy_server_sign_packed = None + + # need to convert from a tensor list to a single tensor + # dist.all_gather only provides a tensor list as the recv/output buffer + recvbuf_sign_server = torch.stack(recvbuf_sign_server) + + cupy_recvbuf_sign_server = self.compression_backend.torch2cupy(recvbuf_sign_server) + + buffer_m.data.copy_( + self.compression_backend.cupy2torch((cupy.unpackbits(cupy_recvbuf_sign_server.flatten())).reshape( + self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + self.compression_backend.cupy2torch(cupy_recvbuf_scale_server)).flatten().data) + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compiler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/compiler.py new file mode 100644 index 0000000000000000000000000000000000000000..be778b83f8bb8f3b9416fdf1d5717532367a8e4b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/compiler.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils.torch import required_torch_version + +try: + from torch.compiler import is_compiling as torch_is_compiling +except ImportError: + try: + from torch._dynamo.external_utils import is_compiling as torch_is_compiling + except ImportError: + # Torch does not have compiler support + torch_is_compiling = lambda: False + + +def is_compile_supported(): + return required_torch_version(min_version=2.1) + + +def disable(func): + if is_compile_supported(): + return torch.compiler.disable(func) + return func + + +def is_compiling(): + return torch_is_compiling() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b8d439bfbf4ad2903d701886ec163f5b55a05f7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b743bf3e021e3d27f8a70a8b82bf7d726917e94d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/cupy.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/cupy.py new file mode 100644 index 0000000000000000000000000000000000000000..7133ac04ed2b65bd496656870aebbf95ac04d3ad --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/cupy.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import cupy +from torch.utils.dlpack import to_dlpack +from torch.utils.dlpack import from_dlpack + + +class CupyBackend(object): + + def __init__(self): + pass + + def torch2cupy(self, tensor): + return cupy.from_dlpack(to_dlpack(tensor)) + + def cupy2torch(self, cupy_tensor): + return from_dlpack(cupy_tensor) + + def compress_by_chunk(self, cupy_bool_tensor, num_chunks): + packed_sign = cupy.packbits(cupy_bool_tensor) + sign_list_packed = cupy.split(packed_sign, num_chunks) + cupy.cuda.get_current_stream().synchronize() + return sign_list_packed diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/config.py new file mode 100644 index 0000000000000000000000000000000000000000..6a18ff193aec1e4d47c1feb1b015e94693547a3e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/config.py @@ -0,0 +1,997 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from typing import Union +from enum import Enum + +import torch +import json +import hjson +import copy +import base64 + +from .constants import * +from .config_utils import ( + get_scalar_param, + dict_raise_error_on_duplicate_keys, + ScientificNotationEncoder, +) +from .zero.config import get_zero_config, ZeroStageEnum +from .activation_checkpointing.config import DeepSpeedActivationCheckpointingConfig +from ..comm.config import DeepSpeedCommsConfig +from ..monitor.config import get_monitor_config +from ..inference.config import WeightQuantConfig +from .precision_config import get_bfloat16_config, get_float16_config +from ..compile.config import CompileConfig + +from deepspeed import comm as dist +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + +from ..git_version_info import version as __version__ +from ..utils import logger + +from ..elasticity import ( + elasticity_enabled, + compute_elastic_config, + ensure_immutable_elastic_config, +) +from ..elasticity.config import ElasticityConfigError +from ..elasticity.constants import ( + ELASTICITY, + IGNORE_NON_ELASTIC_BATCH_INFO, + IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT, + MODEL_PARALLEL_SIZE, + MODEL_PARALLEL_SIZE_DEFAULT, + NUM_GPUS_PER_NODE, + NUM_GPUS_PER_NODE_DEFAULT, +) + +from ..profiling.config import DeepSpeedFlopsProfilerConfig +from ..autotuning.config import DeepSpeedAutotuningConfig +from ..nebula.config import DeepSpeedNebulaConfig + +from ..compression.config import get_compression_config, get_quantize_enabled +from ..compression.constants import * +from .swap_tensor.aio_config import get_aio_config +from .model_checkpointing.config import get_checkpoint_config + +from .tensor_parallel import get_tensor_parallel_config +from .data_pipeline.config import get_data_efficiency_enabled, get_data_efficiency_config, get_curriculum_enabled_legacy, get_curriculum_params_legacy +from .data_pipeline.constants import * + +from ..utils.config import get_timers_config + +TENSOR_CORE_ALIGN_SIZE = 8 + +ADAGRAD_OPTIMIZER = 'adagrad' +ADAM_OPTIMIZER = 'adam' +ADAMW_OPTIMIZER = 'adamw' +LAMB_OPTIMIZER = 'lamb' +ONEBIT_ADAM_OPTIMIZER = 'onebitadam' +ZERO_ONE_ADAM_OPTIMIZER = 'zerooneadam' +ONEBIT_LAMB_OPTIMIZER = 'onebitlamb' +MUADAM_OPTIMIZER = 'muadam' +MUADAMW_OPTIMIZER = 'muadamw' +MUSGD_OPTIMIZER = 'musgd' +LION_OPTIMIZER = 'lion' +DEEPSPEED_OPTIMIZERS = [ + ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, + ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, MUSGD_OPTIMIZER, LION_OPTIMIZER +] + +# extra optimizer parameters for adam/adamw +TORCH_ADAM_PARAM = "torch_adam" + +# default to adamw logic for adam/adamw optimizers unless user explicitly opts out +ADAM_W_MODE = "adam_w_mode" +ADAM_W_MODE_DEFAULT = True + + +class DeepSpeedConfigError(Exception): + pass + + +class DtypeEnum(Enum): + # The torch dtype must always be the first value (so we return torch.dtype) + fp16 = torch.float16, "torch.float16", "fp16", "float16", "half" + fp32 = torch.float32, "torch.float32", "fp32", "float32", "float" + int8 = torch.int8, "torch.int8", "int8" + bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16" + + # Copied from https://stackoverflow.com/a/43210118 + # Allows us to use multiple values for each Enum index and returns first + # listed value when Enum is called + def __new__(cls, *values): + obj = object.__new__(cls) + # first value is canonical value + obj._value_ = values[0] + for other_value in values[1:]: + cls._value2member_map_[other_value] = obj + obj._all_values = values + return obj + + def __repr__(self): + return "<%s.%s: %s>" % ( + self.__class__.__name__, + self._name_, + ", ".join([repr(v) for v in self._all_values]), + ) + + +def get_pld_enabled(param_dict): + if PROGRESSIVE_LAYER_DROP in param_dict.keys(): + return get_scalar_param(param_dict[PROGRESSIVE_LAYER_DROP], PLD_ENABLED, PLD_ENABLED_DEFAULT) + else: + return False + + +def get_pld_params(param_dict): + if PROGRESSIVE_LAYER_DROP in param_dict.keys(): + pld_params = copy.copy(param_dict[PROGRESSIVE_LAYER_DROP]) + pld_params.pop(PLD_ENABLED) + return pld_params + else: + return False + + +def get_amp_enabled(param_dict): + if AMP in param_dict.keys(): + return get_scalar_param(param_dict[AMP], AMP_ENABLED, AMP_ENABLED_DEFAULT) + else: + return False + + +def get_amp_params(param_dict): + if AMP in param_dict.keys(): + amp_params = copy.copy(param_dict[AMP]) + amp_params.pop(AMP_ENABLED) + return amp_params + else: + return False + + +def get_torch_autocast_enabled(param_dict): + if TORCH_AUTOCAST in param_dict.keys(): + return get_scalar_param(param_dict[TORCH_AUTOCAST], TORCH_AUTOCAST_ENABLED, TORCH_AUTOCAST_ENABLED_DEFAULT) + else: + return False + + +def get_torch_autocast_dtype(param_dict): + if TORCH_AUTOCAST in param_dict: + if TORCH_AUTOCAST_DTYPE in param_dict[TORCH_AUTOCAST]: + try: + return DtypeEnum(param_dict[TORCH_AUTOCAST][TORCH_AUTOCAST_DTYPE]).value + except KeyError: + raise ValueError( + f"Invalid dtype for torch autocast: {param_dict[TORCH_AUTOCAST][TORCH_AUTOCAST_DTYPE]}") + return None + + +def get_lower_precision_safe_modules(param_dict): + if TORCH_AUTOCAST in param_dict: + if TORCH_AUTOCAST_LOWER_PRECISION_SAFE_MODULES in param_dict[TORCH_AUTOCAST]: + module_names_with_package = param_dict[TORCH_AUTOCAST][TORCH_AUTOCAST_LOWER_PRECISION_SAFE_MODULES] + if not all(isinstance(module_name, str) for module_name in module_names_with_package): + raise ValueError( + f"Invalid module names for torch autocast: {module_names_with_package}. Expected list of strings.") + return module_names_with_package + return None + + +def get_gradient_accumulation_steps(param_dict): + return get_scalar_param(param_dict, GRADIENT_ACCUMULATION_STEPS, GRADIENT_ACCUMULATION_STEPS_DEFAULT) + + +def get_sparse_gradients_enabled(param_dict): + return get_scalar_param(param_dict, SPARSE_GRADIENTS, SPARSE_GRADIENTS_DEFAULT) + + +def get_communication_data_type(param_dict, + comm_type=COMMUNICATION_DATA_TYPE, + comm_data_type_default=COMMUNICATION_DATA_TYPE_DEFAULT): + val = get_scalar_param(param_dict, comm_type, comm_data_type_default) + val = val.lower() if val is not None else val + if val is None: + return val # we must determine it by other parameters + elif val == "fp32": + return torch.float32 + elif val == "fp16": + return torch.float16 + elif val == "bf16": + return torch.bfloat16 + + raise ValueError(f"Invalid communication_data_type. Supported data types: ['fp16', 'bf16', 'fp32']. Got: {val}") + + +def get_prescale_gradients(param_dict): + return get_scalar_param(param_dict, PRESCALE_GRADIENTS, PRESCALE_GRADIENTS_DEFAULT) + + +def get_gradient_predivide_factor(param_dict): + return get_scalar_param(param_dict, GRADIENT_PREDIVIDE_FACTOR, GRADIENT_PREDIVIDE_FACTOR_DEFAULT) + + +def get_steps_per_print(param_dict): + return get_scalar_param(param_dict, STEPS_PER_PRINT, STEPS_PER_PRINT_DEFAULT) + + +def get_disable_allgather(param_dict): + return get_scalar_param(param_dict, DISABLE_ALLGATHER, DISABLE_ALLGATHER_DEFAULT) + + +def get_dump_state(param_dict): + return get_scalar_param(param_dict, DUMP_STATE, DUMP_STATE_DEFAULT) + + +def get_gradient_clipping(param_dict): + return get_scalar_param(param_dict, GRADIENT_CLIPPING, GRADIENT_CLIPPING_DEFAULT) + + +def get_graph_harvesting(param_dict): + return get_scalar_param(param_dict, GRAPH_HARVESTING, GRAPH_HARVESTING_DEFAULT) + + +def get_sparse_attention(param_dict): + if SPARSE_ATTENTION in param_dict.keys(): + sparsity = param_dict[SPARSE_ATTENTION] + mode = get_sparse_attention_mode(sparsity) + + if mode == SPARSE_DENSE_MODE: + return get_sparse_dense_config(sparsity) + elif mode == SPARSE_FIXED_MODE: + return get_sparse_fixed_config(sparsity) + elif mode == SPARSE_VARIABLE_MODE: + return get_sparse_variable_config(sparsity) + elif mode == SPARSE_BIGBIRD_MODE: + return get_sparse_bigbird_config(sparsity) + elif mode == SPARSE_BSLONGFORMER_MODE: + return get_sparse_bslongformer_config(sparsity) + else: + raise NotImplementedError(f"Given sparsity mode, {mode}, has not been implemented yet!") + + else: + return None + + +def get_sparse_dense_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + return {SPARSE_MODE: SPARSE_DENSE_MODE, SPARSE_BLOCK: block} + + +def get_sparse_fixed_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_local_blocks = get_scalar_param(sparsity, SPARSE_NUM_LOCAL_BLOCKS, SPARSE_NUM_LOCAL_BLOCKS_DEFAULT) + num_global_blocks = get_scalar_param(sparsity, SPARSE_NUM_GLOBAL_BLOCKS, SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT) + attention = get_scalar_param(sparsity, SPARSE_ATTENTION_TYPE, SPARSE_ATTENTION_TYPE_DEFAULT) + horizontal_global_attention = get_scalar_param( + sparsity, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT, + ) + num_different_global_patterns = get_scalar_param( + sparsity, + SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS, + SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS_DEFAULT, + ) + + return { + SPARSE_MODE: SPARSE_FIXED_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_LOCAL_BLOCKS: num_local_blocks, + SPARSE_NUM_GLOBAL_BLOCKS: num_global_blocks, + SPARSE_ATTENTION_TYPE: attention, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION: horizontal_global_attention, + SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS: num_different_global_patterns, + } + + +def get_sparse_variable_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_random_blocks = get_scalar_param(sparsity, SPARSE_NUM_RANDOM_BLOCKS, SPARSE_NUM_RANDOM_BLOCKS_DEFAULT) + local_window_blocks = get_scalar_param(sparsity, SPARSE_LOCAL_WINDOW_BLOCKS, SPARSE_LOCAL_WINDOW_BLOCKS_DEFAULT) + global_block_indices = get_scalar_param(sparsity, SPARSE_GLOBAL_BLOCK_INDICES, SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT) + global_block_end_indices = get_scalar_param( + sparsity, + SPARSE_GLOBAL_BLOCK_END_INDICES, + SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT, + ) + attention = get_scalar_param(sparsity, SPARSE_ATTENTION_TYPE, SPARSE_ATTENTION_TYPE_DEFAULT) + horizontal_global_attention = get_scalar_param( + sparsity, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT, + ) + + return { + SPARSE_MODE: SPARSE_VARIABLE_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_RANDOM_BLOCKS: num_random_blocks, + SPARSE_LOCAL_WINDOW_BLOCKS: local_window_blocks, + SPARSE_GLOBAL_BLOCK_INDICES: global_block_indices, + SPARSE_GLOBAL_BLOCK_END_INDICES: global_block_end_indices, + SPARSE_ATTENTION_TYPE: attention, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION: horizontal_global_attention, + } + + +def get_sparse_bigbird_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_random_blocks = get_scalar_param(sparsity, SPARSE_NUM_RANDOM_BLOCKS, SPARSE_NUM_RANDOM_BLOCKS_DEFAULT) + num_sliding_window_blocks = get_scalar_param( + sparsity, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT, + ) + num_global_blocks = get_scalar_param(sparsity, SPARSE_NUM_GLOBAL_BLOCKS, SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT) + + return { + SPARSE_MODE: SPARSE_BIGBIRD_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_RANDOM_BLOCKS: num_random_blocks, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS: num_sliding_window_blocks, + SPARSE_NUM_GLOBAL_BLOCKS: num_global_blocks, + } + + +def get_sparse_bslongformer_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_sliding_window_blocks = get_scalar_param( + sparsity, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT, + ) + global_block_indices = get_scalar_param(sparsity, SPARSE_GLOBAL_BLOCK_INDICES, SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT) + global_block_end_indices = get_scalar_param( + sparsity, + SPARSE_GLOBAL_BLOCK_END_INDICES, + SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT, + ) + + return { + SPARSE_MODE: SPARSE_BSLONGFORMER_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS: num_sliding_window_blocks, + SPARSE_GLOBAL_BLOCK_INDICES: global_block_indices, + SPARSE_GLOBAL_BLOCK_END_INDICES: global_block_end_indices, + } + + +def get_sparse_attention_mode(param_dict): + if SPARSE_MODE in param_dict.keys(): + return param_dict[SPARSE_MODE] + else: + return SPARSE_MODE_DEFAULT + + +def get_sparse_attention_type(param_dict): + if SPARSE_ATTENTION_TYPE in param_dict.keys(): + return param_dict[SPARSE_ATTENTION_TYPE] + else: + return SPARSE_ATTENTION_TYPE_DEFAULT + + +def get_pipeline_config(param_dict): + """Parses pipeline engine configuration. """ + default_pipeline = { + "stages": "auto", + "partition": "best", + "seed_layers": False, + "activation_checkpoint_interval": 0, + "pipe_partitioned": True, + "grad_partitioned": True, + } + config = default_pipeline + for key, val in param_dict.get("pipeline", {}).items(): + config[key] = val + return config + + +def get_optimizer_name(param_dict): + if OPTIMIZER in param_dict.keys() and TYPE in param_dict[OPTIMIZER].keys(): + return param_dict[OPTIMIZER][TYPE] + else: + return OPTIMIZER_TYPE_DEFAULT + + +def get_optimizer_params(param_dict): + if (get_optimizer_name(param_dict) is not None and OPTIMIZER_PARAMS in param_dict[OPTIMIZER].keys()): + return param_dict[OPTIMIZER][OPTIMIZER_PARAMS] + else: + return None + + +def get_optimizer_gradient_clipping(param_dict): + optimizer_params = get_optimizer_params(param_dict) + if optimizer_params is not None and MAX_GRAD_NORM in optimizer_params.keys(): + return optimizer_params[MAX_GRAD_NORM] + else: + return None + + +def get_optimizer_legacy_fusion(param_dict): + if OPTIMIZER in param_dict.keys() and LEGACY_FUSION in param_dict[OPTIMIZER].keys(): + return param_dict[OPTIMIZER][LEGACY_FUSION] + else: + return LEGACY_FUSION_DEFAULT + + +def get_zero_allow_untested_optimizer(param_dict): + return get_scalar_param(param_dict, ZERO_ALLOW_UNTESTED_OPTIMIZER, ZERO_ALLOW_UNTESTED_OPTIMIZER_DEFAULT) + + +def get_zero_force_ds_cpu_optimizer(param_dict): + return get_scalar_param(param_dict, ZERO_FORCE_DS_CPU_OPTIMIZER, ZERO_FORCE_DS_CPU_OPTIMIZER_DEFAULT) + + +def get_scheduler_name(param_dict): + if SCHEDULER in param_dict.keys() and TYPE in param_dict[SCHEDULER].keys(): + return param_dict[SCHEDULER][TYPE] + else: + return SCHEDULER_TYPE_DEFAULT + + +def get_scheduler_params(param_dict): + if (get_scheduler_name(param_dict) is not None and SCHEDULER_PARAMS in param_dict[SCHEDULER].keys()): + return param_dict[SCHEDULER][SCHEDULER_PARAMS] + else: + return None + + +def get_train_batch_size(param_dict): + return get_scalar_param(param_dict, TRAIN_BATCH_SIZE, TRAIN_BATCH_SIZE_DEFAULT) + + +def get_train_micro_batch_size_per_gpu(param_dict): + return get_scalar_param( + param_dict, + TRAIN_MICRO_BATCH_SIZE_PER_GPU, + TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT, + ) + + +def get_wall_clock_breakdown(param_dict): + return get_scalar_param(param_dict, WALL_CLOCK_BREAKDOWN, WALL_CLOCK_BREAKDOWN_DEFAULT) + + +def get_memory_breakdown(param_dict): + return get_scalar_param(param_dict, MEMORY_BREAKDOWN, MEMORY_BREAKDOWN_DEFAULT) + + +class HybridEngineConfig(DeepSpeedConfigModel): + enabled: bool = False + max_out_tokens: int = 512 + inference_tp_size: int = 1 + release_inference_cache: bool = False + pin_parameters: bool = True + tp_gather_partition_size: int = 8 + + +def get_hybrid_engine_config(param_dict): + hybrid_engine_config_dict = param_dict.get("hybrid_engine", {}) + hybrid_engine_config = HybridEngineConfig(**hybrid_engine_config_dict) + return hybrid_engine_config + + +def get_expert_data_topo_config(param_dict): + return get_scalar_param(param_dict, USE_DATA_BEFORE_EXPERT_PARALLEL, USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT) + + +def get_eigenvalue_config(param_dict): + if get_quantize_enabled(param_dict): + param_dict = param_dict[QUANTIZE_TRAINING] + assert not get_eigenvalue_enabled(param_dict), "Eigenvalue based MoQ is temporarily disabled" + return ( + get_eigenvalue_enabled(param_dict), + get_eigenvalue_verbose(param_dict), + get_eigenvalue_max_iter(param_dict), + get_eigenvalue_tol(param_dict), + get_eigenvalue_stability(param_dict), + get_eigenvalue_gas_boundary_resolution(param_dict), + get_eigenvalue_layer_name(param_dict), + get_eigenvalue_layer_num(param_dict), + ) + else: + return ( + EIGENVALUE_ENABLED_DEFAULT, + EIGENVALUE_VERBOSE_DEFAULT, + EIGENVALUE_MAX_ITER_DEFAULT, + EIGENVALUE_TOL_DEFAULT, + EIGENVALUE_STABILITY_DEFAULT, + EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT, + EIGENVALUE_LAYER_NAME_DEFAULT, + EIGENVALUE_LAYER_NUM_DEFAULT, + ) + + +def get_eigenvalue_enabled(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_ENABLED, EIGENVALUE_ENABLED_DEFAULT) + else: + return EIGENVALUE_ENABLED_DEFAULT + + +def get_eigenvalue_verbose(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_VERBOSE, EIGENVALUE_VERBOSE_DEFAULT) + else: + return EIGENVALUE_VERBOSE_DEFAULT + + +def get_eigenvalue_max_iter(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_MAX_ITER, EIGENVALUE_MAX_ITER_DEFAULT) + else: + return EIGENVALUE_MAX_ITER_DEFAULT + + +def get_eigenvalue_tol(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_TOL, EIGENVALUE_TOL_DEFAULT) + else: + return EIGENVALUE_TOL_DEFAULT + + +def get_eigenvalue_stability(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_STABILITY, EIGENVALUE_STABILITY_DEFAULT) + else: + return EIGENVALUE_STABILITY_DEFAULT + + +def get_eigenvalue_gas_boundary_resolution(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param( + param_dict[EIGENVALUE], + EIGENVALUE_GAS_BOUNDARY_RESOLUTION, + EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT, + ) + else: + return EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT + + +def get_eigenvalue_layer_name(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_LAYER_NAME, EIGENVALUE_LAYER_NAME_DEFAULT) + else: + return EIGENVALUE_LAYER_NAME_DEFAULT + + +def get_eigenvalue_layer_num(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_LAYER_NUM, EIGENVALUE_LAYER_NUM_DEFAULT) + else: + return EIGENVALUE_LAYER_NUM_DEFAULT + + +def get_checkpoint_params(param_dict): + return param_dict.get(CHECKPOINT, {}) + + +def get_data_types_params(param_dict): + return param_dict.get(DATA_TYPES, {}) + + +def get_checkpoint_tag_validation_mode(checkpoint_params): + tag_validation_mode = checkpoint_params.get(CHECKPOINT_TAG_VALIDATION, CHECKPOINT_TAG_VALIDATION_DEFAULT) + tag_validation_mode = tag_validation_mode.upper() + if tag_validation_mode in CHECKPOINT_TAG_VALIDATION_MODES: + return tag_validation_mode + else: + raise DeepSpeedConfigError( + "Checkpoint config contains invalid tag_validation " + f"value of {tag_validation_mode}, expecting one of {CHECKPOINT_TAG_VALIDATION_MODES}") + + +def get_checkpoint_parallel_write_pipeline(checkpoint_params): + par_write_params = checkpoint_params.get(CHECKPOINT_PARALLEL_WRITE, {}) + par_write_pipeline = par_write_params.get(CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE, + CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE_DEFAULT) + if par_write_pipeline in [True, False]: + return par_write_pipeline + else: + raise DeepSpeedConfigError("checkpoint::parallel_write::pipeline_stage " + f"value of '{par_write_pipeline}' is invalid, expecting: true or false") + + +def get_dataloader_drop_last(param_dict): + return get_scalar_param(param_dict, DATALOADER_DROP_LAST, DATALOADER_DROP_LAST_DEFAULT) + + +'''Write deepspeed config files by modifying basic templates. +Can be used for quickly changing parameters via command line parameters.''' + + +class DeepSpeedConfigWriter: + + def __init__(self, data=None): + self.data = data if data is not None else {} + + def add_config(self, key, value): + self.data[key] = value + + def load_config(self, filename): + self.data = json.load(open(filename, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) + + def write_config(self, filename): + with open(filename, "w") as outfile: + json.dump(self.data, outfile) + + +class DeepSpeedConfig(object): + + def __init__(self, config: Union[str, dict], mpu=None, mesh_device=None): + super(DeepSpeedConfig, self).__init__() + if isinstance(config, dict): + self._param_dict = config + elif os.path.exists(config): + self._param_dict = hjson.load(open(config, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) + else: + try: + config_decoded = base64.urlsafe_b64decode(config).decode('utf-8') + self._param_dict = hjson.loads(config_decoded) + except (UnicodeDecodeError, AttributeError): + raise ValueError( + f"Expected a string path to an existing deepspeed config, or a dictionary or a valid base64. Received: {config}" + ) + + try: + self.global_rank = dist.get_rank() + if mpu is not None: + # Ulysses SP + if not hasattr(mpu, "get_data_parallel_world_size"): + self.world_size = dist.get_world_size() / mpu.get_sequence_parallel_world_size() + else: + self.world_size = mpu.get_data_parallel_world_size() + elif mesh_device is not None: + self.world_size = dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) + else: + # HF zero.init case where there is no mpu + if "sequence_parallel_size" in config: + self.world_size = dist.get_world_size() / config["sequence_parallel_size"] + else: + self.world_size = dist.get_world_size() + except: + self.global_rank = 0 + self.world_size = 1 + logger.info(f"Config mesh_device {mesh_device} world_size = {self.world_size}") + # If elastic-mode enabled, update compute + update _param_dict + self.elasticity_enabled = elasticity_enabled(self._param_dict) + if self.elasticity_enabled: + logger.info("DeepSpeed elasticity support enabled") + final_batch_size, valid_gpus, micro_batch_size = compute_elastic_config( + ds_config=self._param_dict, + target_deepspeed_version=__version__, + world_size=self.world_size, + ) + + elastic_dict = self._param_dict[ELASTICITY] + + # Ensure the resource scheduler saw the same elastic config we are using at runtime + ensure_immutable_elastic_config(runtime_elastic_config_dict=elastic_dict) + + self.elastic_model_parallel_size = elastic_dict.get(MODEL_PARALLEL_SIZE, MODEL_PARALLEL_SIZE_DEFAULT) + if self.elastic_model_parallel_size < 1: + raise ElasticityConfigError("Model-Parallel size cannot be less than 1, " + f"given model-parallel size: {self.elastic_model_parallel_size}") + + self.num_gpus_per_node = elastic_dict.get(NUM_GPUS_PER_NODE, NUM_GPUS_PER_NODE_DEFAULT) + if self.num_gpus_per_node < 1: + raise ElasticityConfigError("NUmber of GPUs per node cannot be less than 1, " + f"given number of GPUs per node: {self.num_gpus_per_node}") + + ignore_non_elastic_batch_info = elastic_dict.get(IGNORE_NON_ELASTIC_BATCH_INFO, + IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT) + + if not ignore_non_elastic_batch_info: + batch_params = [ + TRAIN_BATCH_SIZE, + TRAIN_MICRO_BATCH_SIZE_PER_GPU, + GRADIENT_ACCUMULATION_STEPS, + ] + if any(map(lambda t: t in self._param_dict, batch_params)): + raise ElasticityConfigError("One or more batch related parameters were found in your " \ + f"ds_config ({TRAIN_BATCH_SIZE}, {TRAIN_MICRO_BATCH_SIZE_PER_GPU}, and/or " \ + f"{GRADIENT_ACCUMULATION_STEPS}). These parameters *will not be used* since " \ + "elastic training is enabled, which takes control of these parameters. " \ + "If you want to suppress this error (the parameters will be silently ignored) " \ + f"please set {IGNORE_NON_ELASTIC_BATCH_INFO}':true in your elasticity config.") + + # micro_bsz * world_size * gas = total_batch_size + # gas = total_batch_size // (micro_bsz * world_size) + gradient_accu_steps = final_batch_size // (micro_batch_size * self.world_size) + + if TRAIN_BATCH_SIZE in self._param_dict: + logger.warning("[Elasticity] overriding training_batch_size: " + f"{self._param_dict[TRAIN_BATCH_SIZE]} -> {final_batch_size}") + if TRAIN_MICRO_BATCH_SIZE_PER_GPU in self._param_dict: + logger.warning("[Elasticity] overriding train_micro_batch_size_per_gpu: " + f"{self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU]} -> {micro_batch_size}") + if GRADIENT_ACCUMULATION_STEPS in self._param_dict: + logger.warning("[Elasticity] overriding gradient_accumulation_steps: " + f"{self._param_dict[GRADIENT_ACCUMULATION_STEPS]} -> {gradient_accu_steps}") + + logger.info(f"[Elasticity] valid GPU counts: {valid_gpus}") + + self._param_dict[TRAIN_BATCH_SIZE] = final_batch_size + self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = micro_batch_size + self._param_dict[GRADIENT_ACCUMULATION_STEPS] = gradient_accu_steps + + # Pass a copy so that user json is unmodified, e.g. for logging + self._initialize_params(copy.copy(self._param_dict)) + self._configure_train_batch_size() + self._do_sanity_check() + + def _initialize_params(self, param_dict): + self.train_batch_size = get_train_batch_size(param_dict) + self.train_micro_batch_size_per_gpu = get_train_micro_batch_size_per_gpu(param_dict) + self.gradient_accumulation_steps = get_gradient_accumulation_steps(param_dict) + self.steps_per_print = get_steps_per_print(param_dict) + self.dump_state = get_dump_state(param_dict) + + self.disable_allgather = get_disable_allgather(param_dict) + self.communication_data_type = get_communication_data_type(param_dict) + self.seq_parallel_communication_data_type = get_communication_data_type( + param_dict, SEQ_PARALLEL_COMMUNICATION_DATA_TYPE, SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT) + self.prescale_gradients = get_prescale_gradients(param_dict) + self.gradient_predivide_factor = get_gradient_predivide_factor(param_dict) + self.sparse_gradients_enabled = get_sparse_gradients_enabled(param_dict) + + self.zero_config = get_zero_config(param_dict) + self.mics_shard_size = self.zero_config.mics_shard_size + self.mics_hierarchial_params_gather = self.zero_config.mics_hierarchical_params_gather + self.zero_optimization_stage = self.zero_config.stage + self.zero_enabled = self.zero_optimization_stage > 0 + + self.activation_checkpointing_config = DeepSpeedActivationCheckpointingConfig(param_dict) + + self.comms_config = DeepSpeedCommsConfig(param_dict) + self.monitor_config = get_monitor_config(param_dict) + + self.gradient_clipping = get_gradient_clipping(param_dict) + self.float16_config = get_float16_config(param_dict) + self.bfloat16_config = get_bfloat16_config(param_dict) + assert not (self.float16_config.enabled + and self.bfloat16_config.enabled), 'bfloat16 and fp16 modes cannot be simultaneously enabled' + + self.amp_enabled = get_amp_enabled(param_dict) + self.amp_params = get_amp_params(param_dict) + + self.torch_autocast_enabled = get_torch_autocast_enabled(param_dict) + self.torch_autocast_dtype = get_torch_autocast_dtype(param_dict) + self.torch_autocast_lower_precision_safe_modules = get_lower_precision_safe_modules(param_dict) + + self.compression_config = get_compression_config(param_dict) + self.graph_harvesting = get_graph_harvesting(param_dict) + + self.optimizer_name = get_optimizer_name(param_dict) + if (self.optimizer_name is not None and self.optimizer_name.lower() in DEEPSPEED_OPTIMIZERS): + self.optimizer_name = self.optimizer_name.lower() + + self.optimizer_params = get_optimizer_params(param_dict) + self.optimizer_legacy_fusion = get_optimizer_legacy_fusion(param_dict) + + self.zero_allow_untested_optimizer = get_zero_allow_untested_optimizer(param_dict) + + self.zero_force_ds_cpu_optimizer = get_zero_force_ds_cpu_optimizer(param_dict) + + self.scheduler_name = get_scheduler_name(param_dict) + self.scheduler_params = get_scheduler_params(param_dict) + + self.flops_profiler_config = DeepSpeedFlopsProfilerConfig(param_dict) + self.wall_clock_breakdown = (get_wall_clock_breakdown(param_dict) | self.flops_profiler_config.enabled) + self.memory_breakdown = get_memory_breakdown(param_dict) + self.autotuning_config = DeepSpeedAutotuningConfig(param_dict) + + ( + self.eigenvalue_enabled, + self.eigenvalue_verbose, + self.eigenvalue_max_iter, + self.eigenvalue_tol, + self.eigenvalue_stability, + self.eigenvalue_gas_boundary_resolution, + self.eigenvalue_layer_name, + self.eigenvalue_layer_num, + ) = get_eigenvalue_config(param_dict) + + self.use_data_before_expert_parallel_ = get_expert_data_topo_config(param_dict) + self.hybrid_engine = get_hybrid_engine_config(param_dict) + + self.sparse_attention = get_sparse_attention(param_dict) + self.pipeline = get_pipeline_config(param_dict) + + self.pld_enabled = get_pld_enabled(param_dict) + self.pld_params = get_pld_params(param_dict) + + self.curriculum_enabled_legacy = get_curriculum_enabled_legacy(param_dict) + self.curriculum_params_legacy = get_curriculum_params_legacy(param_dict) + + self.data_efficiency_enabled = get_data_efficiency_enabled(param_dict) + self.data_efficiency_config = get_data_efficiency_config(param_dict) + + checkpoint_params = get_checkpoint_params(param_dict) + validation_mode = get_checkpoint_tag_validation_mode(checkpoint_params) + self.checkpoint_tag_validation_enabled = (validation_mode != ValidationMode.IGNORE) + self.checkpoint_tag_validation_fail = validation_mode == ValidationMode.FAIL + self.load_universal_checkpoint = checkpoint_params.get(LOAD_UNIVERSAL_CHECKPOINT, + LOAD_UNIVERSAL_CHECKPOINT_DEFAULT) + + self.use_node_local_storage = checkpoint_params.get(USE_NODE_LOCAL_STORAGE_CHECKPOINT, + USE_NODE_LOCAL_STORAGE_CHECKPOINT_DEFAULT) + + data_types_params = get_data_types_params(param_dict) + self.grad_accum_dtype = data_types_params.get(GRAD_ACCUM_DTYPE, GRAD_ACCUM_DTYPE_DEFAULT) + + par_write_pipe = get_checkpoint_parallel_write_pipeline(checkpoint_params) + self.checkpoint_parallel_write_pipeline = par_write_pipe + + self.aio_config = get_aio_config(param_dict) + + self.dataloader_drop_last = get_dataloader_drop_last(param_dict) + + self.nebula_config = DeepSpeedNebulaConfig(param_dict) + self.checkpoint_config = get_checkpoint_config(param_dict) + + self.weight_quantization_config = WeightQuantConfig( + **param_dict['weight_quantization']) if 'weight_quantization' in param_dict else None + + self.compile_config = CompileConfig(**param_dict.get('compile', {})) + + self.timers_config = get_timers_config(param_dict) + self.tensor_parallel_config = get_tensor_parallel_config(param_dict) + + def _batch_assertion(self): + + train_batch = self.train_batch_size + micro_batch = self.train_micro_batch_size_per_gpu + grad_acc = self.gradient_accumulation_steps + + assert (train_batch > 0), f"Train batch size: {train_batch} has to be greater than 0" + + assert (micro_batch > 0), f"Micro batch size per gpu: {micro_batch} has to be greater than 0" + + assert (grad_acc > 0), f"Gradient accumulation steps: {grad_acc} has to be greater than 0" + + assert train_batch == micro_batch * grad_acc * self.world_size, ( + f"Check batch related parameters. train_batch_size is not equal " + "to micro_batch_per_gpu * gradient_acc_step * world_size " + f"{train_batch} != {micro_batch} * {grad_acc} * {self.world_size}") + + def _set_batch_related_parameters(self): + + train_batch = self.train_batch_size + micro_batch = self.train_micro_batch_size_per_gpu + grad_acc = self.gradient_accumulation_steps + + #print(f"in: train_batch = {train_batch}, micro_batch={micro_batch}") + + # all values are provided nothing needs to be set + if train_batch is not None and micro_batch is not None and grad_acc is not None: + return + + # global_accumulation_steps needs to be set + elif train_batch is not None and micro_batch is not None: + grad_acc = train_batch // micro_batch + grad_acc //= self.world_size + self.gradient_accumulation_steps = grad_acc + + # micro_batch_per_gpu needs to be set + elif train_batch is not None and grad_acc is not None: + micro_batch = train_batch // self.world_size + micro_batch //= grad_acc + self.train_micro_batch_size_per_gpu = micro_batch + + # train_batch_size needs to be set + elif micro_batch is not None and grad_acc is not None: + train_batch_size = micro_batch * grad_acc + train_batch_size *= self.world_size + self.train_batch_size = train_batch_size + + # gradient_accumulation_steps and micro_batch_per_gpus is set + elif train_batch is not None: + self.gradient_accumulation_steps = 1 + self.train_micro_batch_size_per_gpu = train_batch // self.world_size + + # train_batch_size and gradient_accumulation_step is set + elif micro_batch is not None: + self.train_batch_size = micro_batch * self.world_size + self.gradient_accumulation_steps = 1 + + # either none of the three parameters are provided or just gradient_accumulation_step is provided + else: + assert False, \ + 'Either train_batch_size or train_micro_batch_size_per_gpu needs to be provided' + + #print(f"final: {self.train_batch_size=} {self.train_micro_batch_size_per_gpu=} {self.gradient_accumulation_steps=}") + + def _configure_train_batch_size(self): + self._set_batch_related_parameters() + self._batch_assertion() + + def _do_sanity_check(self): + self._do_error_check() + + self._do_warning_check() + + def print_user_config(self): + logger.info(" json = {}".format( + json.dumps( + self._param_dict, + sort_keys=True, + indent=4, + cls=ScientificNotationEncoder, + separators=(",", ":"), + ))) + + def print(self, name): + logger.info("{}:".format(name)) + for arg in sorted(vars(self)): + if arg != "_param_dict": + dots = "." * (29 - len(arg)) + logger.info(" {} {} {}".format(arg, dots, getattr(self, arg))) + + self.print_user_config() + + def _do_error_check(self): + assert (self.train_micro_batch_size_per_gpu + ), "DeepSpeedConfig: {} is not defined".format(TRAIN_MICRO_BATCH_SIZE_PER_GPU) + + assert ( + self.gradient_accumulation_steps), "DeepSpeedConfig: {} is not defined".format(GRADIENT_ACCUMULATION_STEPS) + + if self.zero_enabled: + assert (self.zero_optimization_stage + <= ZeroStageEnum.max_stage), "DeepSpeedConfig: Maximum supported ZeRO stage is {}".format( + ZeroStageEnum.max_stage) + + if self.float16_config.fp16_master_weights_and_grads: + assert self.zero_enabled and self.zero_optimization_stage == ZeroStageEnum.gradients, "Fp16_master_weights_and_grads is only supported with ZeRO Stage 2 for now." + + def _do_warning_check(self): + fp16_enabled = self.float16_config.enabled + + vocabulary_size = self._param_dict.get(VOCABULARY_SIZE, VOCABULARY_SIZE_DEFAULT) + if vocabulary_size and vocabulary_size % TENSOR_CORE_ALIGN_SIZE != 0: + logger.warning( + "DeepSpeedConfig: vocabulary size {} is not aligned to {}, may import tensor core utilization.".format( + vocabulary_size, TENSOR_CORE_ALIGN_SIZE)) + + if (self.optimizer_params is not None and MAX_GRAD_NORM in self.optimizer_params.keys() + and self.optimizer_params[MAX_GRAD_NORM] > 0): + if fp16_enabled: + if self.global_rank == 0: + logger.warning("DeepSpeedConfig: In FP16 mode, DeepSpeed will pass {}:{} to FP16 wrapper".format( + MAX_GRAD_NORM, self.optimizer_params[MAX_GRAD_NORM])) + else: + if self.global_rank == 0: + logger.warning( + "DeepSpeedConfig: In FP32 mode, DeepSpeed does not permit MAX_GRAD_NORM ({}) > 0, setting to zero" + .format(self.optimizer_params[MAX_GRAD_NORM])) + self.optimizer_params[MAX_GRAD_NORM] = 0.0 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/config_utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/config_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2ae6921c363b5201ce191c9630ea9daa3774e19b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/config_utils.py @@ -0,0 +1,212 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Collection of DeepSpeed configuration utilities +""" +import collections +import json +import torch +from functools import reduce +from pydantic import BaseModel, ConfigDict, field_serializer + +from deepspeed.utils import logger + + +class DeepSpeedConfigModel(BaseModel): + """ + This class should be used as a base for all DeepSpeed configs. It extends + pydantic.BaseModel to allow for deprecated fields. To enable this feature, + add deprecated=True to pydantic.Field: + + my_dep_field: int = Field(0, deprecated=True) + + Deprecated Field kwargs: + - deprecated: [True|False], default False + Enables / Disables deprecated fields + - deprecated_msg: str, default "" + Message to include with deprecation warning + - new_param: str, default "" + Name of the field replacing the deprecated field + - set_new_param: [True|False], default True + If new_param is provided, enables setting the value of that param with + deprecated field value + - new_param_fn: callable, default (lambda x: x) + If new_param is provided and set_new_param is True, this function will + modify the value of the deprecated field before placing that value in + the new_param field + + Example: + my_new_field is replacing a deprecated my_old_field. The expected type + for my_new_field is int while the expected type for my_old_field is + str. We want to maintain backward compatibility with our configs, so we + define the fields with: + + class MyExampleConfig(DeepSpeedConfigModel): + my_new_field: int = 0 + my_old_field: str = Field('0', + deprecated=True, + new_param='my_new_field', + new_param_fn=(lambda x: int(x))) + """ + + def __init__(self, strict=False, **data): + if (not strict): # This is temporary until we refactor all DS configs, allows HF to load models + data = {k: v for k, v in data.items() if (v != "auto" or k == "replace_method")} + super().__init__(**data) + self._deprecated_fields_check() + + def _process_deprecated_field(self, dep_field): + # Get information about the deprecated field + pydantic_config = self + fields_set = pydantic_config.model_fields_set + kwargs = type(pydantic_config).model_fields[dep_field].json_schema_extra + new_param_fn = kwargs.get("new_param_fn", lambda x: x) + param_value = new_param_fn(getattr(pydantic_config, dep_field)) + new_field = kwargs.get("new_param", "") + dep_msg = kwargs.get("deprecated_msg", "") + if dep_field in fields_set: + logger.warning(f"Config parameter {dep_field} is deprecated" + + (f" use {new_field} instead" if new_field else "") + (f". {dep_msg}" if dep_msg else "")) + # Check if there is a new param and if it should be set with a value + if new_field and kwargs.get("set_new_param", True): + # Remove the deprecate field if there is a replacing field + try: + delattr(pydantic_config, dep_field) + except Exception as e: + logger.error(f"Tried removing deprecated '{dep_field}' from config") + raise e + + # Set new param value + new_param_nested = new_field.split(".") + if len(new_param_nested) > 1: + # If the new param exists in a subconfig, we need to get + # the fields set for that subconfig + pydantic_config = reduce(getattr, new_param_nested[:-1], pydantic_config) + fields_set = pydantic_config.model_fields_set + new_param_name = new_param_nested[-1] + assert ( + new_param_name not in fields_set + ), f"Cannot provide deprecated parameter '{dep_field}' and replacing parameter '{new_field}' together" + # A custom function for converting the old param value to new param value can be provided + try: + setattr(pydantic_config, new_param_name, param_value) + except Exception as e: + logger.error(f"Tried setting value for '{new_field}' with value from deprecated '{dep_field}'") + raise e + + def _deprecated_fields_check(self): + fields = type(self).model_fields + for field_name, field_info in fields.items(): + if field_info.json_schema_extra and field_info.json_schema_extra.get("deprecated", False): + self._process_deprecated_field(field_name) + + model_config = ConfigDict( + validate_default=True, + validate_assignment=True, + use_enum_values=True, + populate_by_name=True, + extra="forbid", + arbitrary_types_allowed=True, + protected_namespaces=(), + ) + + @field_serializer("dtype", check_fields=False) + def serialize_torch_dtype(dtype: torch.dtype) -> str: + return str(dtype) + + +def get_config_default(config, field_name): + assert field_name in config.model_fields, f"'{field_name}' is not a field in {config}" + assert not config.model_fields.get( + field_name).is_required(), f"'{field_name}' is a required field and does not have a default value" + return config.model_fields.get(field_name).get_default() + + +class pp_int(int): + """ + A wrapper for integers that will return a custom string or comma-formatted + string of the integer. For example, print(pp_int(1e5)) will return + "10,000". This is useful mainly for auto-generated documentation purposes. + """ + + def __new__(cls, val, custom_print_str=None): + inst = super().__new__(cls, val) + inst.custom_print_str = custom_print_str + return inst + + def __repr__(self): + if self.custom_print_str: + return self.custom_print_str + return f"{self.real:,}" + + +# adapted from https://stackoverflow.com/a/50701137/9201239 +class ScientificNotationEncoder(json.JSONEncoder): + """ + This class overrides ``json.dumps`` default formatter. + + This version keeps everything as normal except formats numbers bigger than 1e3 using scientific notation. + + Just pass ``cls=ScientificNotationEncoder`` to ``json.dumps`` to activate it + + """ + + def iterencode(self, o, _one_shot=False, level=0): + indent = self.indent if self.indent is not None else 4 + prefix_close = " " * level * indent + level += 1 + prefix = " " * level * indent + if isinstance(o, bool): + return "true" if o else "false" + elif isinstance(o, float) or isinstance(o, int): + if o > 1e3: + return f"{o:e}" + else: + return f"{o}" + elif isinstance(o, collections.abc.Mapping): + x = [f'\n{prefix}"{k}": {self.iterencode(v, level=level)}' for k, v in o.items()] + return "{" + ", ".join(x) + f"\n{prefix_close}" + "}" + elif isinstance(o, collections.abc.Sequence) and not isinstance(o, str): + return f"[{ f', '.join(map(self.iterencode, o)) }]" + return "\n, ".join(super().iterencode(o, _one_shot)) + + +class DeepSpeedConfigObject(object): + """ + For json serialization + """ + + def repr(self): + return self.__dict__ + + def __repr__(self): + return json.dumps( + self.__dict__, + sort_keys=True, + indent=4, + cls=ScientificNotationEncoder, + ) + + +def get_scalar_param(param_dict, param_name, param_default_value): + return param_dict.get(param_name, param_default_value) + + +def get_list_param(param_dict, param_name, param_default_value): + return param_dict.get(param_name, param_default_value) + + +def get_dict_param(param_dict, param_name, param_default_value): + return param_dict.get(param_name, param_default_value) + + +def dict_raise_error_on_duplicate_keys(ordered_pairs): + """Reject duplicate keys.""" + d = dict((k, v) for k, v in ordered_pairs) + if len(d) != len(ordered_pairs): + counter = collections.Counter([pair[0] for pair in ordered_pairs]) + keys = [key for key, value in counter.items() if value > 1] + raise ValueError("Duplicate keys in DeepSpeed config: {}".format(keys)) + return d diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..7fe5076e5d9043f1553b4f686d4fb2d39afa32d8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/constants.py @@ -0,0 +1,492 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.accelerator import get_accelerator + +############################################# +# Routes +############################################# +ROUTE_TRAIN = "train" +ROUTE_EVAL = "eval" +ROUTE_PREDICT = "predict" +ROUTE_ENCODE = "encode" + +############################################# +# Batch size +############################################# +TRAIN_BATCH_SIZE = "train_batch_size" +TRAIN_BATCH_SIZE_DEFAULT = None + +############################################# +# Sparse attention +############################################# +SPARSE_ATTENTION = "sparse_attention" +SPARSE_DENSE_MODE = "dense" +SPARSE_FIXED_MODE = "fixed" +SPARSE_VARIABLE_MODE = "variable" +SPARSE_BIGBIRD_MODE = "bigbird" +SPARSE_BSLONGFORMER_MODE = "bslongformer" +SPARSE_MODE = "mode" +SPARSE_MODE_DEFAULT = SPARSE_FIXED_MODE +SPARSE_BLOCK = "block" +SPARSE_BLOCK_DEFAULT = 16 +SPARSE_DIFFERENT_LAYOUT_PER_HEAD = "different_layout_per_head" +SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT = False +SPARSE_NUM_LOCAL_BLOCKS = "num_local_blocks" +SPARSE_NUM_LOCAL_BLOCKS_DEFAULT = 4 +SPARSE_NUM_GLOBAL_BLOCKS = "num_global_blocks" +SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT = 1 +SPARSE_ATTENTION_TYPE = "attention" +SPARSE_ATTENTION_TYPE_DEFAULT = "bidirectional" +SPARSE_HORIZONTAL_GLOBAL_ATTENTION = "horizontal_global_attention" +SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT = False +SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS = "num_different_global_patterns" +SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS_DEFAULT = 1 +SPARSE_NUM_RANDOM_BLOCKS = "num_random_blocks" +SPARSE_NUM_RANDOM_BLOCKS_DEFAULT = 0 +SPARSE_LOCAL_WINDOW_BLOCKS = "local_window_blocks" +SPARSE_LOCAL_WINDOW_BLOCKS_DEFAULT = [4] +SPARSE_GLOBAL_BLOCK_INDICES = "global_block_indices" +SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT = [0] +SPARSE_GLOBAL_BLOCK_END_INDICES = "global_block_end_indices" +SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT = None +SPARSE_NUM_SLIDING_WINDOW_BLOCKS = "num_sliding_window_blocks" +SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT = 3 + +############################################# +# Optimizer and lr scheduler +############################################# +OPTIMIZER = "optimizer" +OPTIMIZER_TYPE_DEFAULT = None +OPTIMIZER_PARAMS = "params" +TYPE = "type" +LEGACY_FUSION = "legacy_fusion" +LEGACY_FUSION_DEFAULT = False +SCHEDULER = "scheduler" +SCHEDULER_TYPE_DEFAULT = None +SCHEDULER_PARAMS = "params" +MAX_GRAD_NORM = 'max_grad_norm' + +############################################# +# Optimizer and lr scheduler +############################################# +ZERO_ALLOW_UNTESTED_OPTIMIZER = "zero_allow_untested_optimizer" +ZERO_ALLOW_UNTESTED_OPTIMIZER_DEFAULT = False +ZERO_FORCE_DS_CPU_OPTIMIZER = "zero_force_ds_cpu_optimizer" +ZERO_FORCE_DS_CPU_OPTIMIZER_DEFAULT = True + +# Steps +STEPS_PER_PRINT = "steps_per_print" +STEPS_PER_PRINT_DEFAULT = None + +######################################### +# Training micro batch size per GPU +######################################### +# Batch size for one training step. This is used when the +# TRAIN_BATCH_SIZE cannot fit in GPU memory to determine +# the number of gradient accumulation steps. By default, this +# is set to None. Users can configure in ds_config.json as below example: +TRAIN_MICRO_BATCH_SIZE_PER_GPU = ''' +TRAIN_MICRO_BATCH_SIZE_PER_GPU is defined in this format: +"train_micro_batch_size_per_gpu": 1 +''' +TRAIN_MICRO_BATCH_SIZE_PER_GPU = "train_micro_batch_size_per_gpu" +TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = None + +######################################### +# Gradient Accumulation +######################################### +# Gradient accumulation feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +GRADIENT_ACCUMULATION_FORMAT = ''' +Gradient Accumulation should be of the format: +"gradient_accumulation_steps": 1 +''' +GRADIENT_ACCUMULATION_STEPS = "gradient_accumulation_steps" +GRADIENT_ACCUMULATION_STEPS_DEFAULT = None + +# DeepSpeed CSR gradient sparsity +SPARSE_GRADIENTS = "sparse_gradients" +SPARSE_GRADIENTS_DEFAULT = False + +######################################### +# BFLOAT16 support +######################################### +# BFLOAT16 feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +BFLOAT16_FORMAT = ''' +BFLOAT16 parameters should be of the format: +"bf16": { + "enabled": true, + "immediate_grad_update": false, + "check_overflow": false +} +''' +BFLOAT16 = "bf16" +BFLOAT16_OLD = "bfloat16" # keeping for backwards compatibility + +BFLOAT16_ENABLED = "enabled" +BFLOAT16_ENABLED_DEFAULT = False + +CHECK_OVERFLOW = "check_overflow" +BFLOAT16_CHECK_OVERFLOW_DEFAULT = False + +# BFLOAT16 optimizer immediate gradient update +BFLOAT16_IMMEDIATE_GRAD_UPDATE = "immediate_grad_update" +BFLOAT16_IMMEDIATE_GRAD_UPDATE_DEFAULT = True + +######################################### +# FP16 support +######################################### +# FP16 feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +FP16_FORMAT = ''' +FP16 parameters should be of the format: +"fp16": { + "enabled": true, + "auto_cast": false, + "loss_scale": 0, + "initial_scale_power": 16, + "loss_scale_window": 1000, + "hysteresis": 2, + "consecutive_hysteresis": false, + "min_loss_scale": 1 +} +''' +FP16 = "fp16" + +FP16_ENABLED = "enabled" +FP16_ENABLED_DEFAULT = False + +# FP16 loss scale, zero means using dynamic scaling +FP16_LOSS_SCALE = "loss_scale" +FP16_LOSS_SCALE_DEFAULT = 0 + +FP16_AUTO_CAST = "auto_cast" +FP16_AUTO_CAST_DEFAULT = False + +# FP16 initial dynamic scale loss power +FP16_INITIAL_SCALE_POWER = "initial_scale_power" +FP16_INITIAL_SCALE_POWER_DEFAULT = 16 + +# FP16 loss scale window +FP16_LOSS_SCALE_WINDOW = "loss_scale_window" +FP16_LOSS_SCALE_WINDOW_DEFAULT = 1000 + +# FP16 hysteresis +FP16_HYSTERESIS = "hysteresis" +FP16_HYSTERESIS_DEFAULT = 2 + +# FP16 consecutive hysteresis +FP16_CONSECUTIVE_HYSTERESIS = "consecutive_hysteresis" +FP16_CONSECUTIVE_HYSTERESIS_DEFAULT = False + +# FP16 min loss scale +FP16_MIN_LOSS_SCALE = "min_loss_scale" +FP16_MIN_LOSS_SCALE_DEFAULT = 1 + +# FP16 master and grads +FP16_MASTER_WEIGHTS_AND_GRADS = "fp16_master_weights_and_grads" +FP16_MASTER_WEIGHTS_AND_GRADS_DEFAULT = False + +######################################### +# Apex AMP support +######################################### +# Use Apex AMP for mixed precision support, all parameters (other than 'enabled') will be passed to +# amp.initialize(model, optimizer, **amp_params) +# See apex documentation for supported parameters/features: https://nvidia.github.io/apex/amp.html#apex.amp.initialize +AMP_FORMAT = ''' +"amp" { + "enabled: true, + "opt_level": "O1", + ... +} +''' +AMP = "amp" + +AMP_ENABLED = "enabled" +AMP_ENABLED_DEFAULT = False + +######################################### +# Torch AMP support +######################################### +TORCH_AUTOCAST_FORMAT = ''' +PyTorch autocast config should be of the format: +"torch_autocast": { + "enabled": true, + "dtype": "bfloat16", + "lower_precision_safe_modules": [ + "torch.nn.modules.linear.Linear", + "torch.nn.modules.conv.Conv2d" + ] +} +''' +TORCH_AUTOCAST = "torch_autocast" + +TORCH_AUTOCAST_ENABLED = "enabled" +TORCH_AUTOCAST_ENABLED_DEFAULT = False +TORCH_AUTOCAST_DTYPE = "dtype" +TORCH_AUTOCAST_LOWER_PRECISION_SAFE_MODULES = "lower_precision_safe_modules" + +######################################### +# Gradient clipping +######################################### +# Gradient clipping. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +GRADIENT_CLIPPING_FORMAT = ''' +Gradient clipping should be enabled as: +"gradient_clipping": 1.0 +''' +GRADIENT_CLIPPING = 'gradient_clipping' +GRADIENT_CLIPPING_DEFAULT = 0. + +######################################### +# Capture graph for short kernels sequences +######################################### +# Graph harvesting. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +GRAPH_HARVESTING_FORMAT = ''' +Graph harvesting should be enabled as: +"graph_harvesting": true +''' +GRAPH_HARVESTING = 'graph_harvesting' +GRAPH_HARVESTING_DEFAULT = False + +######################################### +# Communication data type +######################################### +# Supported types: ['none', 'fp16', 'fp32'] +# By default, this feature is not enabled ('none' value) +# Users can configure in ds_config.json as below example: +COMMUNICATION_DATA_TYPE_FORMAT = ''' +Communication data type should be set as: +"communication_data_type": "fp32" +''' +COMMUNICATION_DATA_TYPE = "communication_data_type" +COMMUNICATION_DATA_TYPE_DEFAULT = None + +########################################################### +# Gradient communication data type for sequence parallelism +########################################################### +# Supported types: ['fp16', 'bf16','fp32'] +# Default value is fp32 +# Users can configure in ds_config.json as below example: +SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_FORMAT = ''' +Optional comm data type for seq paralleism should be set as: +"seq_parallel_communication_data_type": "fp32" +''' +SEQ_PARALLEL_COMMUNICATION_DATA_TYPE = "seq_parallel_communication_data_type" + +if get_accelerator().device_name == 'cuda' and get_accelerator().communication_backend_version() >= (2, 27, 3): + # nccl>=2.27.3 uses fp32 accumulation for half precision inputs, so there is no need to waste compute and memory to manually upcast to fp32 unless the user wants it and then override + SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = None +else: + SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = "fp32" + +SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = "fp32" + +######################################### +# Scale/predivide gradients before allreduce +######################################### +# Prescale gradients. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +PRESCALE_GRADIENTS_FORMAT = ''' +Gradient prescaling should be enabled as: +"prescale_gradients": true +''' +PRESCALE_GRADIENTS = "prescale_gradients" +PRESCALE_GRADIENTS_DEFAULT = False + +GRADIENT_PREDIVIDE_FACTOR_FORMAT = ''' +Gradient predivide factor should be enabled as: +"gradient_predivide_factor": 1.0 +''' +GRADIENT_PREDIVIDE_FACTOR = "gradient_predivide_factor" +GRADIENT_PREDIVIDE_FACTOR_DEFAULT = 1.0 + +######################################### +# Disable AllGather +######################################### +# Disable AllGather. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +DISABLE_ALLGATHER_FORMAT = ''' +Disable AllGather should be enabled as: +"disable_allgather": true +''' +DISABLE_ALLGATHER = "disable_allgather" +DISABLE_ALLGATHER_DEFAULT = False + +######################################### +# Dump DeepSpeed state +######################################### +# Dump State. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +DUMP_STATE_FORMAT = ''' +Dump state should be enabled as: +"dump_state": true +''' +DUMP_STATE = 'dump_state' +DUMP_STATE_DEFAULT = False + +######################################### +# Vocabulary size +######################################### +# Vocabulary size. +# Users can configure in ds_config.json as below example: +VOCABULARY_SIZE_FORMAT = ''' +Vocabulary size can be specified as: +"vocabulary_size": 1024 +''' +VOCABULARY_SIZE = 'vocabulary_size' +VOCABULARY_SIZE_DEFAULT = None + +######################################### +# Wall block breakdown +######################################### +# Wall clock breakdown. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +WALL_CLOCK_BREAKDOWN_FORMAT = ''' +Wall block breakdown should be enabled as: +"wall_clock_breakdown": true +''' +WALL_CLOCK_BREAKDOWN = 'wall_clock_breakdown' +WALL_CLOCK_BREAKDOWN_DEFAULT = False + +MEMORY_BREAKDOWN = 'memory_breakdown' +MEMORY_BREAKDOWN_DEFAULT = False + +######################################### +# Eigenvalue +######################################### +# Eigenvalue computation. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +EIGENVALUE_FORMAT = ''' +Tensorboard can be specified as: +"eigenvalue": { + "enabled": true, + "verbose": true, + "max_iter": 100, + "tol": 1e-2, + "stability": 1e-6 +} +''' +EIGENVALUE = "eigenvalue" + +# Tensorboard enable signal +EIGENVALUE_ENABLED = "enabled" +EIGENVALUE_ENABLED_DEFAULT = False + +EIGENVALUE_VERBOSE = "verbose" +EIGENVALUE_VERBOSE_DEFAULT = False + +EIGENVALUE_MAX_ITER = "max_iter" +EIGENVALUE_MAX_ITER_DEFAULT = 100 + +EIGENVALUE_TOL = "tol" +EIGENVALUE_TOL_DEFAULT = 1e-2 + +EIGENVALUE_STABILITY = "stability" +EIGENVALUE_STABILITY_DEFAULT = 1e-6 + +EIGENVALUE_GAS_BOUNDARY_RESOLUTION = "gas_boundary_resolution" +EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT = 1 + +EIGENVALUE_LAYER_NAME = "layer_name" +EIGENVALUE_LAYER_NAME_DEFAULT = "bert.encoder.layer" + +EIGENVALUE_LAYER_NUM = "layer_num" +EIGENVALUE_LAYER_NUM_DEFAULT = 0 + +######################################### +# Progressive Layer Drop (PLD) +######################################### +PROGRESSIVE_LAYER_DROP = "progressive_layer_drop" + +# PLD enable signal +PLD_ENABLED = "enabled" +PLD_ENABLED_DEFAULT = False + +PLD_THETA = "theta" +PLD_THETA_DEFAULT = 1.0 + +PLD_GAMMA = "gamma" +PLD_GAMMA_DEFAULT = 0.001 + + +######################################### +# Validation modes +######################################### +class ValidationMode: + WARN = "WARN" + IGNORE = "IGNORE" + FAIL = "FAIL" + + +######################################### +# Checkpoint config params +######################################### +# "checkpoint": { +# tag_validation=["Ignore"|"Warn"|"Fail"] +# load_universal=false +# use_node_local_storage=false +# parallel_write: { +# pipeline_stage: [True|False] +# } +# } +CHECKPOINT = "checkpoint" +CHECKPOINT_TAG_VALIDATION = "tag_validation" +CHECKPOINT_TAG_VALIDATION_DEFAULT = ValidationMode.WARN +CHECKPOINT_TAG_VALIDATION_MODES = [ValidationMode.WARN, ValidationMode.IGNORE, ValidationMode.FAIL] + +LOAD_UNIVERSAL_CHECKPOINT = "load_universal" +LOAD_UNIVERSAL_CHECKPOINT_DEFAULT = False + +USE_NODE_LOCAL_STORAGE_CHECKPOINT = "use_node_local_storage" +USE_NODE_LOCAL_STORAGE_CHECKPOINT_DEFAULT = False + +CHECKPOINT_PARALLEL_WRITE = "parallel_write" +CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE = "pipeline_stage" +CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE_DEFAULT = False + +######################################### +# Data types config params +######################################### +# "data_types": { +# grad_accum_dtype=["bf16"|"fp16"|"fp32"] +# } +# } + +DATA_TYPES = "data_types" +GRAD_ACCUM_DTYPE = "grad_accum_dtype" +GRAD_ACCUM_DTYPE_DEFAULT = None + +######################################### +# Drop the last incomplete Batch +# ######################################### +# dataloader_drop_last. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +DATALOADER_DROP_LAST_FORMAT = ''' +The last incomplete batch can be dropped by setting: +"dataloader_drop_last": True +''' +DATALOADER_DROP_LAST = "dataloader_drop_last" +DATALOADER_DROP_LAST_DEFAULT = False + +######################################### +# PIPELINE PARALLELISM +######################################### +PIPE_REPLICATED = 'ds_pipe_replicated' + +######################################### +# DATA PARALLELISM +######################################### +DATA_PARALLEL_GROUP = "data_parallel_group" +GLOBAL_RANK = "global_rank" + +######################################### +# EXPERT-DATA PARALLELISM TOPO Config +######################################### +USE_DATA_BEFORE_EXPERT_PARALLEL = "use_data_before_expert_parallelism" +USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42d17520f9eadaf688bbf529515a763529f34331 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c38643a6e000c21441cac2db44c1fe842112eec Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a430f1619fbc870ec0189e4b0c1cea0dd886338 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7bfa2bfeff8e1416ac25670620e9bf05a320774 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/config.py new file mode 100644 index 0000000000000000000000000000000000000000..690ce97034e452c2dbb451c3f71ebfaa579ad02d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/config.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .constants import * +import copy +from ..config_utils import get_scalar_param + + +# TODO: Reducing config verbosity by returning None or {} when disabled. +# One challenge is that we still need to somehow include the default values, +# for example the *_ENABLED has default of false. +def get_data_efficiency_config(param_dict): + output = {} + output[DATA_EFFICIENCY_ENABLED] = get_data_efficiency_enabled(param_dict) + output[DATA_EFFICIENCY_SEED] = get_data_efficiency_seed(param_dict) + if DATA_EFFICIENCY not in param_dict.keys(): + param_dict[DATA_EFFICIENCY] = {} + sub_param_dict = param_dict[DATA_EFFICIENCY] + output[DATA_SAMPLING] = get_data_sampling(sub_param_dict) + output[DATA_ROUTING] = get_data_routing(sub_param_dict) + return output + + +def get_data_efficiency_enabled(param_dict): + if DATA_EFFICIENCY in param_dict.keys(): + return get_scalar_param(param_dict[DATA_EFFICIENCY], DATA_EFFICIENCY_ENABLED, DATA_EFFICIENCY_ENABLED_DEFAULT) + else: + return False + + +def get_data_efficiency_seed(param_dict): + if DATA_EFFICIENCY in param_dict.keys(): + return get_scalar_param(param_dict[DATA_EFFICIENCY], DATA_EFFICIENCY_SEED, DATA_EFFICIENCY_SEED_DEFAULT) + else: + return DATA_EFFICIENCY_SEED_DEFAULT + + +def get_data_sampling(param_dict): + sub_param_dict = param_dict.get(DATA_SAMPLING, {}) + output = copy.copy(sub_param_dict) + output[DATA_SAMPLING_ENABLED] = get_data_sampling_enabled(param_dict) + output[DATA_SAMPLING_NUM_EPOCHS] = get_data_sampling_num_epochs(param_dict) + output[DATA_SAMPLING_NUM_WORKERS] = get_data_sampling_num_workers(param_dict) + output[DATA_SAMPLING_PIN_MEMORY] = get_data_sampling_pin_memory(param_dict) + output[CURRICULUM_LEARNING] = get_curriculum_learning(sub_param_dict) + output[DYNAMIC_BATCHING] = get_dynamic_batching(sub_param_dict) + return output + + +def get_data_sampling_enabled(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_ENABLED, DATA_SAMPLING_ENABLED_DEFAULT) + else: + return False + + +def get_data_sampling_num_epochs(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_NUM_EPOCHS, DATA_SAMPLING_NUM_EPOCHS_DEFAULT) + else: + return DATA_SAMPLING_NUM_EPOCHS_DEFAULT + + +def get_data_sampling_num_workers(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_NUM_WORKERS, + DATA_SAMPLING_NUM_WORKERS_DEFAULT) + else: + return DATA_SAMPLING_NUM_WORKERS_DEFAULT + + +def get_data_sampling_pin_memory(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_PIN_MEMORY, DATA_SAMPLING_PIN_MEMORY_DEFAULT) + else: + return DATA_SAMPLING_PIN_MEMORY_DEFAULT + + +def get_curriculum_learning(param_dict): + output = {} + output[CURRICULUM_LEARNING_ENABLED] = get_curriculum_learning_enabled(param_dict) + if CURRICULUM_LEARNING not in param_dict.keys(): + param_dict[CURRICULUM_LEARNING] = {} + sub_param_dict = param_dict[CURRICULUM_LEARNING] + if output[CURRICULUM_LEARNING_ENABLED]: + assert CURRICULUM_LEARNING_METRICS in sub_param_dict.keys( + ), f"Curriculum learning is enabled, {CURRICULUM_LEARNING_METRICS} must be specified" + for key, val in get_curriculum_learning_params(param_dict).items(): + output[key] = val + return output + + +def get_dynamic_batching(param_dict): + output = copy.copy(param_dict.get(DYNAMIC_BATCHING, {})) + output[DYNAMIC_BATCHING_ENABLED] = bool(output.get(DYNAMIC_BATCHING_ENABLED, DYNAMIC_BATCHING_ENABLED_DEFAULT)) + output[DYNAMIC_BATCHING_LR_SCALING_METHOD] = str( + output.get(DYNAMIC_BATCHING_LR_SCALING_METHOD, DYNAMIC_BATCHING_LR_SCALING_METHOD_DEFAULT)) + output[DYNAMIC_BATCHING_MIN_BATCH_SIZE] = int( + output.get(DYNAMIC_BATCHING_MIN_BATCH_SIZE, DYNAMIC_BATCHING_MIN_BATCH_SIZE_DEFAULT)) + output[DYNAMIC_BATCHING_MAX_BATCH_SIZE] = int(output[DYNAMIC_BATCHING_MAX_BATCH_SIZE]) \ + if DYNAMIC_BATCHING_MAX_BATCH_SIZE in output.keys() \ + else DYNAMIC_BATCHING_MAX_BATCH_SIZE_DEFAULT + output[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER] = str( + output.get(DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER, DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER_DEFAULT)) + if output[DYNAMIC_BATCHING_ENABLED]: + assert DYNAMIC_BATCHING_MAX_TOKENS in output.keys( + ), f"Dynamic batching is enabled, so {DYNAMIC_BATCHING_MAX_TOKENS} must be specified" + output[DYNAMIC_BATCHING_MAX_TOKENS] = int(output[DYNAMIC_BATCHING_MAX_TOKENS]) + output[DYNAMIC_BATCHING_VERBOSE] = bool(output.get(DYNAMIC_BATCHING_VERBOSE, False)) + return output + + +def get_curriculum_learning_enabled(param_dict): + if CURRICULUM_LEARNING in param_dict.keys(): + return get_scalar_param(param_dict[CURRICULUM_LEARNING], CURRICULUM_LEARNING_ENABLED, + CURRICULUM_LEARNING_ENABLED_DEFAULT) + else: + return False + + +def get_curriculum_learning_params(param_dict): + if CURRICULUM_LEARNING in param_dict.keys(): + curriculum_learning_params = copy.copy(param_dict[CURRICULUM_LEARNING]) + curriculum_learning_params.pop(CURRICULUM_LEARNING_ENABLED) + return curriculum_learning_params + else: + return {} + + +def get_curriculum_enabled_legacy(param_dict): + if CURRICULUM_LEARNING_LEGACY in param_dict.keys(): + return get_scalar_param(param_dict[CURRICULUM_LEARNING_LEGACY], CURRICULUM_ENABLED_LEGACY, + CURRICULUM_ENABLED_DEFAULT_LEGACY) + else: + return False + + +def get_curriculum_params_legacy(param_dict): + if CURRICULUM_LEARNING_LEGACY in param_dict.keys(): + curriculum_params = copy.copy(param_dict[CURRICULUM_LEARNING_LEGACY]) + curriculum_params.pop(CURRICULUM_ENABLED_LEGACY) + return curriculum_params + else: + return False + + +def get_data_routing(param_dict): + output = {} + output[DATA_ROUTING_ENABLED] = get_data_routing_enabled(param_dict) + if DATA_ROUTING not in param_dict.keys(): + param_dict[DATA_ROUTING] = {} + sub_param_dict = param_dict[DATA_ROUTING] + output[RANDOM_LTD] = get_random_ltd(sub_param_dict) + + return output + + +def get_data_routing_enabled(param_dict): + if DATA_ROUTING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_ROUTING], DATA_ROUTING_ENABLED, DATA_ROUTING_ENABLED_DEFAULT) + else: + return False + + +def get_random_ltd(param_dict): + output = {} + output[RANDOM_LTD_ENABLED] = RANDOM_LTD_ENABLED_DEFAULT + output[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE] = {} + output[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][ + RANDOM_LTD_LAYER_TOKEN_LR_ENABLED] = RANDOM_LTD_LAYER_TOKEN_LR_ENABLED_DEFAULT + if get_random_ltd_enabled(param_dict): + output[RANDOM_LTD_ENABLED] = get_random_ltd_enabled(param_dict) + for key, val in get_random_ltd_params(param_dict).items(): + output[key] = val + return output + + +def get_random_ltd_enabled(param_dict): + if RANDOM_LTD in param_dict.keys(): + return get_scalar_param(param_dict[RANDOM_LTD], RANDOM_LTD_ENABLED, RANDOM_LTD_ENABLED_DEFAULT) + else: + return False + + +def get_random_ltd_params(param_dict): + if RANDOM_LTD in param_dict.keys(): + random_ltd_params = copy.copy(param_dict[RANDOM_LTD]) + random_ltd_params.pop(RANDOM_LTD_ENABLED) + return random_ltd_params + else: + return {} diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..73cc69c1f606a3c77e6969cc32c57ef7cd8e5111 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/constants.py @@ -0,0 +1,136 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Data efficiency library + See sample config at https://www.deepspeed.ai/docs/config-json/data-efficiency +""" +DATA_EFFICIENCY = "data_efficiency" +DATA_EFFICIENCY_ENABLED = "enabled" +DATA_EFFICIENCY_ENABLED_DEFAULT = False +DATA_EFFICIENCY_SEED = "seed" +DATA_EFFICIENCY_SEED_DEFAULT = 1234 + +######################################### +# Data efficiency - Data Sampling +######################################### +DATA_SAMPLING = "data_sampling" +DATA_SAMPLING_ENABLED = "enabled" +DATA_SAMPLING_ENABLED_DEFAULT = False +DATA_SAMPLING_NUM_EPOCHS = "num_epochs" +DATA_SAMPLING_NUM_EPOCHS_DEFAULT = 1000 +DATA_SAMPLING_NUM_WORKERS = "num_workers" +DATA_SAMPLING_NUM_WORKERS_DEFAULT = 0 +DATA_SAMPLING_PIN_MEMORY = "pin_memory" +DATA_SAMPLING_PIN_MEMORY_DEFAULT = False + +######################################### +# Data efficiency - Data Sampling - Curriculum Learning +######################################### +CURRICULUM_LEARNING = "curriculum_learning" +CURRICULUM_LEARNING_ENABLED = "enabled" +CURRICULUM_LEARNING_ENABLED_DEFAULT = False +CURRICULUM_LEARNING_CLUSTER_PATH = "data_cluster_path" +CURRICULUM_LEARNING_METRICS = "curriculum_metrics" +CURRICULUM_LEARNING_SAMPLE_PATH = "index_to_sample_path" +CURRICULUM_LEARNING_METRIC_PATH = "index_to_metric_path" +CURRICULUM_LEARNING_CLUSTERING_TYPE = "clustering_type" +CURRICULUM_LEARNING_SINGLE_CLUSTER = "single_cluster" +CURRICULUM_LEARNING_CLUSTER_PREFIX = "cluster" +CURRICULUM_LEARNING_DIFFICULTY_TYPE = "difficulty_type" +CURRICULUM_LEARNING_VALUE_BASED = "value" +CURRICULUM_LEARNING_PERCENTILE_BASED = "percentile" +CURRICULUM_LEARNING_MIN_DIFFICULTY = "min_difficulty" +CURRICULUM_LEARNING_MAX_DIFFICULTY = "max_difficulty" +CURRICULUM_LEARNING_SCHEDULE_TYPE = "schedule_type" +CURRICULUM_LEARNING_SCHEDULE_CONFIG = "schedule_config" +CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY = "difficulty" +CURRICULUM_LEARNING_SCHEDULE_MAX_STEP = "max_step" +CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP = "total_curriculum_step" +CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP = "difficulty_step" +CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE = "root_degree" +CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE = "fixed_discrete" +CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT = "fixed_root" +CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR = "fixed_linear" +CURRICULUM_LEARNING_SCHEDULE_CUSTOM = "custom" +CURRICULUM_LEARNING_CURRENT_DIFFICULTY = "current_difficulty" + +CURRICULUM_LEARNING_BATCH = "batch" +CURRICULUM_LEARNING_CONSUMED_SAMPLES = "consumed_samples" +CURRICULUM_LEARNING_STEP = "curriculum_step" +CURRICULUM_LEARNING_CURRENT_DIFFICULTIES = "current_difficulties" +CURRICULUM_LEARNING_DATA_CLUSTER_PATHS = "data_cluster_paths" +CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION = "data_cluster_current_position" +CURRICULUM_LEARNING_NP_RNG_STATE = "np_rng_state" + +######################################### +# Data efficiency - Dynamic batching and LR scaling +######################################### +DYNAMIC_BATCHING = "dynamic_batching" +DYNAMIC_BATCHING_ENABLED = "enabled" +DYNAMIC_BATCHING_ENABLED_DEFAULT = False +DYNAMIC_BATCHING_METRICS_PATH = "metrics_path" +DYNAMIC_BATCHING_LR_SCALING_METHOD = "lr_scaling_method" # "linear" / "sqrt" / "none" +DYNAMIC_BATCHING_LR_SCALING_METHOD_DEFAULT = "linear" +DYNAMIC_BATCHING_MIN_BATCH_SIZE = "min_batch_size" +DYNAMIC_BATCHING_MIN_BATCH_SIZE_DEFAULT = 1 +DYNAMIC_BATCHING_MAX_BATCH_SIZE = "max_batch_size" +DYNAMIC_BATCHING_MAX_BATCH_SIZE_DEFAULT = None +DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER = "sequence_picking_order" # "random" / "seqlen" / "dataloader" +DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER_DEFAULT = "dataloader" # "random" / "seqlen" / "dataloader" +DYNAMIC_BATCHING_MAX_TOKENS = "max_tokens" +DYNAMIC_BATCHING_VERBOSE = "verbose" + +######################################### +# Curriculum Learning legacy implementation +######################################### +CURRICULUM_LEARNING_LEGACY = "curriculum_learning" + +CURRICULUM_ENABLED_LEGACY = "enabled" +CURRICULUM_ENABLED_DEFAULT_LEGACY = False + +######################################### +# Data efficiency - Data Routing +######################################### +DATA_ROUTING = "data_routing" +DATA_ROUTING_ENABLED = "enabled" +DATA_ROUTING_ENABLED_DEFAULT = False + +######################################### +# Data efficiency - Data Routing - Random LTD +######################################### +RANDOM_LTD = "random_ltd" +RANDOM_LTD_ENABLED = "enabled" +RANDOM_LTD_ENABLED_DEFAULT = False + +RANDOM_LTD_MODEL_MASK_NAME = "model_mask_name" +RANDOM_LTD_MODEL_TYPE = "model_type" +RANDOM_LTD_MICRO_BATCH_SIZE = "micro_batch_size" +RANDOM_LTD_GLOBAL_BATCH_SIZE = "global_batch_size" +RANDOM_LTD_SAMPLE_INDEX = "sample_idx" +RANDOM_LTD_ATTENTION_MASK = "attention_mask" +RANDOM_LTD_HIDDEN_STATE_ORDER = "hidden_state_order" +RANDOM_LTD_LAYER_NUM = "random_ltd_layer_num" +RANDOM_LTD_LAYER_ID = "random_ltd_layer_id" +RANDOM_LTD_TOTAL_LAYER_NUM = "total_layer_num" +RANDOM_LTD_CONSUMED_LAYER_TOKENS = "consumed_layer_tokens" + +# scheduler +RANDOM_LTD_SCHEDULER = "random_ltd_schedule" +RANDOM_LTD_MAX_VALUE = "max_value" +RANDOM_LTD_MIN_VALUE = "min_value" +RANDOM_LTD_CURRENT_VALUE = "current_value" +RANDOM_LTD_SCHEDULE_CONFIG = "schedule_config" +RANDOM_LTD_INCREASE_STEP = "seq_per_step" +RANDOM_LTD_REQUIRE_STEP = "require_steps" +RANDOM_LTD_SCHEDULER_TYPE = "schedule_type" +RANDOM_LTD_CURR_STEP = "current_steps" + +# learning rate schedulers +RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE = "layer_token_lr_schedule" +RANDOM_LTD_LAYER_TOKEN_LR_ENABLED = "enabled" +RANDOM_LTD_LAYER_TOKEN_LR_ENABLED_DEFAULT = False +RANDOM_LTD_TOTAL_LAYER_TOKENS = "total_layer_tokens" +RANDOM_LTD_WARMUP_TYPE = "warmup_type" +RANDOM_LTD_WARMUP_LAYER_TOKENS = "warmup_layer_tokens" diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..23d747957dc4647e06fad0a94e5e4b071b6f6e23 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +from deepspeed.utils import logger +from .constants import * + + +class CurriculumScheduler(object): + + def __init__(self, config): + super().__init__() + self.state = {} + assert CURRICULUM_LEARNING_MIN_DIFFICULTY in config, \ + f"Curriculum learning requires the config '{CURRICULUM_LEARNING_MIN_DIFFICULTY}'" + assert CURRICULUM_LEARNING_MAX_DIFFICULTY in config, \ + f"Curriculum learning requires the config '{CURRICULUM_LEARNING_MAX_DIFFICULTY}'" + assert CURRICULUM_LEARNING_SCHEDULE_TYPE in config, \ + f"Curriculum learning requires the config '{CURRICULUM_LEARNING_SCHEDULE_TYPE}'" + self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY] = config[CURRICULUM_LEARNING_MIN_DIFFICULTY] + self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] = config[CURRICULUM_LEARNING_MAX_DIFFICULTY] + self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = config[CURRICULUM_LEARNING_MIN_DIFFICULTY] + self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] = config[CURRICULUM_LEARNING_SCHEDULE_TYPE] + self.first_step = True + if config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE: + """ + The schedule_config is a list of difficulty and a list of max + step belonging to each difficulty. Example json config: + "schedule_config": { + "difficulty": [1,2,3], + "max_step": [5,10] + } + The "max_step" has one less element than "difficulty", because + the last difficulty will be used for all following steps. + The self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] is a dictionary of + difficulty : [max step for this difficulty, next difficulty]. + """ + assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_discrete schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY}'" + assert CURRICULUM_LEARNING_SCHEDULE_MAX_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_discrete schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_MAX_STEP}'" + assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_MAX_STEP]) > 0 + assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY]) > 0 + assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY]) == len( + config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_MAX_STEP]) + 1 + self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT: + """ + The schedule_config includes: + total_curriculum_step: how many steps the curriculum learning takes to go + from min difficulty to max difficulty. + difficulty_step: the difficulty level determined every time must + be a multiple of this difficulty_step. This is used to determine + the step of difficulty increase, and to ensure the use of NVIDIA + Tensor Core acceleration (requires multiple of 8 (FP16) or + 16 (INT8)). + root_degree: the degree of the root function. Degree of 2 means + square root and degree of 3 means cube root. Degree of 1 is + equivalent to linear. + "schedule_config": { + "total_curriculum_step": 30000, + "difficulty_step": 8, + "root_degree": 2 + } + """ + assert CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP}'" + assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP}'" + assert CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE}'" + if config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP] % 8 != 0: + logger.warning( + f'When using seqlen metric, the difficulty_step for curriculum learning has to be multiple of 8 (for FP16 data) or 16 (for INT8 data) to enable NVIDIA Tensor Core acceleration. Disregard this warning if this is unrelated to your metric/hardware.' + ) + self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR: + """ + The schedule_config is the same as CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT but without the + root_degree. + "schedule_config": { + "total_curriculum_step": 30000, + "difficulty_step": 8 + } + """ + assert CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_linear schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP}'" + assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_linear schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP}'" + if config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP] % 8 != 0: + logger.warning( + f'When using seqlen metric, the difficulty_step for curriculum learning has to be multiple of 8 (for FP16 data) or 16 (for INT8 data) to enable NVIDIA Tensor Core acceleration. Disregard this warning if this is unrelated to your metric/hardware.' + ) + self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_CUSTOM: + """ + Fully customized schedule. User need to provide a custom schedule + function by using the set_custom_curriculum_learning_schedule API + in deepspeed/runtime/engine.py + """ + self.custom_get_difficulty = None + else: + raise RuntimeError('Unsupported curriculum schedule type') + + def get_current_difficulty(self): + return self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] + + def set_current_difficulty(self, difficulty): + self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = difficulty + + def set_custom_get_difficulty(self, schedule_function): + self.custom_get_difficulty = schedule_function + + def get_state(self): + return self.state + + def set_state(self, state): + self.state = state + + def __fixed_discrete_get_difficulty(self, global_steps): + s_state = self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + if global_steps > s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP][-1]: + return s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY][-1] + for i in range(len(s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP])): + if global_steps <= s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP][i]: + return s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY][i] + + def __fixed_root_get_difficulty(self, global_steps, root_degree=None): + s_state = self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + if root_degree is None: + root_degree = s_state[CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE] + next_difficulty = (float(global_steps) / s_state[CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP])**(1.0 / root_degree) + next_difficulty = math.floor( + next_difficulty * + (self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] - self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) + + self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) + next_difficulty -= (next_difficulty % s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP]) + next_difficulty = min(next_difficulty, self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY]) + return next_difficulty + + def get_difficulty(self, global_steps): + if self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE: + return self.__fixed_discrete_get_difficulty(global_steps) + elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR: + return self.__fixed_root_get_difficulty(global_steps, 1) + elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT: + return self.__fixed_root_get_difficulty(global_steps) + elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_CUSTOM: + return self.custom_get_difficulty(global_steps) + else: + raise RuntimeError('Unsupported curriculum schedule type') + + def update_difficulty(self, global_steps): + if self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] < self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY]: + self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = self.get_difficulty(global_steps) + return self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b9c1c5c49c7dc7d59762c4c08f4b5c4171e75d9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..164053097aff638211a00394303d9a447dc5f616 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d728dbbee2d39cd0ccb30aba76a2b4bfe8136cb Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47555803744b1c39fa8580eef887f7070fc68374 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f64c99133a047df7c11c034240b480a250ee7d4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..780a88c3d5a2fb016f2c59bbe3a88fbaff6499b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.utils import logger +from torch import Tensor +from torch.nn import Module +from ..constants import * +from deepspeed.ops.random_ltd.dropping_utils import gpt_sample_tokens, bert_sample_tokens, GatherTokens, ScatterTokens + + +#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586 +class RandomLayerTokenDrop(Module): + """ + A layer wrapper for random LTD + """ + + def __init__(self, layer: Module): + super(RandomLayerTokenDrop, self).__init__() + self.random_ltd_layer = layer + self.reserved_length = None #config['max_value'] + self.random_ltd_scheduler = None + self.max_length = None + self.reserved_length = -1 + self.curr_seq = -1 + self.batch_first = False + + def init_config(self, config, scheduler, random_ltd_layer_id): + self.random_ltd_scheduler = scheduler + self.random_ltd_layer_id = random_ltd_layer_id + self.max_length = self.random_ltd_scheduler.state[RANDOM_LTD_MAX_VALUE] + + self.mask_name = config[RANDOM_LTD_MODEL_MASK_NAME] + self.micro_bs = config[RANDOM_LTD_MICRO_BATCH_SIZE] + self.random_ltd_num_layer = self.random_ltd_scheduler.random_ltd_layer_num + hs_order = config[RANDOM_LTD_HIDDEN_STATE_ORDER] + self.model_type = config[RANDOM_LTD_MODEL_TYPE] + + if hs_order == 'batch_seq_dim': + self.get_hidden_tensor_shape = self.get_bsh + self.batch_first = True + elif hs_order == 'seq_batch_dim': + self.get_hidden_tensor_shape = self.get_sbh + self.batch_first = False + else: + logger.warning( + "************For now, we only support batch_seq_dim or seq_batch_dim inputs. You can easily \ + your own input dimension orders************") + raise NotImplementedError + + if self.model_type == 'encoder': + self.index_generator = bert_sample_tokens + elif self.model_type == 'decoder': + self.index_generator = gpt_sample_tokens + else: + logger.warning("************For now, we only support encoder-only or decoder-only models************") + raise NotImplementedError + + def get_bsh(self, hidden_stats): + self.curr_seq, self.curr_micro_batch = hidden_stats.size()[1], hidden_stats.size()[0] + + def get_sbh(self, hidden_stats): + self.curr_seq, self.curr_micro_batch = hidden_stats.size()[0], hidden_stats.size()[1] + + def forward(self, hidden_states, **kwargs) -> Tensor: + if self.random_ltd_scheduler is not None: + self.reserved_length = self.random_ltd_scheduler.get_current_seq() + self.get_hidden_tensor_shape(hidden_states) + if self.training and self.random_ltd_scheduler is not None and self.reserved_length < self.curr_seq: + if self.mask_name is not None: + mask = kwargs[self.mask_name] + else: + mask = None + if self.random_ltd_layer_id == 0: + sampled_indices, part_attention_mask = self.index_generator(self.reserved_length,\ + self.curr_seq, \ + self.curr_micro_batch, \ + self.random_ltd_num_layer, \ + hidden_states.device, mask) + self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] = sampled_indices + self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] = part_attention_mask + else: + sampled_indices = self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] + part_attention_mask = self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] + + hidden_states, part_hidden_states = GatherTokens.apply(hidden_states, + sampled_indices[self.random_ltd_layer_id, :, :], + self.batch_first) + if self.mask_name is not None: + if self.model_type == 'encoder': + kwargs[self.mask_name] = part_attention_mask[self.random_ltd_layer_id] + else: + kwargs[self.mask_name] = part_attention_mask + + outputs = self.random_ltd_layer(part_hidden_states, **kwargs) + + if isinstance(outputs, tuple): + hidden_states = ScatterTokens.apply(hidden_states, outputs[0], + sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first) + my_list = list(outputs) + my_list[0] = hidden_states + return tuple(my_list) + elif isinstance(outputs, Tensor): + hidden_states = ScatterTokens.apply(hidden_states, outputs, + sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first) + return hidden_states + else: + logger.warning("************For now, we only support tuple and tensor output. \ + You need to adjust the output according to the layer in your model************") + raise NotImplementedError + else: + return self.random_ltd_layer(hidden_states, **kwargs) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py new file mode 100644 index 0000000000000000000000000000000000000000..150182d77bcfda20b1aa1aabd4b8785542ca9d1b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .basic_layer import RandomLayerTokenDrop +from collections import OrderedDict +from deepspeed.compression.helper import recursive_getattr, recursive_setattr + + +def convert_to_random_ltd(model, convert_type): + if hasattr(model, 'module'): + c_model = model.module + else: + c_model = model + + for name, module in c_model.named_modules(): + + if isinstance(module, convert_type): + old_module = recursive_getattr(c_model, name) + new_module = RandomLayerTokenDrop(old_module) + recursive_setattr(c_model, name, new_module) + + model.random_ltd_initialize() + return model + + +def save_without_random_ltd(model): + if hasattr(model, 'module'): + c_model = model.module + else: + c_model = model + + model_dic = c_model.state_dict() + return remove_random_ltd_state_dict(model_dic) + + +def remove_random_ltd_state_dict(state_dict): + new_state_dict = OrderedDict() + for key, value in state_dict.items(): + if '.random_ltd_layer' in key: + new_key = ''.join(key.split('.random_ltd_layer')) + else: + new_key = key + new_state_dict[new_key] = value + return new_state_dict diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..8a0b031d4f633976f438d5151973e58afa77712e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py @@ -0,0 +1,107 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math + +from deepspeed.utils import logger +# from deepspeed.runtime.lr_schedules import WarmupLR +from ..constants import * + +#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586 + + +class BaseScheduler(object): + + def __init__(self): + self.state = {} + + def __fixed_root_get_value(self, global_steps, root_degree=None): + s_state = self.state[RANDOM_LTD_SCHEDULE_CONFIG] + if root_degree is None: + root_degree = s_state['root_degree'] + next_seq = (float(global_steps) / s_state[RANDOM_LTD_REQUIRE_STEP])**(1.0 / root_degree) + next_seq = math.floor(next_seq * (self.state[RANDOM_LTD_MAX_VALUE] - self.state[RANDOM_LTD_MIN_VALUE]) + + self.state[RANDOM_LTD_MIN_VALUE]) + next_seq -= (next_seq % s_state[RANDOM_LTD_INCREASE_STEP]) + next_seq = min(next_seq, self.state[RANDOM_LTD_MAX_VALUE]) + return next_seq + + def get_value(self, global_steps): + if self.state[RANDOM_LTD_SCHEDULER_TYPE] == 'fixed_linear': + return self.__fixed_root_get_value(global_steps, 1) + else: + raise RuntimeError('Unsupported random LTD schedule type') + + +class RandomLTDScheduler(BaseScheduler): + + def __init__(self, config): + super().__init__() + self.model_layer_num = config[RANDOM_LTD_TOTAL_LAYER_NUM] + self.random_ltd_layer_num = config[RANDOM_LTD_LAYER_NUM] + self.config_schedule = config[RANDOM_LTD_SCHEDULER] + self.global_batch_size = config[RANDOM_LTD_GLOBAL_BATCH_SIZE] + self.reset_to_init() + + if config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: + logger.warning("**********Work In Progress************") + raise NotImplementedError + + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0 + + # self.first_step = True + def get_total_layer_tokens(self, train_iters): + for step in range(train_iters): + self.update_seq(step) + return self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] + + def reset_to_init(self): + if self.config_schedule is not None: + self.state[RANDOM_LTD_MIN_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE] + self.state[RANDOM_LTD_MAX_VALUE] = self.config_schedule[RANDOM_LTD_MAX_VALUE] + self.state[RANDOM_LTD_CURRENT_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE] + self.state[RANDOM_LTD_SCHEDULE_CONFIG] = self.config_schedule[RANDOM_LTD_SCHEDULE_CONFIG] + self.state[RANDOM_LTD_SCHEDULER_TYPE] = self.config_schedule[RANDOM_LTD_SCHEDULER_TYPE] + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0 + self.state[RANDOM_LTD_CURR_STEP] = -1 + + def get_current_seq(self): + return self.state[RANDOM_LTD_CURRENT_VALUE] + + def set_current_seq(self, seq_length): + self.state[RANDOM_LTD_CURRENT_VALUE] = seq_length + + def get_random_ltd_layer_num(self): + return self.random_ltd_layer_num + + def get_state(self): + return self.state + + def set_state(self, state): + self.state = state + + def update_seq(self, global_steps): + if self.state[RANDOM_LTD_CURRENT_VALUE] < self.state[RANDOM_LTD_MAX_VALUE]: + self.state[RANDOM_LTD_CURRENT_VALUE] = self.get_value(global_steps) + if global_steps != self.state[RANDOM_LTD_CURR_STEP]: + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] += self.global_batch_size*(self.state[RANDOM_LTD_CURRENT_VALUE] * self.random_ltd_layer_num \ + + self.state[RANDOM_LTD_MAX_VALUE] * (self.model_layer_num - self.random_ltd_layer_num)) + self.state[RANDOM_LTD_CURR_STEP] = global_steps + + def state_dict(self): + return { + RANDOM_LTD_CONSUMED_LAYER_TOKENS: self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS], + RANDOM_LTD_CURR_STEP: self.state[RANDOM_LTD_CURR_STEP], + RANDOM_LTD_CURRENT_VALUE: self.state[RANDOM_LTD_CURRENT_VALUE], + RANDOM_LTD_MIN_VALUE: self.state[RANDOM_LTD_MIN_VALUE], + RANDOM_LTD_MAX_VALUE: self.state[RANDOM_LTD_MAX_VALUE], + } + + def load_state_dict(self, state_dict): + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = state_dict[RANDOM_LTD_CONSUMED_LAYER_TOKENS] + self.state[RANDOM_LTD_CURR_STEP] = state_dict[RANDOM_LTD_CURR_STEP] + self.state[RANDOM_LTD_CURRENT_VALUE] = state_dict[RANDOM_LTD_CURRENT_VALUE] + self.state[RANDOM_LTD_MIN_VALUE] = state_dict[RANDOM_LTD_MIN_VALUE] + self.state[RANDOM_LTD_MAX_VALUE] = state_dict[RANDOM_LTD_MAX_VALUE] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..afcfef2ef4dc04279b7fa3ee49916519f33ffbec --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + + +def bsh_decoder_gather(reserved_length, hidden_states, mask): + # random-layer-token-drop + rand_list = [] + part_hidden_states = [] # batch, seq, hidden ## different from megatron + for k in range(hidden_states.size(0)): + B_tmp = torch.randperm(hidden_states.size(1), device=hidden_states.device)[:reserved_length] + B = B_tmp.sort()[0] + rand_list.append(B) + part_hidden_states.append(hidden_states[k:k + 1, B, :]) + + part_hidden_states = torch.cat(part_hidden_states, dim=0) + part_mask = mask[:, :, :reserved_length, :reserved_length] + return part_hidden_states, rand_list, part_mask + + +def bsh_decoder_scatter(hidden_states, part_hidden_states, rand_list): + for k in range(hidden_states.size(0)): + hidden_states[k, rand_list[k], :] = part_hidden_states[k, :, :] + return hidden_states diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4047518331986d5b489904158578900a000246a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ee54cd1ec6962063e1ec303c0149e1d62c8d4d9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50356307a566ce38fa7dec96be530f5af22dbca5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8e04aa78bd52f3a9e6fc1a78100fe14533b1aa4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa34ee0ebb20b74193ea439f699b773fbd322066 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd86079e7a4a286ba58496256000d75d4a1af2e2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py new file mode 100644 index 0000000000000000000000000000000000000000..f82c684ec6e229597736831d7ea56ca8c3bdbaa0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py @@ -0,0 +1,885 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import sys +from collections import defaultdict +import csv +import time +from multiprocessing import Process, Manager +import numpy as np +import torch +from torch.utils.data import BatchSampler, SequentialSampler, DataLoader, Subset + +import deepspeed.comm as dist +from deepspeed.utils import logger +from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset, valid_dtypes +from deepspeed.runtime.data_pipeline.data_sampling.utils import split_dataset, split_index, create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype + + +class DataAnalyzer(object): + + def __init__(self, + dataset, + num_workers=1, + worker_id=0, + num_threads=1, + num_threads_reduce=1, + specific_threads=[], + batch_size=1, + metric_names=[], + metric_functions=[], + metric_types=[], + metric_dtypes=[], + save_path="./", + collate_fn=None, + custom_map_init=None, + custom_map_update=None, + custom_map_finalize=None, + custom_reduce=None, + sample_indices=None): + super().__init__() + self.dataset = dataset + self.num_workers = num_workers + self.worker_id = worker_id + self.num_threads = num_threads + self.num_threads_reduce = num_threads_reduce + self.specific_threads = specific_threads + self.batch_size = batch_size + self.metric_names = metric_names + self.metric_functions = metric_functions + self.metric_types = metric_types + self.metric_dtypes = metric_dtypes + self.save_path = save_path + self.collate_fn = collate_fn + self.custom_map_init = custom_map_init + self.custom_map_update = custom_map_update + self.custom_map_finalize = custom_map_finalize + self.custom_reduce = custom_reduce + self.sample_indices = sample_indices + + def init_metric_results(self, thread_id, metric_names, metric_types, metric_dtypes, save_path, worker_id): + metric_results = [] + for m_idx in range(len(metric_names)): + metric_name, metric_type, metric_dtype = metric_names[m_idx], \ + metric_types[m_idx], metric_dtypes[m_idx] + assert metric_dtype in valid_dtypes, f"metric_dtype {metric_dtype} not supported. Supported dtypes {valid_dtypes}" + metric_save_path = f"{save_path}/{metric_name}/worker{worker_id}_thread{thread_id}/" + os.makedirs(metric_save_path, exist_ok=True) + if metric_type == 'single_value_per_sample': + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" + sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_dtype) + metric_to_sample_fname = f"{metric_save_path}/{metric_name}_metric_to_sample" + os.system(f"rm -rf {metric_to_sample_fname}*") + metric_to_sample_dict = defaultdict(list) + metric_results.append({ + "sample_to_metric_fname": sample_to_metric_fname, + "sample_to_metric_builder": sample_to_metric_builder, + "metric_to_sample_fname": metric_to_sample_fname, + "metric_to_sample_dict": metric_to_sample_dict + }) + elif metric_type == 'accumulate_value_over_samples': + metric_value = None + metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" + metric_results.append({"metric_value": metric_value, "metric_value_fname": metric_value_fname}) + return metric_results + + def update_metric_results(self, + data, + metric_types, + metric_dtypes, + metric_functions, + metric_results, + batch_start_idx=0): + for m_idx in range(len(metric_types)): + metric_type, metric_dtype, metric_function, metric_result = metric_types[m_idx], \ + metric_dtypes[m_idx], metric_functions[m_idx], metric_results[m_idx] + metric_values = metric_function(data) + + assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \ + "metric_function must return a tensor or array" + assert metric_values.dtype == metric_dtype, \ + f"metric_function result dtype {metric_values.dtype} does not match metric_dtype {metric_dtype}" + if isinstance(metric_values, np.ndarray): + metric_values = torch.from_numpy(metric_values) + + if metric_type == 'single_value_per_sample': + for row in range(metric_values.size()[0]): + sample_idx = batch_start_idx + row # sample idx following dataset iteration order + if isinstance(data, dict) and 'index' in data: # Megatron use case, idx provided in 'index' field + sample_idx = data['index'][row][0].item() + elif self.sample_indices is not None: # user defined shuffling of indices + sample_idx = self.sample_indices[sample_idx] + metric_result["sample_to_metric_builder"].add_item(metric_values[row].reshape(-1)) + metric_result["metric_to_sample_dict"][metric_values[row].item()].append(sample_idx) + for m_value in metric_result["metric_to_sample_dict"]: + if len(metric_result["metric_to_sample_dict"][m_value]) > 100: + metric_fname = metric_result["metric_to_sample_fname"] + with open(f"{metric_fname}_{m_value}.csv", 'a') as f: + writer = csv.writer(f) + writer.writerows([metric_result["metric_to_sample_dict"][m_value]]) + metric_result["metric_to_sample_dict"][m_value] = [] + elif metric_type == 'accumulate_value_over_samples': + if metric_result["metric_value"] is None: + metric_result["metric_value"] = metric_values + else: + metric_result["metric_value"].add_(metric_values) + + def finalize_metric_results(self, metric_types, metric_dtypes, metric_results): + for m_idx in range(len(metric_types)): + metric_type, metric_dtype, metric_result = metric_types[m_idx], \ + metric_dtypes[m_idx], metric_results[m_idx] + if metric_type == 'single_value_per_sample': + metric_fname = metric_result["sample_to_metric_fname"] + close_mmap_dataset_builder(metric_result["sample_to_metric_builder"], metric_fname) + for m_value in metric_result["metric_to_sample_dict"]: + if len(metric_result["metric_to_sample_dict"][m_value]) > 0: + metric_fname = metric_result["metric_to_sample_fname"] + with open(f"{metric_fname}_{m_value}.csv", 'a') as f: + writer = csv.writer(f) + writer.writerows([metric_result["metric_to_sample_dict"][m_value]]) + metric_result["metric_to_sample_dict"][m_value] = [] + elif metric_type == 'accumulate_value_over_samples': + if metric_result["metric_value"] is not None: + metric_value_builder = create_mmap_dataset_builder(metric_result["metric_value_fname"], + metric_dtype) + metric_value_builder.add_item(metric_result["metric_value"].reshape(-1)) + close_mmap_dataset_builder(metric_value_builder, metric_result["metric_value_fname"]) + + def run_map_helper(self, thread_id): + start_idx, end_idx = self.thread_splits[thread_id][0], \ + self.thread_splits[thread_id][1] + logger.info(f"worker {self.worker_id} thread {thread_id}: start working " \ + f"on data subset {start_idx} to {end_idx}") + thread_dataset = Subset(self.dataset, list(range(start_idx, end_idx))) + sampler = BatchSampler(SequentialSampler(thread_dataset), batch_size=self.batch_size, drop_last=False) + iterator = iter( + DataLoader(thread_dataset, + batch_sampler=sampler, + num_workers=0, + collate_fn=self.collate_fn, + pin_memory=False)) + if self.custom_map_init is None: + metric_results = self.init_metric_results(thread_id, self.metric_names, self.metric_types, + self.metric_dtypes, self.save_path, self.worker_id) + else: + metric_results = self.custom_map_init(thread_id, self.metric_names, self.metric_types, self.metric_dtypes, + self.save_path, self.worker_id) + total_sample = len(thread_dataset) + processed_sample = 0 + start = time.time() + while True: + try: + data = next(iterator) + batch_start_idx = start_idx + processed_sample + if self.custom_map_update is None: + self.update_metric_results(data, self.metric_types, self.metric_dtypes, self.metric_functions, + metric_results, batch_start_idx) + else: + self.custom_map_update(data, self.metric_types, self.metric_dtypes, self.metric_functions, + metric_results, batch_start_idx) + processed_sample += len(data) + duration = (time.time() - start) / 3600.0 + remain_duration = duration * total_sample / processed_sample - duration + logger.info( + f"worker {self.worker_id} thread {thread_id}: {processed_sample} " \ + f"out of {total_sample} processed in {duration:.2f} hr, " \ + f"estimated to finish in {remain_duration:.2f} hr") + except StopIteration: + logger.info(f"worker {self.worker_id} thread {thread_id}: reach end of file") + break + if self.custom_map_finalize is None: + self.finalize_metric_results(self.metric_types, self.metric_dtypes, metric_results) + else: + self.custom_map_finalize(self.metric_types, self.metric_dtypes, metric_results) + logger.info(f"worker {self.worker_id} thread {thread_id}: finished") + + def run_map(self): + self.worker_splits, self.thread_splits = split_dataset(self.dataset, self.num_workers, self.worker_id, + self.num_threads) + if len(self.specific_threads) > 0: + threads_to_run = self.specific_threads + else: + threads_to_run = list(range(self.num_threads)) + if self.num_threads > 1: + p = [] + for thread in threads_to_run: + p.append(Process(target=self.run_map_helper, args=(thread, ))) + p[thread].start() + + for thread in threads_to_run: + p[thread].join() + else: + assert self.num_threads == 1 + self.run_map_helper(0) + + def get_metric_value_percentiles(self, metric_name, num_sample_per_value, total_num_samples): + logger.info(f"Checking the value percentiles of metric {metric_name}...") + processed_samples = 0 + current_percentile = 5 + for key in sorted(num_sample_per_value.keys()): + processed_samples += num_sample_per_value[key] + if processed_samples >= total_num_samples * current_percentile / 100.0: + logger.info(f"Metric {metric_name} {current_percentile}th percentile: {key}") + current_percentile += 5 + + def merge_gather_map_stats(self, num_workers, num_threads, num_threads_reduce, t_idx_reduce, metric_save_path, + metric_name, return_dict): + results = [] + for w_idx in range(num_workers): + for t_idx in range(num_threads): + if (w_idx * num_threads + t_idx) % num_threads_reduce == t_idx_reduce: + w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" + w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric" + w_sample_to_metric = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True) + unique_v = list(np.unique(w_sample_to_metric)) + sample_to_metric_count = len(w_sample_to_metric) + logger.info(f"Finished gathering map stats from worker {w_idx} thread {t_idx}.") + results.append([unique_v, sample_to_metric_count]) + return_dict[t_idx_reduce] = results + + def merge_sample_to_metric(self, t_idx_reduce, metric_save_path, metric_name, metric_value_dtype, + map_worker_thread): + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}" + sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype) + for w_t in map_worker_thread: + w_metric_save_path = f"{metric_save_path}/worker{w_t[0]}_thread{w_t[1]}/" + w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric" + w_data = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True) + for row in range(len(w_data)): + sample_to_metric_builder.add_item(torch.tensor(w_data[row].astype(np.int64), dtype=torch.long)) + logger.info(f"Finished merge_sample_to_metric from worker {w_t[0]} thread {w_t[1]}.") + close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname) + + def merge_metric_to_sample(self, t_idx_reduce, metric_save_path, metric_name, sample_idx_dtype, metric_value_dtype, + unique_metric_values, num_workers, num_threads): + index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}" + index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype) + index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}" + index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype) + for unique_v in unique_metric_values: + samples = [] + for w_idx in range(num_workers): + for t_idx in range(num_threads): + w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" + w_metric_to_sample_fname = f"{w_metric_save_path}/{metric_name}_metric_to_sample_{unique_v}.csv" + if os.path.isfile(w_metric_to_sample_fname): + with open(w_metric_to_sample_fname, 'r') as f: + datareader = csv.reader(f) + for row in datareader: + samples += [int(x) for x in row] + index_to_sample_builder.add_item(torch.tensor(samples, dtype=torch.long)) + index_to_metric_builder.add_item(torch.tensor([unique_v], dtype=torch.long)) + logger.info(f"Finished reducing metric {metric_name} value {unique_v}.") + close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname) + close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname) + + def merge_map_results(self, dataset, metric_names, metric_types, save_path, num_workers, num_threads, + num_threads_reduce): + total_num_samples = len(dataset) + sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1) + logger.info( + f"Total number of data samples: {total_num_samples}. Will use {sample_idx_dtype} to store the sample indexes." + ) + for m_idx in range(len(metric_names)): + metric_name, metric_type = metric_names[m_idx], metric_types[m_idx] + if metric_type == 'single_value_per_sample': + metric_save_path = f"{save_path}/{metric_name}/" + sample_to_metric_count = 0 + unique_metric_values = set([]) + manager = Manager() + return_dict = manager.dict() + p = [] + for t_idx_reduce in range(num_threads_reduce): + p.append( + Process(target=self.merge_gather_map_stats, + args=( + num_workers, + num_threads, + num_threads_reduce, + t_idx_reduce, + metric_save_path, + metric_name, + return_dict, + ))) + p[t_idx_reduce].start() + for t_idx_reduce in range(num_threads_reduce): + p[t_idx_reduce].join() + for t_idx_reduce in range(num_threads_reduce): + results = return_dict[t_idx_reduce] + for res in results: + unique_metric_values = unique_metric_values.union(set(res[0])) + sample_to_metric_count += res[1] + value_max = max(unique_metric_values) + value_min = min(unique_metric_values) + assert sample_to_metric_count == total_num_samples, "The number of samples in map result files are not correct. It's possible that some map worker didn't finish successfully." + metric_value_dtype = find_fit_int_dtype(value_min, value_max) + logger.info( + f"Metric {metric_name} has values between {value_min} and {value_max}. Will use {metric_value_dtype} to store the metric values." + ) + + # sample_to_metric + map_worker_thread = [] + for w_idx in range(num_workers): + for t_idx in range(num_threads): + map_worker_thread.append([w_idx, t_idx]) + thread_splits = split_index(0, len(map_worker_thread), num_threads_reduce) + p = [] + for t_idx_reduce in range(num_threads_reduce): + start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1] + p.append( + Process(target=self.merge_sample_to_metric, + args=( + t_idx_reduce, + metric_save_path, + metric_name, + metric_value_dtype, + map_worker_thread[start_idx:end_idx], + ))) + p[t_idx_reduce].start() + for t_idx_reduce in range(num_threads_reduce): + p[t_idx_reduce].join() + + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" + sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype) + for t_idx_reduce in range(num_threads_reduce): + chunk_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}" + logger.info(f"Merging file {chunk_fname}") + sample_to_metric_builder.merge_file_(chunk_fname) + close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname) + sample_to_metric = MMapIndexedDataset(sample_to_metric_fname, skip_warmup=True) + assert len(sample_to_metric) == total_num_samples + + # metric_to_sample + unique_metric_values = list(sorted(unique_metric_values)) + thread_splits = split_index(0, len(unique_metric_values), num_threads_reduce) + p = [] + for t_idx_reduce in range(num_threads_reduce): + start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1] + p.append( + Process(target=self.merge_metric_to_sample, + args=( + t_idx_reduce, + metric_save_path, + metric_name, + sample_idx_dtype, + metric_value_dtype, + unique_metric_values[start_idx:end_idx], + num_workers, + num_threads, + ))) + p[t_idx_reduce].start() + for t_idx_reduce in range(num_threads_reduce): + p[t_idx_reduce].join() + index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" + index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype) + index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" + index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype) + for t_idx_reduce in range(num_threads_reduce): + chunk_is_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}" + logger.info(f"Merging file {chunk_is_fname}") + index_to_sample_builder.merge_file_(chunk_is_fname) + chunk_im_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}" + logger.info(f"Merging file {chunk_im_fname}") + index_to_metric_builder.merge_file_(chunk_im_fname) + close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname) + close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname) + + num_sample_per_value = DataAnalyzer.output_index_to_sample_percentile( + index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, total_num_samples, + sample_idx_dtype) + self.get_metric_value_percentiles(metric_name, num_sample_per_value, total_num_samples) + elif metric_type == 'accumulate_value_over_samples': + metric_save_path = f"{save_path}/{metric_name}/" + metric_value = None + for w_idx in range(num_workers): + for t_idx in range(num_threads): + w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" + w_metric_value_fname = f"{w_metric_save_path}/{metric_name}_metric_value" + w_metric_value = MMapIndexedDataset(w_metric_value_fname, skip_warmup=True) + if metric_value is None: + metric_value = np.copy(w_metric_value[0]) + else: + metric_value += np.copy(w_metric_value[0]) + value_max = int(max(metric_value)) + value_min = int(min(metric_value)) + metric_value_dtype = find_fit_int_dtype(value_min, value_max) + metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" + metric_value_builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype) + metric_value_builder.add_item(torch.tensor(metric_value.astype(np.int64), dtype=torch.long)) + close_mmap_dataset_builder(metric_value_builder, metric_value_fname) + + @staticmethod + def output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, + total_num_samples, sample_idx_dtype): + """ read index_to_metric and index_to_sample files and write distribution to index_to_sample_percentage_merged """ + num_sample_per_value = {} + index_to_sample = MMapIndexedDataset(index_to_sample_fname, skip_warmup=True) + index_to_metric = MMapIndexedDataset(index_to_metric_fname, skip_warmup=True) + index_to_sample_merged_fname = f"{metric_save_path}/{metric_name}_index_to_sample_percentile_merged" + index_to_sample_merged_builder = create_mmap_dataset_builder(index_to_sample_merged_fname, sample_idx_dtype) + for v_idx in range(len(index_to_sample)): + if v_idx > 0: + assert index_to_metric[v_idx] > index_to_metric[v_idx - 1] + num_sample_per_value[index_to_metric[v_idx][0]] = len(index_to_sample[v_idx]) + assert sum(list(num_sample_per_value.values())) == total_num_samples + merge_step = max(1, len(index_to_sample) // 100) + for v_idx in range(0, len(index_to_sample), merge_step): + merged_samples = np.copy( + np.concatenate(index_to_sample[v_idx:min(len(index_to_sample), (v_idx + merge_step))], axis=None)) + index_to_sample_merged_builder.add_item(torch.tensor(merged_samples.astype(np.int64), dtype=torch.long)) + logger.info(f"Finished merging index_to_sample {v_idx} to {v_idx+merge_step}.") + close_mmap_dataset_builder(index_to_sample_merged_builder, index_to_sample_merged_fname) + return num_sample_per_value + + def run_reduce(self): + if self.custom_reduce is None: + self.merge_map_results(self.dataset, self.metric_names, self.metric_types, self.save_path, + self.num_workers, self.num_threads, self.num_threads_reduce) + else: + self.custom_reduce(self.dataset, self.metric_names, self.metric_types, self.save_path, self.num_workers, + self.num_threads, self.num_threads_reduce) + + def run_map_reduce(self, comm_group=None): + self.run_map() + # wait for the mapping operation, where all nodes outputs their own (partial) result files + dist.barrier(group=comm_group) + if self.worker_id == 0: + self.run_reduce() + # wait for the reduce, where rank 0 merges all (partial) files. Dataset can then be used by all nodes. + dist.barrier(group=comm_group) + + +class DistributedDataAnalyzer(object): + + def __init__( + self, + dataset, + num_workers=1, + num_threads=1, + worker_id=0, + batch_size=1, + metric_names=[], + metric_functions=[], + metric_types=[], + save_path="./", + collate_fn=None, + device='cuda', + comm_group=None, + sample_indices=None, + ) -> None: + self.dataset = dataset + self.batch_size = batch_size + self.metric_names = metric_names + self.metric_functions = metric_functions + self.metric_types = metric_types + self.save_path = save_path + self.collate_fn = collate_fn + self.device = device + self.sample_indices = sample_indices + self.num_threads = num_threads + self.worker_id = worker_id + + if not dist.is_initialized(): + dist.init_distributed() + + # comm_group and worker_id+num_workers are mutually exclusive + self.comm_group = comm_group + if self.comm_group is None: + # self.comm_group = deepspeed.utils.groups._clone_world_group() + self.num_workers = num_workers + self.worker_id = worker_id + else: + self.num_workers = self.comm_group.size() + self.worker_id = self.comm_group.rank() + + if self.worker_id == 0: + logger.info(f"Distributed data analyzer initialized with {self.num_workers} workers.") + + def run_map_helper(self, thread_id=0, metric_queues=None): + thread_start_idx, thread_end_idx = self.thread_splits[thread_id][0], self.thread_splits[thread_id][1] + worker_dataset = Subset(self.dataset, list(range(thread_start_idx, thread_end_idx))) + sampler = BatchSampler(SequentialSampler(worker_dataset), batch_size=self.batch_size, drop_last=False) + dataloader = DataLoader(dataset=worker_dataset, + batch_sampler=sampler, + num_workers=0, + collate_fn=self.collate_fn, + pin_memory=False) + + # set initial results list + metric_results = [] + for metric_type in self.metric_types: + assert metric_type in ['single_value_per_sample', 'accumulate_value_over_samples'], \ + f"metric_type {metric_type} not implemented." + metric_results.append([] if metric_type == 'single_value_per_sample' else None) + + # iterate dataloader and store metric results + batch_start_idx = thread_start_idx + for data in dataloader: + for m_idx in range(len(self.metric_names)): + metric_type, metric_function = self.metric_types[m_idx], self.metric_functions[m_idx] + metric_values = metric_function(data) + assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \ + "metric_function must return a tensor or array" + if isinstance(metric_values, np.ndarray): + metric_values = torch.from_numpy(metric_values) + assert metric_values.dtype in valid_dtypes, \ + f"metric_function result dtype {metric_values.dtype} not supported. Supported dtypes {valid_dtypes}" + + if metric_type == 'single_value_per_sample': + for row in range(metric_values.size()[0]): + value = metric_values[row].item() + sample_idx = batch_start_idx + row # sample idx following dataset iteration order + if isinstance(data, dict) and 'index' in data: # Megatron use case + sample_idx = data['index'][row][0].item() + elif self.sample_indices is not None: # user defined shuffling of indices + sample_idx = self.sample_indices[sample_idx] + metric_results[m_idx].append((value, sample_idx)) + elif metric_type == 'accumulate_value_over_samples': + if metric_results[m_idx] is None: + metric_results[m_idx] = metric_values + else: + metric_results[m_idx].add_(metric_values) + batch_start_idx += len(data) + + if self.num_threads == 1: + return metric_results + + # copy metric_results to the shared queue + assert metric_queues + for m_idx in range(len(self.metric_names)): + results = metric_results[m_idx] + if torch.is_tensor(results): + results = results.item() if results.dim() == 0 else results.tolist() + try: + metric_queues[m_idx].put((thread_id, results)) + except Exception as e: + logger.error(f"Error putting metric results to queue: {e}") + sys.exit(1) + + def run_map_reduce(self): + + # setup individual dataloaders + self.worker_splits, self.thread_splits = split_dataset(self.dataset, + self.num_workers, + self.worker_id, + num_threads=self.num_threads) + node_start_idx, node_end_idx = self.worker_splits[self.worker_id] + logger.info(f"worker {self.worker_id} working on data subset {node_start_idx} to {node_end_idx}.") + + if self.num_threads in [0, 1, None]: + metric_results = self.run_map_helper() + metric_results = [torch.tensor(m).to(self.device) for m in metric_results] + else: + + # create a shared queue of results per metric to be populated by individual threads + with Manager() as manager: + metric_queues = [manager.Queue() for _ in self.metric_names] + threads = [ + Process(target=self.run_map_helper, args=(t, metric_queues)) for t in range(self.num_threads) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + # gather results from shared queues into metric_results + metric_results = [None for _ in self.metric_names] + for m_idx, (queue, metric_type) in enumerate(zip(metric_queues, self.metric_types)): + while not queue.empty(): + t_idx, t_results = queue.get() + t_start_idx, t_end_idx = self.thread_splits[t_idx] + if t_start_idx >= t_end_idx: # no results from this thread + continue #corner case for small datasets and high thread count + t_results = torch.tensor(t_results) + if metric_type == 'single_value_per_sample': + # add thread results to the metric_results list, ordered by thread idx + if metric_results[m_idx] is None: # initialize if needed + metric_results[m_idx] = torch.zeros(node_end_idx - node_start_idx, + t_results.size(1)).to(self.device) + metric_results[m_idx][t_start_idx - node_start_idx:t_end_idx - node_start_idx] = t_results + else: + if metric_results[m_idx] is None: # initialize if needed + metric_results[m_idx] = torch.zeros(t_results.size()).to(self.device) + metric_results[m_idx].add_(t_results) + + # compute dtype for sample ids + total_num_samples = len(self.dataset) + sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1) + logger.info(f"Total number of data samples: {total_num_samples}.") + logger.info(f"Will use {sample_idx_dtype} to store the sample indexes.") + + for m_idx in range(len(self.metric_names)): + metric_values, metric_name, metric_type = \ + metric_results[m_idx], self.metric_names[m_idx], self.metric_types[m_idx] + metric_save_path = f"{self.save_path}/{metric_name}/" + os.makedirs(metric_save_path, exist_ok=True) + + if metric_type == 'single_value_per_sample': + + # Compute sample and metric value dtypes based on range + values, samples = metric_values[:, 0], metric_values[:, 1] + value_min, value_max = Dist.min_max(values, self.comm_group) + sample_min, sample_max = Dist.min_max(samples, self.comm_group) + metric_value_dtype = find_fit_int_dtype(value_min, value_max) + sample_value_dtype = find_fit_int_dtype(sample_min, sample_max) + + # sample_to_metric maps sample ids to metric values, as a list of metric values + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" + values = [torch.tensor([x]) for x in metric_values[:, 0]] + self.file_write_ordered(values, sample_to_metric_fname, metric_value_dtype) + + # distributed sorting by values, gives an ordered disjoint subset of keys on nodes + metric_values = Dist.sample_sort(metric_values, self.comm_group, self.num_workers) + metric_to_samples_dict = {} + if len(metric_values) > 0: + for value, sample in metric_values: + if value.item() not in metric_to_samples_dict: + metric_to_samples_dict[value.item()] = [] + metric_to_samples_dict[value.item()].append(sample.item()) + + # index_to_metric and index_to_sample serialize a dicitonary from metric to samples + # index_to_metric stores a key per row, index_to_sample stores the values per row + values = [torch.tensor([x]) for x in metric_to_samples_dict.keys()] + samples = [torch.tensor(metric_to_samples_dict[x]) for x in metric_to_samples_dict.keys()] + index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" #dict keys + index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" #dict values + self.file_write_ordered(values, index_to_metric_fname, metric_value_dtype) + self.file_write_ordered(samples, index_to_sample_fname, sample_value_dtype) + + if self.worker_id == 0: + DataAnalyzer.output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, + metric_name, metric_save_path, total_num_samples, + sample_idx_dtype) + dist.barrier(self.comm_group) + + elif metric_type == 'accumulate_value_over_samples': + metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" + dist.reduce(metric_values, dst=0, op=dist.ReduceOp.SUM, group=self.comm_group) + metric_value_dtype = find_fit_int_dtype(metric_values.min(), metric_values.max()) + + if self.worker_id == 0: + builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype) + builder.add_item(metric_values.cpu()) + close_mmap_dataset_builder(builder, metric_value_fname) + dist.barrier(self.comm_group) + + def file_write_ordered(self, tensor_list, fname, numpy_dtype): + """ MPI_file_write_ordered extended to write a list of tensors, by one rank, iteratively """ + + # each node has a list of rows (tensors) to be written to the file. + # we will serialize it in order to communicate it in one comm step. + + tkwargs = dict(dtype=torch.int64, device=self.device) + + # 1. gather on rank 0 the number of rows to be sent/recv + row_count = torch.tensor([len(tensor_list)], **tkwargs) + row_counts = torch.zeros(self.num_workers, **tkwargs) + dist.all_gather_into_tensor(row_counts, row_count, group=self.comm_group) + assert row_counts[self.worker_id] == row_count == len(tensor_list), "all_gather failed" + + # 2. gather on rank 0 the sizes of the rows to be sent/recv + row_len = torch.tensor([len(l) for l in tensor_list], **tkwargs) + row_lens = Dist.gather_v(row_len, 0, self.comm_group, self.num_workers, self.worker_id) + + # 4. gather on rank 0 of the total size (sum of all row lengths) to be received + size = torch.tensor([sum(row_len).item()], **tkwargs) + sizes = torch.zeros(self.num_workers, **tkwargs) + dist.all_gather_into_tensor(sizes, size, group=self.comm_group) + assert sizes[self.worker_id] == size.item(), "all_gather did not return the same sizes" #sanity check + + # method to deserializes a buffer into rows of different lengths and write them to file + def write_buffer_to_file(buff, src, builder): + assert self.worker_id == 0, "only rank 0 can write to file" + + # collect all buffers and write them at once + buff = buff.cpu().detach().numpy() + row_offsets = np.cumsum([0] + row_lens[src].tolist()) + arr_list = [] + for i in range(len(row_lens[src])): + arr_list.append(buff[row_offsets[i]:row_offsets[i + 1]]) + builder.add_items(arr_list) + + # 5. rank 0 prepares output folder and file + if self.worker_id == 0: + os.makedirs(os.path.dirname(fname), exist_ok=True) + builder = create_mmap_dataset_builder(fname, numpy_dtype) + + # iterate through ranks that have data to be sent/recv/written + for src in [rank for rank, count in enumerate(row_counts) if count > 0]: + + dist.barrier(group=self.comm_group) + if self.worker_id == 0 and src == 0: # rank 0's write its own data + buffer = torch.cat(tensor_list, dim=0).to(self.device) + write_buffer_to_file(buffer, 0, builder) + elif self.worker_id == 0 and src > 0: # rank 0 receives other rank's data and writes it + buffer = torch.empty(sizes[src].item(), dtype=buffer.dtype, device=buffer.device) + err = dist.recv(buffer, src=src, group=self.comm_group, tag=src) + assert err == src and len(buffer) > 0, "recv failed" + write_buffer_to_file(buffer, src, builder) + elif self.worker_id == src: # current rank sends data to rank 0 + buffer = torch.cat(tensor_list, dim=0).to(self.device) + dist.send(buffer, 0, group=self.comm_group, tag=src) + + # rank 0 closes the file + if self.worker_id == 0: + close_mmap_dataset_builder(builder, fname) # close file + dist.barrier(self.comm_group) + + +class Dist: + """ auxiliary class to perform distributed operations on tensors""" + + @staticmethod + def min_max(tensor, comm_group): + """ given a distributed tensor, return the min/max values across all ranks""" + + value_min, value_max = tensor.min(), tensor.max() + dist.reduce(value_min, 0, op=dist.ReduceOp.MIN, group=comm_group) + dist.reduce(value_max, 0, op=dist.ReduceOp.MAX, group=comm_group) + return value_min.item(), value_max.item() + + @staticmethod + def gather_v(tensor, dst, comm_group, num_workers, worker_id): + """ MPI_Gatherv. gather tensors of variable sizes in a single rank """ + + # gather the number of rows to be sent/recv + size = torch.tensor([len(tensor)], dtype=torch.int64, device=tensor.device) + sizes = torch.zeros(num_workers, dtype=torch.int64, device=tensor.device) + dist.all_gather_into_tensor(sizes, size, group=comm_group) + assert sizes[worker_id] == size, "all_gather failed" + + # all_gather requires all tensors to be of same size so we need to pad them + max_size = max(sizes).item() + buffer = torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) + buffer[0:size] = tensor.data + buffer_list = None + if worker_id == 0: # create padded recv buffers + buffer_list = [torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) for _ in range(num_workers)] + dist.gather(buffer, buffer_list, dst=dst, group=comm_group) + + # revert padding and return value + if worker_id == 0: + buffer_list = [r[:s.item()] for r, s in zip(buffer_list, sizes)] + return buffer_list + + @staticmethod + def sample_sort(tensor, comm_group, num_workers, n_samples=100): + """ perform a distributed random sort of a tensor, and returns the sorted partial tensor""" + device, dims = tensor.device, tensor.size()[1] + + # 1 - sort rows by first column, then second column, then third, etc... + tensor = torch.tensor(sorted(tensor.tolist()), dtype=tensor.dtype, device=tensor.device) + + # 2 - collect few samples per rank + idx = torch.round(torch.linspace(0, len(tensor) - 1, n_samples)).to(int) + samples = tensor[idx][:, 0].contiguous().to(device) #only first column, all but last row + + # 2 - Allgather samples + all_samples = [torch.zeros(n_samples, dtype=samples.dtype, device=device) for _ in range(num_workers)] + dist.all_gather(all_samples, samples, group=comm_group) + all_samples = torch.cat(all_samples, dim=0).to(device) + + # 3 - Sort all samples and collect the ranges of each rank as equidistant + all_samples = all_samples.sort()[0] + idx = torch.round(torch.linspace(0, len(all_samples) - 1, num_workers + 1)).to(int) + ranges = all_samples[idx] # range of each rank r as ranges[r] <= x < ranges[r+1] + ranges[-1] += 1 # increase upper limit of last rank so that x < ranges[r+1]. + + # 4 - collect elements to send to each rank, based on the rank ranges + send = [] + for rank in range(num_workers): + mask = (tensor[:, 0] >= ranges[rank]) & (tensor[:, 0] < ranges[rank + 1]) + send.append(tensor[mask]) + + # 5. all to all to communicate the sizes to be sent/recv + send_count = [torch.tensor([len(s) * dims], dtype=torch.int64, device=device) for s in send] + recv_count = list(torch.empty([num_workers], dtype=torch.int64, device=device).chunk(num_workers)) + dist.all_to_all(recv_count, send_count, group=comm_group) + + # 6. all-to-all-v to communicate the elements to be sent/recv as a single tensor + send = torch.cat(send, dim=0).flatten().to(device) + recv = torch.zeros(sum(recv_count), dtype=send.dtype).to(device) + send_count = [s.item() for s in send_count] # convert to list of ints + recv_count = [r.item() for r in recv_count] + dist.all_to_all_single(recv, send, recv_count, send_count, group=comm_group) + del send + + # 7. the received tensor is the 1D disjoint subset of the distributed tensor. + # We will recover the original dimensionality and sort it by columns again. + recv = recv.view(-1, dims) + recv = torch.tensor(sorted(recv.tolist()), dtype=recv.dtype, device=recv.device) + return recv + + +def test_compare_both_data_analyzers(dataset): + """ given a dataset, compare file and memory based data analyser""" + + id = lambda t: t.to(torch.int64) # identity + batch_sum = lambda t: id(t).sum() #sum batch + num_threads = 4 + kwargs = dict( + dataset=dataset, + batch_size=2**10, + worker_id=int(os.environ['RANK']), + num_workers=int(os.environ['WORLD_SIZE']), + metric_names=["mod", "batch_sum"], + metric_functions=[id, batch_sum], + metric_types=['single_value_per_sample', 'accumulate_value_over_samples'], + num_threads=num_threads, + ) + + dda = DistributedDataAnalyzer( + save_path="./output_dist", + device=f"cuda:{int(os.environ['LOCAL_RANK'])}", + **kwargs, + ) + start_time = time.time() + dda.run_map_reduce() + if dda.worker_id == 0: + print("DistributedDataAnalyzer runtime: %s seconds " % (time.time() - start_time)) + + da = DataAnalyzer(num_threads_reduce=num_threads, + save_path="./output_disk", + metric_dtypes=[torch.int64, torch.int64], + **kwargs) + start_time = time.time() + da.run_map_reduce() + if da.worker_id == 0: + print("DataAnalyzer runtime: %s seconds " % (time.time() - start_time)) + + output_paths = [ + "batch_sum/batch_sum_metric_value.bin", "batch_sum/batch_sum_metric_value.idx", \ + "mod/mod_index_to_metric.bin", "mod/mod_index_to_metric.idx", \ + "mod/mod_index_to_sample.bin", "mod/mod_index_to_sample.idx", \ + "mod/mod_index_to_sample_percentile_merged.bin", "mod/mod_index_to_sample_percentile_merged.idx", \ + "mod/mod_sample_to_metric.bin", "mod/mod_sample_to_metric.idx" + ] + + if dda.worker_id == 0: + for path in output_paths: + with open(os.path.join(da.save_path, path), 'rb') as f1, \ + open(os.path.join(dda.save_path, path), 'rb') as f2: + # if files have suffix .bin, they should be identical + if path.endswith(".bin"): + assert f1.read() == f2.read(), f"files {path} are not identical." + elif f1.read() != f2.read(): + print(f"files {path} are not identical.") + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + + class TestDataset(torch.utils.data.Dataset): + + def __init__(self, size=10_000_000): + self.values = [(x + 7) % 10_000 for x in range(size)] + self.size = size + + __len__ = lambda self: self.size + __getitem__ = lambda self, idx: self.values[idx] + + test_compare_both_data_analyzers(TestDataset()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..100bef3f7946c8d2b7280e234abc5e625b6c68f9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py @@ -0,0 +1,349 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +coding=utf-8 + Copyright (c) 2020, 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. +Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/data_samplers.py +""" + +import torch +import os +import numpy as np + +import deepspeed.comm as dist +from deepspeed.utils import logger +from deepspeed.accelerator import get_accelerator +from ..constants import * +from ..curriculum_scheduler import CurriculumScheduler +from .indexed_dataset import MMapIndexedDataset +from .utils import create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype + + +class DeepSpeedDataSampler(object): + + def __init__(self, + data_efficiency_config, + one_epoch_total_samples, + micro_batch_size, + data_parallel_rank, + data_parallel_size, + data_parallel_group, + gradient_accumulation_steps, + global_rank, + drop_last=True): + # Keep a copy of input params for later use. + self.data_efficiency_config = data_efficiency_config + self.one_epoch_total_samples = one_epoch_total_samples + self.index_dtype = find_fit_int_dtype(0, one_epoch_total_samples) + self.total_samples = one_epoch_total_samples * self.data_efficiency_config[DATA_SAMPLING][ + DATA_SAMPLING_NUM_EPOCHS] + self.micro_batch_size = micro_batch_size + self.data_parallel_rank = data_parallel_rank + self.data_parallel_group = data_parallel_group + self.micro_batch_times_data_parallel_size = \ + self.micro_batch_size * data_parallel_size + self.gradient_accumulation_steps = gradient_accumulation_steps + self.global_batch_size = self.micro_batch_times_data_parallel_size * \ + self.gradient_accumulation_steps + self.global_rank = global_rank + self.drop_last = drop_last + self.np_rng = np.random.default_rng(self.data_efficiency_config[DATA_EFFICIENCY_SEED]) + self.state = {} + self.batch = [] + self.consumed_samples = 0 + if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]: + self.curriculum_step = 0 + self.current_difficulties = {} + self.data_cluster_paths = [] + self.data_cluster_current_position = [] + self.curriculum_schedulers = {} + self.curriculum_index_to_sample = {} + self.curriculum_index_to_metric = {} + self.difficulty_type = {} + self.clustering_type = {} + self.data_1epoch_size = None + if self.global_rank == 0: + self.data_clusters = [] + self.data_cluster_sizes = [] + cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + if not os.path.exists(cluster_path): + os.makedirs(cluster_path) + for metric in self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS]: + self.curriculum_schedulers[metric] = CurriculumScheduler( + data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][metric]) + self.difficulty_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_DIFFICULTY_TYPE] + self.clustering_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_CLUSTERING_TYPE] + if self.global_rank == 0: + if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: + self.curriculum_index_to_sample[metric] = MMapIndexedDataset( + data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS] + [metric][CURRICULUM_LEARNING_SAMPLE_PATH], + skip_warmup=True) + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + self.curriculum_index_to_metric[metric] = MMapIndexedDataset( + data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS] + [metric][CURRICULUM_LEARNING_METRIC_PATH], + skip_warmup=True) + + # Sanity checks. + assert self.total_samples > 0, \ + 'no sample to consume: {}'.format(self.total_samples) + assert self.micro_batch_size > 0 + assert data_parallel_size > 0 + assert self.data_parallel_rank < data_parallel_size, \ + 'data_parallel_rank should be smaller than data size: {}, ' \ + '{}'.format(self.data_parallel_rank, data_parallel_size) + + def __len__(self): + return self.total_samples + + def set_custom_curriculum_learning_schedule(self, schedule_func_dict): + for metric in self.curriculum_schedulers: + if metric in schedule_func_dict: + self.curriculum_schedulers[metric].set_custom_get_difficulty(schedule_func_dict[metric]) + + def get_start_end_idx(self, batch_len=None): + """ + given the length of a minibatch (defaults to micro-batch size * data_parallel_size), + return the start and end indices of the current data parallel rank + """ + batch_len = batch_len or self.micro_batch_times_data_parallel_size + start_idx_fn = lambda r: round(r * batch_len / self.data_parallel_group.size()) + start_idx = start_idx_fn(self.data_parallel_rank) + end_idx = start_idx_fn(self.data_parallel_rank + 1) + return start_idx, end_idx + + def get_sample_based_on_metric_value(self, metric, value_start, value_end): + new_samples = None + for row in range(len(self.curriculum_index_to_sample[metric])): + if self.curriculum_index_to_metric[metric][row] <= value_end and self.curriculum_index_to_metric[metric][ + row] > value_start: + row_samples = np.copy(self.curriculum_index_to_sample[metric][row]) + new_samples = row_samples if new_samples is None else np.concatenate( + (new_samples, row_samples), axis=None) + return new_samples + + def get_sample_based_on_metric_percentile(self, metric, percentile_start, percentile_end): + new_samples = None + if self.data_1epoch_size is None: + self.data_1epoch_size = sum(len(x) for x in self.curriculum_index_to_sample[metric]) + max_percentile = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][ + metric][CURRICULUM_LEARNING_MAX_DIFFICULTY] + sample_per_percentile = self.data_1epoch_size // max_percentile + start_count = sample_per_percentile * percentile_start + end_count = sample_per_percentile * percentile_end + if percentile_end == max_percentile: + end_count = self.data_1epoch_size + current_count = 0 + for row in range(len(self.curriculum_index_to_sample[metric])): + row_size = len(self.curriculum_index_to_sample[metric][row]) + if current_count + row_size > start_count: + row_start = max(0, start_count - current_count) + if current_count + row_size <= end_count: + row_end = row_size + else: + row_end = end_count - current_count + row_samples = np.copy(self.curriculum_index_to_sample[metric][row][row_start:row_end]) + new_samples = row_samples if new_samples is None else np.concatenate( + (new_samples, row_samples), axis=None) + current_count += row_size + if current_count >= end_count: + break + return new_samples + + def get_new_cluster(self, previous_difficulties): + cluster_fname = CURRICULUM_LEARNING_CLUSTER_PREFIX + for metric in self.curriculum_schedulers: + cluster_fname = f"{cluster_fname}_{metric}{self.current_difficulties[metric]}" + cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + cluster_path = f"{cluster_path}/{cluster_fname}" + if self.global_rank == 0: + new_cluster = None + need_clustering = 0 + for metric in self.clustering_type: + if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: + need_clustering += 1 + if need_clustering > 1: + for metric in self.curriculum_schedulers: + if self.clustering_type[metric] == CURRICULUM_LEARNING_SINGLE_CLUSTER: + metric_cluster = np.arange(start=0, + stop=self.one_epoch_total_samples, + step=1, + dtype=self.index_dtype) + else: + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + metric_cluster = self.get_sample_based_on_metric_value(metric, float('-inf'), + self.current_difficulties[metric]) + elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: + metric_cluster = self.get_sample_based_on_metric_percentile( + metric, 0, self.current_difficulties[metric]) + new_cluster = metric_cluster if new_cluster is None else \ + np.intersect1d(new_cluster, metric_cluster, assume_unique=True) + for cluster in self.data_clusters: + new_cluster = np.setdiff1d(new_cluster, cluster[0], assume_unique=True) + else: + if len(self.data_clusters) == 0: + new_cluster = np.arange(start=0, stop=self.one_epoch_total_samples, step=1, dtype=self.index_dtype) + for metric in self.curriculum_schedulers: + if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + new_cluster = self.get_sample_based_on_metric_value(metric, previous_difficulties[metric], + self.current_difficulties[metric]) + elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: + new_cluster = self.get_sample_based_on_metric_percentile( + metric, previous_difficulties[metric], self.current_difficulties[metric]) + if new_cluster is not None and len(new_cluster) > 0: + logger.info( + f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) with size {len(new_cluster)} generated." + ) + self.np_rng.shuffle(new_cluster) + cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype) + cluster_builder.add_item_numpy(new_cluster) + close_mmap_dataset_builder(cluster_builder, cluster_path) + self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True)) + self.data_cluster_sizes.append(len(self.data_clusters[-1][0])) + else: + logger.info( + f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) has no matched data thus skipped." + ) + dist.barrier(group=self.data_parallel_group) + if os.path.isfile(f"{cluster_path}.bin"): + self.data_cluster_paths.append(cluster_fname) + self.data_cluster_current_position.append(0) + + def sample_from_clusters(self): + num_clusters = len(self.data_clusters) + weight_sum = sum(self.data_cluster_sizes) + weights = [x / weight_sum for x in self.data_cluster_sizes] + samples = self.np_rng.choice(num_clusters, self.global_batch_size, replace=True, p=weights) + samples = np.bincount(samples, minlength=num_clusters) + return samples + + def reshuffle_clusters(self, cidx): + cluster_fname = self.data_cluster_paths[cidx] + cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + cluster_path = f"{cluster_path}/{cluster_fname}" + cluster = np.copy(self.data_clusters[cidx][0]) + self.np_rng.shuffle(cluster) + cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype) + cluster_builder.add_item_numpy(cluster) + close_mmap_dataset_builder(cluster_builder, cluster_path) + self.data_clusters[cidx] = MMapIndexedDataset(cluster_path, skip_warmup=True) + + def get_sample_from_cluster(self, cidx, num_samples): + start_idx = self.data_cluster_current_position[cidx] + samples = list(np.copy(self.data_clusters[cidx][0][start_idx:(start_idx + num_samples)])) + self.data_cluster_current_position[cidx] += num_samples + if len(samples) < num_samples: + num_samples_remained = num_samples - len(samples) + logger.info(f"reshuffling cluster {cidx}.") + self.reshuffle_clusters(cidx) + samples += list(np.copy(self.data_clusters[cidx][0][:num_samples_remained])) + self.data_cluster_current_position[cidx] = num_samples_remained + return samples + + def get_next_global_batch(self): + if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]: + self.curriculum_step += 1 + new_cluster = False + previous_difficulties = {} + for metric in self.curriculum_schedulers: + next_difficulty = self.curriculum_schedulers[metric].update_difficulty(self.curriculum_step) + if metric not in self.current_difficulties or \ + next_difficulty != self.current_difficulties[metric]: + new_cluster = True + if metric in self.current_difficulties: + previous_difficulties[metric] = self.current_difficulties[metric] + else: + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + previous_difficulties[metric] = float('-inf') + elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: + previous_difficulties[metric] = 0 + self.current_difficulties[metric] = next_difficulty + if new_cluster: + self.get_new_cluster(previous_difficulties) + if self.global_rank == 0: + samples_per_cluster = self.sample_from_clusters() + batch = [] + for cidx in range(len(samples_per_cluster)): + batch += self.get_sample_from_cluster(cidx, samples_per_cluster[cidx]) + self.np_rng.shuffle(batch) + + # broadcast tensor must have same shape across participants. So we fill batch with -1s when not full + assert len(batch) <= self.global_batch_size + batch += [-1] * (self.global_batch_size - len(batch)) + batch = torch.tensor(batch, device=get_accelerator().current_device_name(), dtype=torch.long).view(-1) + else: + batch = torch.empty(self.global_batch_size, + device=get_accelerator().current_device_name(), + dtype=torch.long) + dist.broadcast(batch, 0, group=self.data_parallel_group) + batch = batch[batch != -1] # remove trailing -1s used to fill incomplete batch tensor + self.batch = batch.tolist() + + def __iter__(self): + while self.consumed_samples <= self.total_samples: + if len(self.batch) == 0: + self.get_next_global_batch() + current_batch = self.batch[:self.micro_batch_times_data_parallel_size] + self.batch = self.batch[self.micro_batch_times_data_parallel_size:] + if len(current_batch) == self.micro_batch_times_data_parallel_size or \ + (len(current_batch) > 0 and not self.drop_last): + start_idx, end_idx = self.get_start_end_idx(len(current_batch)) + yield current_batch[start_idx:end_idx] + self.consumed_samples += len(current_batch) + current_batch = [] + + def state_dict(self): + return { + CURRICULUM_LEARNING_BATCH: self.batch, + CURRICULUM_LEARNING_CONSUMED_SAMPLES: self.consumed_samples, + CURRICULUM_LEARNING_STEP: self.curriculum_step, + CURRICULUM_LEARNING_CURRENT_DIFFICULTIES: self.current_difficulties, + CURRICULUM_LEARNING_DATA_CLUSTER_PATHS: self.data_cluster_paths, + CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION: self.data_cluster_current_position, + CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state() + } + + def load_state_dict(self, state_dict): + self.batch = state_dict[CURRICULUM_LEARNING_BATCH] + self.consumed_samples = state_dict[CURRICULUM_LEARNING_CONSUMED_SAMPLES] + self.curriculum_step = state_dict[CURRICULUM_LEARNING_STEP] + self.current_difficulties = state_dict[CURRICULUM_LEARNING_CURRENT_DIFFICULTIES] + self.data_cluster_paths = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_PATHS] + self.data_cluster_current_position = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION] + np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE]) + cluster_root_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + # Backward compatibility: previously data_cluster_paths were stored as + # absolute paths. Now we changed it to just the file name so that even + # if user moved the cluster files, the checkpoint loading still works + # as long as user set the correct new CURRICULUM_LEARNING_CLUSTER_PATH + # in deepspeed json config. + for idx in range(len(self.data_cluster_paths)): + if '/' in self.data_cluster_paths[idx]: + self.data_cluster_paths[idx] = self.data_cluster_paths[idx].split('/')[-1] + if self.global_rank == 0: + for cluster_fname in self.data_cluster_paths: + cluster_path = f"{cluster_root_path}/{cluster_fname}" + self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True)) + self.data_cluster_sizes.append(len(self.data_clusters[-1][0])) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..872d05de01453aea76fc850de77da3acbbcea76c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py @@ -0,0 +1,627 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py +""" + +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# copied from fairseq/fairseq/data/indexed_dataset.py +# Removed IndexedRawTextDataset since it relied on Fairseq dictionary +# other slight modifications to remove fairseq dependencies +# Added document index to index file and made it accessible. +# An empty sentence no longer separates documents. + +# Some of the fixes/improvements are adopted from +# https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/megatron/data/indexed_dataset.py + +from functools import lru_cache +import os +import shutil +import struct +from itertools import accumulate + +import numpy as np +import torch + + +def __best_fitting_dtype(vocab_size=None): + if vocab_size is not None and vocab_size < 65500: + return np.uint16 + else: + return np.int32 + + +def get_available_dataset_impl(): + return ['lazy', 'cached', 'mmap'] + + +def infer_dataset_impl(path): + if IndexedDataset.exists(path): + with open(index_file_path(path), 'rb') as f: + magic = f.read(8) + if magic == IndexedDataset._HDR_MAGIC: + return 'cached' + elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]: + return 'mmap' + else: + return None + else: + print(f"Dataset does not exist: {path}") + print("Path should be a basename that both .idx and .bin can be appended to get full filenames.") + return None + + +def make_builder(out_file, impl, vocab_size=None): + if impl == 'mmap': + return MMapIndexedDatasetBuilder(out_file, dtype=__best_fitting_dtype(vocab_size)) + else: + return IndexedDatasetBuilder(out_file) + + +def make_dataset(path, impl, skip_warmup=False): + if not IndexedDataset.exists(path): + print(f"Dataset does not exist: {path}") + print("Path should be a basename that both .idx and .bin can be appended to get full filenames.") + return None + if impl == 'infer': + impl = infer_dataset_impl(path) + if impl == 'lazy' and IndexedDataset.exists(path): + return IndexedDataset(path) + elif impl == 'cached' and IndexedDataset.exists(path): + return IndexedCachedDataset(path) + elif impl == 'mmap' and MMapIndexedDataset.exists(path): + return MMapIndexedDataset(path, skip_warmup) + print(f"Unknown dataset implementation: {impl}") + return None + + +def dataset_exists(path, impl): + if impl == 'mmap': + return MMapIndexedDataset.exists(path) + else: + return IndexedDataset.exists(path) + + +def read_longs(f, n): + a = np.empty(n, dtype=np.int64) + f.readinto(a) + return a + + +def write_longs(f, a): + f.write(np.array(a, dtype=np.int64)) + + +# valid metric_dtypes as numpy and torch types +dtypes = { + 1: (np.uint8, torch.uint8), + 2: (np.int8, torch.int8), + 3: (np.int16, torch.int16), + 4: (np.int32, torch.int32), + 5: (np.int64, torch.int64), + 6: (np.uint16, None), + 7: (np.uint32, None), + 8: (np.uint64, None), +} + +valid_dtypes = set([dt[0] for dt in dtypes.values()] + [dt[1] for dt in dtypes.values() if dt[1] is not None]) + + +def code(dtype): + for c, (np_dt, torch_dt) in dtypes.items(): + if dtype in [np_dt, torch_dt]: + return c + raise ValueError(f"{dtype} not supported. Supported types: {valid_dtypes}") + + +def index_file_path(prefix_path): + return prefix_path + '.idx' + + +def data_file_path(prefix_path): + return prefix_path + '.bin' + + +def create_doc_idx(sizes): + doc_idx = [0] + for i, s in enumerate(sizes): + if s == 0: + doc_idx.append(i + 1) + return doc_idx + + +class IndexedDataset(torch.utils.data.Dataset): + """Loader for IndexedDataset""" + _HDR_MAGIC = b'TNTIDX\x00\x00' + + def __init__(self, path): + super().__init__() + self.path = path + self.data_file = None + self.read_index(path) + + def read_index(self, path): + with open(index_file_path(path), 'rb') as f: + magic = f.read(8) + assert magic == self._HDR_MAGIC, ('Index file doesn\'t match expected format. ' + 'Make sure that --dataset-impl is configured properly.') + version = f.read(8) + assert struct.unpack('= self._len: + raise IndexError('index out of range') + + def __del__(self): + if self.data_file: + self.data_file.close() + + # @lru_cache(maxsize=8) + def __getitem__(self, idx): + if not self.data_file: + self.read_data(self.path) + if isinstance(idx, int): + i = idx + self.check_index(i) + tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]] + a = np.empty(tensor_size, dtype=self.dtype) + self.data_file.seek(self.data_offsets[i] * self.element_size) + self.data_file.readinto(a) + return a + elif isinstance(idx, slice): + start, stop, step = idx.indices(len(self)) + if step != 1: + raise ValueError("Slices into indexed_dataset must be contiguous") + sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]] + size = sum(sizes) + a = np.empty(size, dtype=self.dtype) + self.data_file.seek(self.data_offsets[start] * self.element_size) + self.data_file.readinto(a) + offsets = list(accumulate(sizes)) + sents = np.split(a, offsets[:-1]) + return sents + + def __len__(self): + return self._len + + def num_tokens(self, index): + return self.sizes[index] + + def size(self, index): + return self.sizes[index] + + @staticmethod + def exists(path): + return (os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path))) + + @property + def supports_prefetch(self): + return False # avoid prefetching to save memory + + +class IndexedCachedDataset(IndexedDataset): + + def __init__(self, path): + super().__init__(path) + self.cache = None + self.cache_index = {} + + @property + def supports_prefetch(self): + return True + + def prefetch(self, indices): + if all(i in self.cache_index for i in indices): + return + if not self.data_file: + self.read_data(self.path) + indices = sorted(set(indices)) + total_size = 0 + for i in indices: + total_size += self.data_offsets[i + 1] - self.data_offsets[i] + self.cache = np.empty(total_size, dtype=self.dtype) + ptx = 0 + self.cache_index.clear() + for i in indices: + self.cache_index[i] = ptx + size = self.data_offsets[i + 1] - self.data_offsets[i] + a = self.cache[ptx:ptx + size] + self.data_file.seek(self.data_offsets[i] * self.element_size) + self.data_file.readinto(a) + ptx += size + if self.data_file: + # close and delete data file after prefetch so we can pickle + self.data_file.close() + self.data_file = None + + # @lru_cache(maxsize=8) + def __getitem__(self, idx): + if isinstance(idx, int): + i = idx + self.check_index(i) + tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]] + a = np.empty(tensor_size, dtype=self.dtype) + ptx = self.cache_index[i] + np.copyto(a, self.cache[ptx:ptx + a.size]) + return a + elif isinstance(idx, slice): + # Hack just to make this work, can optimizer later if necessary + sents = [] + for i in range(*idx.indices(len(self))): + sents.append(self[i]) + return sents + + +class IndexedDatasetBuilder(object): + + def __init__(self, out_file, dtype=np.int32): + self.out_file = open(out_file, 'wb') + self.dtype = dtype + self.data_offsets = [0] + self.dim_offsets = [0] + self.sizes = [] + self.element_size = self.dtype().itemsize + self.doc_idx = [0] + + def add_item(self, tensor): + bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype)) + self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size) + for s in tensor.size(): + self.sizes.append(s) + self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size())) + + def end_document(self): + self.doc_idx.append(len(self.sizes)) + + def merge_file_(self, another_file): + index = IndexedDataset(another_file) + assert index.dtype == self.dtype + + doc_offset = len(self.sizes) + + begin = self.data_offsets[-1] + for data_offset in index.data_offsets[1:]: + self.data_offsets.append(begin + data_offset) + self.sizes.extend(index.sizes) + begin = self.dim_offsets[-1] + for dim_offset in index.dim_offsets[1:]: + self.dim_offsets.append(begin + dim_offset) + self.doc_idx.extend((doc_offset + index.doc_idx)[1:]) + + with open(data_file_path(another_file), 'rb') as f: + while True: + data = f.read(1024) + if data: + self.out_file.write(data) + else: + break + + def finalize(self, index_file): + self.out_file.close() + index = open(index_file, 'wb') + index.write(b'TNTIDX\x00\x00') + index.write(struct.pack(' [0, 10, 30, 35] + if arr.size > 1: + arr[1:] = arr[:-1] + if arr.size > 0: + arr[0] = 0 + + +def get_pointers_with_total(sizes, elemsize, dtype): + """Return a numpy array of type np.dtype giving the byte offsets. + + Multiplies values in the sizes array by elemsize (bytes), + and then computes an exclusive scan to get byte offsets. + Returns the total number of bytes as second item in a tuple. + """ + + # scale values in sizes array by elemsize to get sizes in bytes + pointers = np.array(sizes, dtype=dtype) + pointers *= elemsize + np.cumsum(pointers, axis=0, out=pointers) + + # get total number of bytes from all sizes (last element) + bytes_last = pointers[-1] if len(sizes) > 0 else 0 + + # convert to byte offsets + exscan_from_cumsum_(pointers) + + return pointers, bytes_last + + +class MMapIndexedDataset(torch.utils.data.Dataset): + + class Index(object): + _HDR_MAGIC = b'MMIDIDX\x00\x00' + + @classmethod + def writer(cls, path, dtype): + + class _Writer(object): + + def __enter__(self): + self._file = open(path, 'wb') + + self._file.write(cls._HDR_MAGIC) + self._file.write(struct.pack('= 0: + if max_value <= 255: + return np.uint8 + elif max_value <= 65535: + return np.uint16 + elif max_value <= 4294967295: + return np.uint32 + else: + return np.uint64 + else: + if max_value <= 127 and min_value >= -128: + return np.int8 + elif max_value <= 32767 and min_value >= -32768: + return np.int16 + elif max_value <= 2147483647 and min_value >= -2147483648: + return np.int32 + else: + return np.int64 + + +def split_index(start_idx, end_idx, num_partitions): + partition_boundaries = np.linspace(start_idx, end_idx, dtype=int, num=num_partitions + 1) + return [(partition_boundaries[i], partition_boundaries[i + 1]) for i in range(num_partitions)] + + +def split_dataset(dataset, num_workers, worker_id, num_threads): + worker_splits = split_index(0, len(dataset), num_workers) + thread_splits = split_index(worker_splits[worker_id][0], worker_splits[worker_id][1], num_threads) + return worker_splits, thread_splits + + +def create_mmap_dataset_builder(fname, dtype): + logger.info(f"Creating mmap dataset builder at {fname}.") + return MMapIndexedDatasetBuilder(f"{fname}.bin", dtype=dtype) + + +def close_mmap_dataset_builder(builder, fname): + builder.end_document() + builder.finalize(f"{fname}.idx") + logger.info(f"Finalized mmap dataset builder at {fname}.") diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py new file mode 100644 index 0000000000000000000000000000000000000000..c9a39bbc53b509479aa2c18f3ddfc2e0abdb7a3c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py @@ -0,0 +1,492 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# support/questions/maintenance: github user @brunomaga or @deepspeedai/deepspeed + +import random +import torch +import os +import numpy as np +from torch.optim.lr_scheduler import LRScheduler +from torch.optim.optimizer import Optimizer +from torch.utils.data import DataLoader, DistributedSampler +from deepspeed.utils import logger +from deepspeed.runtime.pipe.engine import PipelineEngine +from deepspeed.runtime.data_pipeline.constants import * +from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset +from deepspeed.runtime.data_pipeline.data_sampling.data_analyzer import DistributedDataAnalyzer +import pathlib + + +def batch_by_seqlens( + seqlens, + max_tokens, + sequence_ids_per_mb=None, + min_batch_size=1, + max_batch_size=None, + sequence_picking_order="dataloader", + effective_batch_size=1, + required_microbatches_of_same_size=False, + verbose=False, + seed=None, +): + """ + Yield mini-batches of indices bucketed by size. Batches may contain sequences of different lengths. + Similar to "Attention is all you need", Section 5.1: + "sequence pairs were batched together by approximate sequence length. Each training batch + contained a set of sequence pairs containing approximately X source tokens and X target tokens" + + Arguments: + - `seqlens`: a list of difficulties (metric values) for every sample in the dataset; + - `max_tokens`: maximum cap in total difficulty in a batch; + - `min_batch_size`: smallest allowed size of a batch; + - `min_batch_size`: largest allowed size of a batch; + - `sequence_picking_order`: order in which to process samples: "dataloader" (default), "random" or "seqlen" (ascending) + - `effective_batch_size`: effective batch size; + - `required_microbatches_of_same_size`: enable if each mini-batch (in a total of `batch_size_multiple` + micro-batches per batch), should have all micro-batches with the same batch size ie the same + number of sequences. + - `verbose`: print debug information; + - `seed`: random seed for reproducibility; + + Returns: + - `microbatch_ids`: list of tuple of batch id and samples ids per microbatch + - `batch_sizes`: the effective batch size of each batch, used for to compute the scaled LR + - `batch_max_seqlens`: the max seqlen across all microbatches in a batch + """ + + assert sequence_picking_order in ["random", "seqlen", "dataloader"] + if sequence_ids_per_mb is None: + metrics = list(zip(seqlens, range(len(seqlens)))) # use all samples + else: + metrics = list(zip(np.array(seqlens)[sequence_ids_per_mb], sequence_ids_per_mb)) + + if sequence_picking_order == 'random': + metric_random = random.Random(seed) + metric_random.shuffle(metrics) + if sequence_picking_order == 'seqlen': + metrics = sorted(metrics) + + # go through metrics, warn user, and filter samples that alone exceed the max batch threshold + long_ids = [idx for val, idx in metrics if val > max_tokens] + if len(long_ids) > 0: + logger.warning(f"Data indices {long_ids} ignored as metrics exceed {max_tokens}.") + logger.info(f"Original dataset length: {len(metrics)}. New dataset length: {len(long_ids)}") + metrics = [m for m in metrics if m[1] not in long_ids] + + def is_microbatch_valid(metrics): + if min_batch_size and len(metrics) < min_batch_size: return False # insufficient sample count + if max_batch_size and len(metrics) > max_batch_size: return False # too many samples + if sum([m[0] for m in metrics]) > max_tokens: return False # exceeds max + return True + + # go through all samples and pack then in microbatches of metric sums below the threshold + # `required_microbatches_of_same_size` means all minibatches in a batch must be of equal size + equal_size_multiple = effective_batch_size if required_microbatches_of_same_size else 1 + microbatches = [] + batch_init = 0 + while batch_init < len(metrics): + + # we iterate over possible effective batch sizes (groups of microbatches of same size) + valid_batch_end = batch_init + for batch_end in range(batch_init + equal_size_multiple, len(metrics), equal_size_multiple): + + # attempt effective batch + batch = metrics[batch_init:batch_end] + + # pick interleaved samples for each microbatch to help with load balancing + # (in the ordered use case), and to replicate what the distributed sampler does. + mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)] + + # if they are all valid micro-batches, keep them until you find longer mbatches, if any + is_batch_valid = all([is_microbatch_valid(mb) for mb in mbs]) + if is_batch_valid: + valid_batch_end = batch_end + + if batch_init == valid_batch_end: break # last batch is not valid (size zero), so we are done + batch = metrics[batch_init:valid_batch_end] + mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)] + batch_init += sum([len(l) for l in mbs]) + microbatches += mbs + + # make sure we give the same number of (micro-)batches to each dataloader by trimming the dataset + assert len(microbatches) >= effective_batch_size, "not enough datapoints to create a single sample per dataloader" + microbatches = microbatches[:len(microbatches) - len(microbatches) % effective_batch_size] + + #compute the effective batch size for each microbatch. + batch_sizes, batch_max_seqlens, microbatch_ids = [], [], [] + for rank in range(0, len(microbatches), effective_batch_size): + batch_id = rank // effective_batch_size + mbs = microbatches[rank:rank + effective_batch_size] + # compute the number of samples (not tokens) in this batch (not microbatch) + n_sequences = sum([len(mb) for mb in mbs]) + # compute the longest sequence (as number of tokens) in this batch (not microbatch) + sequence_ids_per_mb = [[m[1] for m in metrics] for metrics in mbs] + sequence_lens_per_mb = [[m[0] for m in metrics] for metrics in mbs] + batch_max_seqlen = max([max(seqlens) for seqlens in sequence_lens_per_mb]) + batch_and_mb_ids = zip([batch_id] * effective_batch_size, sequence_ids_per_mb) + batch_sizes.append(n_sequences) + batch_max_seqlens.append(batch_max_seqlen) + microbatch_ids += batch_and_mb_ids + if verbose: + n_tokens_per_mb = [sum([m[0] for m in mb]) for mb in mbs] + n_sequences_per_mb = [len(mb) for mb in mbs] + assert all([n <= max_tokens for n in n_tokens_per_mb]), "size of microbatch exceeds max tokens" + logger.info( + f"Batch id {batch_id} contains in total {len(mbs)} microbatches or {n_sequences} sequences. "\ + f"n_sequences per microbatch {n_sequences_per_mb}. "\ + f"n_tokens per microbatch {n_tokens_per_mb}. "\ + f"sequence ids per microbatch: {sequence_ids_per_mb}. "\ + f"sequence lengths per microbatch: {sequence_lens_per_mb}.") + + # return the sample ids of each microbatch, and the batch sizes + assert len(batch_sizes) == len(microbatch_ids) // effective_batch_size + return microbatch_ids, batch_sizes, batch_max_seqlens + + +def scale_lr(base_batch_size, batch_size, base_lr=1, method="linear"): + """ given a reference lr and batch_size, compute the new LR for a given batch size """ + if method == "linear": + # Linear Scaling Rule: "When the minibatch size is multiplied by k, multiply the learning + # rate by k" (Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, Goyal et al) + return base_lr * batch_size / base_batch_size + if method == "sqrt": + # Square Root scaling: "when multiplying the batch size by k, multiply the learning rate + # by √k, to keep the variance in the gradient expectation constant" + # (A. Krizhevsky. One weird trick for parallelizing convolutional neural networks) + return base_lr * torch.sqrt(batch_size / base_batch_size) + elif method == None or method.upper() == "NONE": + return base_lr + raise ValueError("Unknown scaling method: {}".format(method)) + + +def dataloader_for_variable_batch_size( + dataset, + microbatch_ids, + batch_max_seqlens, + dataloader_rank=0, + dataloader_batch_size=1, + dataloader_num_replicas=1, + dataloader_collate_fn=None, + dataloader_num_workers=2, + dataloader_pin_memory=False, + required_microbatches_of_same_seqlen=False, + sample_padding_fn=None, +): + + # equidistantly distribute the microbatches across the replicas in an interleaved fashion. + sampler = DistributedSampler( + dataset=microbatch_ids, + num_replicas=dataloader_num_replicas, + rank=dataloader_rank, + shuffle=False, + drop_last=False, + ) + + # collate function wraps user-defined collate function to the variable batch data + def collate_fn_wrapper(list_microbatch_ids): + # each batch is a list of sample ids that fill up to the max tokens per batch + # we return the collated batch of all dataset samples of all input batches. + batch = [] + for batch_id, microbatch_ids in list_microbatch_ids: + batch_data = [dataset[idx] for idx in microbatch_ids] + if required_microbatches_of_same_seqlen: + assert sample_padding_fn is not None, \ + "padding dataloader_padding_fn must be provided if required_microbatches_of_same_seqlen is True" + max_seqlen = batch_max_seqlens[batch_id] + assert all([len(sample) <= max_seqlen for sample in batch_data]), \ + "some samples are longer than the computed max seqlen for the batch those samples belong to" + batch_data = [sample_padding_fn(sample, max_seqlen) for sample in batch_data] + batch += batch_data + return dataloader_collate_fn(batch) if dataloader_collate_fn else batch + + dataloader = DataLoader( + dataset=microbatch_ids, + batch_size=dataloader_batch_size, + sampler=sampler, + num_workers=dataloader_num_workers, + collate_fn=collate_fn_wrapper, + pin_memory=dataloader_pin_memory, + ) + + deepspeed_io_kwargs = dict( + dataset=microbatch_ids, + batch_size=dataloader_batch_size, + pin_memory=dataloader_pin_memory, + data_sampler=sampler, + collate_fn=collate_fn_wrapper, + num_local_io_workers=dataloader_num_workers, + ) + + return dataloader, deepspeed_io_kwargs + + +class VariableBatchSizeLR(LRScheduler): + """ an LR scheduler that scales the LR of a given scheduler's LR """ + + @property + def optimizer(self): + return self.base_lr_scheduler.optimizer + + def __init__(self, + lr_scheduler, + base_batch_size, + batch_sizes, + dataloader, + lr_scaling_method="linear", + last_epoch=-1, + verbose=False): + self.batch_sizes = batch_sizes + self.base_batch_size = base_batch_size + self.lr_scaling_method = lr_scaling_method + self.dataloader = dataloader + self.base_lr_scheduler = lr_scheduler + # the following exist in LRScheduler but not in DeepSpeed's LRScheduler so we redefine them here + self.base_lrs = self.base_lr_scheduler.get_lr() + self.last_epoch = last_epoch + self.verbose = verbose + self.step(0) # scale LR for first sample in the dataloader + + def state_dict(self): + return { + 'base_lr_scheduler': self.base_lr_scheduler.state_dict() + } | { + 'base_batch_size': self.base_batch_size, + 'lr_scaling_method': self.lr_scaling_method, + 'batch_sizes': self.batch_sizes, + 'base_lrs': self.base_lrs, + 'last_epoch': self.last_epoch, + 'verbose': self.verbose, + } + + def load_state_dict(self, state_dict): + self.base_lr_scheduler.load_state_dict(state_dict['base_lr_scheduler']) + self.base_batch_size = state_dict['base_batch_size'] + self.lr_scaling_method = state_dict['lr_scaling_method'] + self.batch_sizes = state_dict['batch_sizes'] + self.base_lrs = state_dict['base_lrs'] + self.last_epoch = state_dict['last_epoch'] + self.verbose = state_dict['verbose'] + + def get_last_lr(self): + return self.base_lr_scheduler._last_lr + + def get_lr(self): + return [group['lr'] for group in self.base_lr_scheduler.optimizer.param_groups] + + def step(self, epoch=None): + # call the base scheduler's step method to get LR for next epoch + # Note: optimizer.step precedes lr_scheduler.step(), so the stepping workflow is: + # init: lr_scheduler.step(0) --> set LR for epoch 0 + # epoch 0: optimizer.step(); lr_scheduler.step(1) --> set LR for epoch 1 + # epoch 1: optimizer.step(); lr_scheduler.step(2) --> set LR for epoch 2 + + # reset unscaled LRs (to the original scheduler's one) to be able to step the base LR scheduler + # Note: epoch==0: reset LR scheduler; epoch==None: scale LR for next epoch; + unscaled_lrs = self.base_lrs if epoch == 0 else self.get_last_lr() + for group, lr in zip(self.base_lr_scheduler.optimizer.param_groups, unscaled_lrs): + group['lr'] = lr + + self.base_lr_scheduler.step(epoch) # set unscaled lr, _step_count, last_epoch, _last_lr for new epoch + + # scale the learning rate for the the next iteration for each parameter group. + self.last_epoch = self.last_epoch + 1 if epoch is None else epoch + # batch sizes are precomputed and stored in batch_sizes se we loop around to get the next one + batch_size = self.batch_sizes[self.last_epoch % len(self.batch_sizes)] + for group in self.base_lr_scheduler.optimizer.param_groups: + group['lr'] = scale_lr(self.base_batch_size, batch_size, group['lr'], self.lr_scaling_method) + + if self.verbose: + logger.info( + f"Next batch id {self.last_epoch}. "\ + f"Reference batch_size {self.base_batch_size} and lr {unscaled_lrs}. "\ + f"Scaled batch_size {batch_size} and lr {self.get_lr()}.") + + +def lr_scheduler_for_variable_batch_size(base_batch_size, + batch_sizes, + dataloader, + lr_scheduler_or_optimizer, + lr_scaling_method='linear', + verbose=False): + """ + returns a class that provides an LR scheduler that scales the learning rate at every + iteration taking into account the batch size of that iteration. + If learning rate is constant, ie no LR scheduler, then the base LR will be taken from the + constant LR values in the optimizer param groups. Otherwise from the scheduler's LR. + + Arguments: + - `base_batch_size`: the batch size that the base LR in the optimizer or scheduler refers to; + - `lr_scaling_method`: method to use to scale LR - see `scale_lr()`; + - `lr_scheduler_or_optimizer`: one instance of `LRScheduler` or `Optimizer` to be used as base; + - `batch_sizes`: the effective batch size of each batch in the dataloader; + + Returns the new LRScheduler + """ + + class StubLRScheduler(LRScheduler): + """ a stub LR scheduler that does not change the LR, keeps it constant """ + + def get_lr(self) -> float: + return self.base_lrs + + if isinstance(lr_scheduler_or_optimizer, Optimizer): + lr_scheduler = StubLRScheduler(lr_scheduler_or_optimizer) + elif hasattr(lr_scheduler_or_optimizer, 'optimizer'): #LRScheduler or DeepSpeed 'object' schedulers + assert isinstance(lr_scheduler_or_optimizer.optimizer, Optimizer) + lr_scheduler = lr_scheduler_or_optimizer + else: + raise ValueError("Unknown type for lr_scheduler_or_optimizer: {}".format(type(lr_scheduler_or_optimizer))) + + return VariableBatchSizeLR(lr_scheduler=lr_scheduler, + base_batch_size=base_batch_size, + batch_sizes=batch_sizes, + dataloader=dataloader, + lr_scaling_method=lr_scaling_method, + verbose=verbose) + + +def get_dataloader_and_lr_scheduler_for_variable_batch_size_deepspeed(dataset, + engine, + dataset_seqlens=None, + dataset_filter_ids=None, + dataloader_collate_fn=None, + sample_padding_fn=None, + batch_seqlens_fn=None): + """ + a simplified call to get_dataloader_and_lr_scheduler_for_variable_batch_size for the deepspeed runtime. + Needs the seqlens of every sample. It will try three alternatives: + - if `dataset_seqlens` is provided by user, use that. + - otherwise, looks for the seqlen metric path (in the connfig) that contains the output of the Data Analyzer + - otherwise, use the user-provided function `batch_seqlens_fn` and call Data Analyzer to output seqlen metric + See `batch_by_seqlens()` for arguments and more documentation. + """ + data_efficiency_config = engine._config.data_efficiency_config + data_sampling_config = data_efficiency_config[DATA_SAMPLING] + batching_config = data_sampling_config[DYNAMIC_BATCHING] + assert batching_config[DYNAMIC_BATCHING_ENABLED], "Dynamic batching is not enabled in the config" + + if dataset_seqlens is None: + # In seqlen provided by user, look for the seqlen metric that was output by the Data Analyzer + # (see the main in deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py for an example) + metrics_path = batching_config[DYNAMIC_BATCHING_METRICS_PATH] + sample_to_seqlen_path = os.path.join(metrics_path, "seqlen/seqlen_sample_to_metric") + if not (os.path.exists(f"{sample_to_seqlen_path}.bin") and os.path.exists(f"{sample_to_seqlen_path}.idx")): + # if the metric files are not found, we run the DataAnalyzer to write the metric files + msg = f"Cannot find metric files for sequence length in {sample_to_seqlen_path}.idx or *.bin." + msg += " We will run data analyzer to generated them..." + logger.warning(msg) + + if batch_seqlens_fn is None: + raise ValueError("sample_seqlen_fn must be provided if dataset_seqlens is not provided") + + DistributedDataAnalyzer( + dataset=dataset, + metric_functions=[batch_seqlens_fn], + collate_fn=dataloader_collate_fn, + batch_size=2**10, # batch size for map-reduce, not training + num_workers=engine.world_size, + worker_id=engine.global_rank, + save_path=pathlib.Path(metrics_path), + metric_types=['single_value_per_sample'], + metric_names=["seqlen"], + device=engine.device, + ).run_map_reduce() + + dataset_seqlens = MMapIndexedDataset(sample_to_seqlen_path, skip_warmup=True) + assert len(dataset_seqlens) == len(dataset), \ + "Seqlens size does not match the input dataset size. If you changed the dataset, delete the metrics_path folder." + + # TODO we are copying all seqlens into memory, we should adapt the code to use an iterative streamer + # and use the other files output by DataAnalyzer that returns an ordered dictionary of seqlen to sample ids + dataset_seqlens = np.array(list(dataset_seqlens), dtype=np.int64).flatten() # from Nx1 to N + + dataloader, lr_scheduler, deepspeed_io_kwargs = get_dataloader_and_lr_scheduler_for_variable_batch_size( + dataset=dataset, + dataset_filter_ids=dataset_filter_ids, + dataset_seqlens=dataset_seqlens, + effective_batch_size=engine.train_batch_size(), + max_tokens=batching_config[DYNAMIC_BATCHING_MAX_TOKENS], + lr_scaling_method=batching_config[DYNAMIC_BATCHING_LR_SCALING_METHOD], + sequence_picking_order=batching_config[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER], + min_batch_size=batching_config[DYNAMIC_BATCHING_MIN_BATCH_SIZE], + max_batch_size=batching_config[DYNAMIC_BATCHING_MAX_BATCH_SIZE], + dataloader_batch_size=engine.train_micro_batch_size_per_gpu(), + dataloader_rank=engine.data_parallel_group.rank(), + dataloader_num_replicas=engine.data_parallel_group.size(), + dataloader_num_workers=data_sampling_config[DATA_SAMPLING_NUM_WORKERS], + dataloader_collate_fn=dataloader_collate_fn, + dataloader_pin_memory=data_sampling_config[DATA_SAMPLING_PIN_MEMORY], + sample_padding_fn=sample_padding_fn, + lr_scheduler_or_optimizer=engine.lr_scheduler or engine.optimizer, + required_microbatches_of_same_size=isinstance(engine, PipelineEngine), + required_microbatches_of_same_seqlen=isinstance(engine, PipelineEngine), + verbose=batching_config[DYNAMIC_BATCHING_VERBOSE], + seed=data_efficiency_config[DATA_EFFICIENCY_SEED], + ) + return dataloader, lr_scheduler, deepspeed_io_kwargs + + +def get_dataloader_and_lr_scheduler_for_variable_batch_size( + dataset, + dataset_seqlens, + max_tokens, + effective_batch_size, + dataset_filter_ids=None, + lr_scaling_method="linear", + min_batch_size=1, + max_batch_size=None, + sequence_picking_order="dataloader", + dataloader_batch_size=1, + dataloader_rank=0, + dataloader_num_replicas=1, + dataloader_num_workers=0, + dataloader_collate_fn=None, + dataloader_pin_memory=False, + lr_scheduler_or_optimizer=None, + required_microbatches_of_same_size=False, + required_microbatches_of_same_seqlen=False, + sample_padding_fn=None, + verbose=False, + seed=None, +): + """ returns a dataloader and LR scheduler for the variable batch size. see `batch_by_seqlens()` for details. """ + + # effective_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation_steps * number of dataloaders + microbatch_ids, batch_sizes, batch_max_seqlens = batch_by_seqlens( + seqlens=dataset_seqlens, + max_tokens=max_tokens, + sequence_ids_per_mb=dataset_filter_ids, + min_batch_size=min_batch_size, + max_batch_size=max_batch_size, + sequence_picking_order=sequence_picking_order, + effective_batch_size=effective_batch_size, + required_microbatches_of_same_size=required_microbatches_of_same_size, + verbose=verbose, + seed=seed, + ) + + dataloader, deepspeed_io_kwargs = dataloader_for_variable_batch_size( + dataset=dataset, + microbatch_ids=microbatch_ids, + batch_max_seqlens=batch_max_seqlens, + dataloader_rank=dataloader_rank, + dataloader_num_replicas=dataloader_num_replicas, + dataloader_batch_size=dataloader_batch_size, + dataloader_collate_fn=dataloader_collate_fn, + dataloader_num_workers=dataloader_num_workers, + dataloader_pin_memory=dataloader_pin_memory, + required_microbatches_of_same_seqlen=required_microbatches_of_same_seqlen, + sample_padding_fn=sample_padding_fn, + ) + + lr_scheduler = lr_scheduler_for_variable_batch_size(base_batch_size=effective_batch_size, + batch_sizes=batch_sizes, + lr_scaling_method=lr_scaling_method, + lr_scheduler_or_optimizer=lr_scheduler_or_optimizer, + dataloader=dataloader, + verbose=verbose) + + return dataloader, lr_scheduler, deepspeed_io_kwargs diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/dataloader.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..499473b4ced81ba5ae5a447e32e03a2971b63c8f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/dataloader.py @@ -0,0 +1,162 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch.utils.data import DataLoader, RandomSampler +from torch.utils.data.distributed import DistributedSampler +from deepspeed.accelerator import get_accelerator + +from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler +from deepspeed.runtime.data_pipeline.constants import CURRICULUM_LEARNING, \ + DATA_EFFICIENCY, DATA_SAMPLING_NUM_WORKERS +from deepspeed.runtime.constants import GRADIENT_ACCUMULATION_STEPS, \ + DATA_PARALLEL_GROUP, GLOBAL_RANK + + +class RepeatingLoader: + + def __init__(self, loader): + """Wraps an iterator to allow for infinite iteration. This is especially useful + for DataLoader types that we wish to automatically restart upon completion. + + Args: + loader (iterator): The data loader to repeat. + """ + self.loader = loader + self.data_iter = iter(self.loader) + + def __iter__(self): + return self + + def __next__(self): + try: + batch = next(self.data_iter) + except StopIteration: + self.data_iter = iter(self.loader) + batch = next(self.data_iter) + return batch + + +class DeepSpeedDataLoader(object): + + def __init__(self, + dataset, + batch_size, + pin_memory, + local_rank, + tput_timer, + collate_fn=None, + num_local_io_workers=None, + data_sampler=None, + data_parallel_world_size=None, + data_parallel_rank=None, + dataloader_drop_last=False, + deepspeed_dataloader_config={}): + self.deepspeed_dataloader_config = deepspeed_dataloader_config + self.tput_timer = tput_timer + self.batch_size = batch_size + self.curriculum_learning_enabled = False + if CURRICULUM_LEARNING in deepspeed_dataloader_config: + self.curriculum_learning_enabled = deepspeed_dataloader_config[CURRICULUM_LEARNING] + + if self.curriculum_learning_enabled: + data_sampler = DeepSpeedDataSampler(self.deepspeed_dataloader_config[DATA_EFFICIENCY], + len(dataset), + self.batch_size, + data_parallel_rank, + data_parallel_world_size, + self.deepspeed_dataloader_config[DATA_PARALLEL_GROUP], + self.deepspeed_dataloader_config[GRADIENT_ACCUMULATION_STEPS], + self.deepspeed_dataloader_config[GLOBAL_RANK], + drop_last=dataloader_drop_last) + device_count = get_accelerator().device_count() + num_local_io_workers = self.deepspeed_dataloader_config[DATA_SAMPLING_NUM_WORKERS] + else: + if local_rank >= 0: + if data_sampler is None: + data_sampler = DistributedSampler(dataset=dataset, + num_replicas=data_parallel_world_size, + rank=data_parallel_rank) + device_count = 1 + else: + if data_sampler is None: + data_sampler = RandomSampler(dataset) + device_count = get_accelerator().device_count() + batch_size *= device_count + + if num_local_io_workers is None: + num_local_io_workers = 2 * device_count + + self.num_local_io_workers = num_local_io_workers + self.data_sampler = data_sampler + self.dataset = dataset + self.collate_fn = collate_fn + self.device_count = device_count + self.batch_size = batch_size + self.pin_memory = pin_memory + self.data = None + self.dataloader_drop_last = dataloader_drop_last + self.post_process_func = None + + if self.dataloader_drop_last: + self.len = len(self.data_sampler) // self.batch_size + else: + from math import ceil + self.len = ceil(len(self.data_sampler) / self.batch_size) + + def __iter__(self): + self._create_dataloader() + return self + + def __len__(self): + return self.len + + def __next__(self): + if self.tput_timer: + self.tput_timer.start() + if self.curriculum_learning_enabled: + data = next(self.data_iterator) + if self.post_process_func is not None: + data = self.post_process_func(data, self.data_sampler.state_dict()) + return data + else: + return next(self.data) + + def _create_dataloader(self): + if self.curriculum_learning_enabled: + if self.collate_fn is None: + self.dataloader = DataLoader(self.dataset, + pin_memory=self.pin_memory, + batch_sampler=self.data_sampler, + num_workers=self.num_local_io_workers) + else: + self.dataloader = DataLoader(self.dataset, + pin_memory=self.pin_memory, + batch_sampler=self.data_sampler, + collate_fn=self.collate_fn, + num_workers=self.num_local_io_workers) + self.data_iterator = iter(self.dataloader) + return self.dataloader + else: + if self.collate_fn is None: + self.dataloader = DataLoader(self.dataset, + batch_size=self.batch_size, + pin_memory=self.pin_memory, + sampler=self.data_sampler, + num_workers=self.num_local_io_workers, + drop_last=self.dataloader_drop_last) + else: + self.dataloader = DataLoader(self.dataset, + batch_size=self.batch_size, + pin_memory=self.pin_memory, + sampler=self.data_sampler, + collate_fn=self.collate_fn, + num_workers=self.num_local_io_workers, + drop_last=self.dataloader_drop_last) + self.data = (x for x in self.dataloader) + + return self.dataloader + + +# DataLoader([(torch.randn(3, 3), torch.tensor(i % 2)) for i in range(10)], batch_size=2)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84896d3405ff4c4ffa848a370325c25006f66a8a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e4848577a040c5205a63d6d9ed99b2c38d8e898 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fde5ed460d63794267dd092b8dfeeb92fc4342c0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/async_linear.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/async_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..8e01da500409cd603c17af646e902cfe2a82525b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/async_linear.py @@ -0,0 +1,140 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Adapted from https://github.com/NVIDIA/Megatron-LM/blob/23.08/megatron/core/tensor_parallel/layers.py + +import torch +from torch.nn.parameter import Parameter +import torch.nn.functional as F +from deepspeed.accelerator import get_accelerator +import deepspeed.comm as dist +from typing import Callable + +TP_group = None + + +class DominoAsyncColumnParallelLinearImpl(torch.autograd.Function): + + @staticmethod + def forward(ctx, inp, weight, bias, handle_dic, h_id): # inp: (b, s, k), weight: (m, k), bias (m) + ctx.save_for_backward(inp, weight, bias) + ctx.handle_dic = handle_dic + ctx.h_id = h_id + output = torch.matmul(inp, weight.t()) # (b, s, k) @ (k, m) -> (b, s, m) + if bias is not None: # bias (m) + output = output + bias + return output + + @staticmethod + def backward(ctx, grad_output): + inp, weight, bias = ctx.saved_tensors + grad_input = grad_weight = grad_bias = None + grad_input = torch.matmul(grad_output, weight) # (b, s, m) @ (m, k) -> (b, s, k) + handle = dist.all_reduce(grad_input, group=TP_group, async_op=True) + ctx.handle_dic[ctx.h_id] = handle + grad_output = grad_output.view(grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2]) # (b*s, m) + + inp = inp.view(inp.shape[0] * inp.shape[1], inp.shape[2]) # (b*s, k) + grad_weight = torch.matmul(grad_output.t(), inp) # (m, b*s) @ (b*s, k) -> (m, k) + + if bias is not None: + grad_bias = grad_output.sum(dim=0) # (b*s, m) -> (m) + return grad_input, grad_weight, grad_bias, None, None + + +class DominoAsyncColumnParallelLinear(torch.nn.Module): + + def __init__(self, + input_size, + output_size, + _tp_group, + config, + init_method: Callable, + bias=True, + skip_bias_add=False): + super(DominoAsyncColumnParallelLinear, self).__init__() + + self.skip_bias_add = skip_bias_add + + global TP_group + if TP_group == None: + TP_group = _tp_group + + self.weight = Parameter( + torch.empty( + output_size, + input_size, + device=get_accelerator().current_device_name(), + dtype=config.params_dtype, + )) + if config.perform_initialization: + init_method(self.weight) + + if bias: + self.bias = Parameter( + torch.empty(output_size, device=get_accelerator().current_device_name(), dtype=config.params_dtype)) + + if config.perform_initialization: + with torch.no_grad(): + self.bias.zero_() + else: + self.register_parameter('bias', None) + + def forward(self, input_: torch.Tensor, handle_dic, h_id): + + bias = self.bias if not self.skip_bias_add else None + + output = DominoAsyncColumnParallelLinearImpl.apply(input_, self.weight, bias, handle_dic, h_id) + + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + +class RowParallelLinearNoComm(torch.nn.Module): + + def __init__( + self, + input_size: int, + output_size: int, + config, + init_method: Callable, + bias: bool = True, + stride: int = 1, + skip_bias_add: bool = False, + ): + super(RowParallelLinearNoComm, self).__init__() + + self.skip_bias_add = skip_bias_add + + self.weight = Parameter( + torch.empty( + output_size, + input_size, + device=get_accelerator().current_device_name(), + dtype=config.params_dtype, + )) + if config.perform_initialization: + init_method(self.weight) + if bias: + self.bias = Parameter( + torch.empty( + output_size, + device=get_accelerator().current_device_name(), + dtype=config.params_dtype, + )) + + if config.perform_initialization: + with torch.no_grad(): + self.bias.zero_() + else: + self.register_parameter('bias', None) + + def forward(self, input_): + bias = self.bias if not self.skip_bias_add else None + + output = F.linear(input_, self.weight, bias) + + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/transformer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..3dfb133373b5aa7686f1dd9aac903add10b4e84d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/transformer.py @@ -0,0 +1,605 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +import enum +import deepspeed.comm as dist + +from .async_linear import DominoAsyncColumnParallelLinear, RowParallelLinearNoComm + + +class LayerType(enum.Enum): + encoder = 1 + decoder = 2 + + +class AttnType(enum.Enum): + self_attn = 1 + cross_attn = 2 + + +class AttnMaskType(enum.Enum): + padding = 1 + causal = 2 + + +class ModelType(enum.Enum): + encoder_or_decoder = 1 + encoder_and_decoder = 2 + + +class DominoUtil: + + BATCH_0 = "BATCH0" + + BATCH_1 = "BATCH1" + + HANDLE_DIC = {"BATCH0": None, "BATCH1": None} + + +class DominoModule(torch.nn.Module): + """extensions of torch Module.""" + + def __init__(self, ): + super(DominoModule, self).__init__() + + +def _Wait_bwd_comm(input_, dic_, h_id): + return NoOper.apply(input_, dic_, h_id) + + +class NoOper(torch.autograd.Function): + + @staticmethod + def symbolic(graph, input_, handle_dic, h_id): + return input_ + + @staticmethod + def forward(ctx, input_, handle_dic, h_id): + ctx.handle_dic = handle_dic + ctx.h_id = h_id + return input_ + + @staticmethod + def backward(ctx, grad_output): + handle = ctx.handle_dic[ctx.h_id] + handle.wait() + return grad_output, None, None + + +class CoreAttention(DominoModule): + + def __init__(self, config, tp_world_size, attn_mask_type=AttnMaskType.causal): + super(CoreAttention, self).__init__() + + self.attn_mask_type = attn_mask_type + + projection_size = config.kv_channels * config.num_attention_heads + + # Per attention head and per partition values. + assert projection_size % tp_world_size == 0, f"projection size {projection_size} should be multiple of TP world size {tp_world_size}" + self.hidden_size_per_partition = projection_size // tp_world_size + self.attention_dropout_rate = config.attention_dropout + + def forward(self, query_layer, key_layer, value_layer, attention_mask): + + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, + key_layer, + value_layer, + attn_mask=None, + dropout_p=self.attention_dropout_rate, + is_causal=True, + scale=None) + + # [b, np, sq, hn] --> [sq, b, np, hn] + context_layer = context_layer.permute(2, 0, 1, 3).contiguous() + + # [sq, b, np, hn] --> [sq, b, hp] + new_context_layer_shape = context_layer.size()[:-2] + \ + (self.hidden_size_per_partition,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer + + +class ShardedAttention(DominoModule): + """Sharded self-attention layer class. + Only support self attention and causal attention mask for now. + """ + + def __init__(self, + config, + mpu, + apply_rotary_pos_emb, + layer_number, + attention_type=AttnType.self_attn, + attn_mask_type=AttnMaskType.causal): + super(ShardedAttention, self).__init__() + + assert attention_type == AttnType.self_attn, "Only support self_attn for now!" + + self.layer_number = max(1, layer_number) + self.attention_type = attention_type + self.attn_mask_type = attn_mask_type + self.params_dtype = config.params_dtype + self.apply_rotary_pos_emb = apply_rotary_pos_emb + + query_projection_size = config.kv_channels * config.num_attention_heads + kv_projection_size = config.kv_channels * config.num_attention_heads + + tp_world_size = mpu.get_tensor_model_parallel_world_size() + self.hidden_size_per_attention_head = query_projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads // tp_world_size + + qkv_projection_per_partition = (query_projection_size + 2 * kv_projection_size) // tp_world_size + + self.query_key_value = DominoAsyncColumnParallelLinear(config.hidden_size, + qkv_projection_per_partition, + mpu.get_tensor_model_parallel_group(), + config=config, + init_method=config.init_method, + bias=config.add_bias_linear) + + self.core_attention = CoreAttention(config, tp_world_size, self.attn_mask_type) + + query_projection_size_per_partition = query_projection_size // tp_world_size + + # Output. + self.dense = RowParallelLinearNoComm(query_projection_size_per_partition, + config.hidden_size, + config=config, + init_method=config.output_layer_init_method, + bias=config.add_bias_linear, + skip_bias_add=True) + + def forward(self, hidden_states, attention_mask, micro_batch_num, rotary_pos_emb=None): + # hidden_states: [sq, b, h] + + mixed_x_layer, _ = self.query_key_value(hidden_states, DominoUtil.HANDLE_DIC, micro_batch_num) + + new_tensor_shape = mixed_x_layer.size()[:-1] + ( + self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head, + ) + + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + mixed_x_layer = mixed_x_layer.permute(1, 2, 0, 3).contiguous() + + (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, [ + self.hidden_size_per_attention_head, self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head + ], + dim=3) + + query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, + self.hidden_size_per_attention_head) + + if rotary_pos_emb is not None: + if isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = rotary_pos_emb + else: + rotary_pos_emb = ((rotary_pos_emb, ) * 2) + q_pos_emb, k_pos_emb = rotary_pos_emb + query_layer = self.apply_rotary_pos_emb(query_layer, q_pos_emb) + key_layer = self.apply_rotary_pos_emb(key_layer, k_pos_emb) + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + output, bias = self.dense(context_layer) + return output, bias + + def domino_core_attention_forward(self, mixed_x_layer, attention_mask, rotary_pos_emb=None): + # hidden_states: [sq, b, h] + + # To illustrate the difference between intra-layer overlap and inter-layer overlap + # mixed_x_layer, _ = self.query_key_value(hidden_states, handle_dic, micro_batch_num) + + new_tensor_shape = mixed_x_layer.size()[:-1] + ( + self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head, + ) + + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + mixed_x_layer = mixed_x_layer.permute(1, 2, 0, 3).contiguous() + + (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, [ + self.hidden_size_per_attention_head, self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head + ], + dim=3) + + query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, + self.hidden_size_per_attention_head) + + if rotary_pos_emb is not None: + if isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = rotary_pos_emb + else: + rotary_pos_emb = ((rotary_pos_emb, ) * 2) + q_pos_emb, k_pos_emb = rotary_pos_emb + query_layer = self.apply_rotary_pos_emb(query_layer, q_pos_emb) + key_layer = self.apply_rotary_pos_emb(key_layer, k_pos_emb) + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + # output, bias = self.dense(context_layer) + # return output, bias + + return context_layer + + +class bias_dropout_add(torch.nn.Module): + + def __init__(self, prob: float): + super(bias_dropout_add, self).__init__() + self.dropout = torch.nn.Dropout(prob) + + def forward(self, x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + if bias is not None: + x = x + bias + out = self.dropout(x) + out = out + residual + return out + + +class DominoTransformerLayer(DominoModule): + """A domino single transformer layer. + [s, b, h] -> [s, b, h] + """ + + def __init__(self, + config, + mpu, + apply_rotary_pos_emb, + layer_number, + layer_type=LayerType.encoder, + self_attn_mask_type=AttnMaskType.causal, + drop_path_rate=0.): + + super(DominoTransformerLayer, self).__init__() + self.layer_number = layer_number + self.layer_type = layer_type + + self.apply_residual_connection_post_layernorm \ + = config.apply_residual_connection_post_layernorm + + self.llama_model = False + + self.input_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) + + # Self attention. + self.self_attention = ShardedAttention(config, + mpu, + apply_rotary_pos_emb, + layer_number, + attention_type=AttnType.self_attn, + attn_mask_type=self_attn_mask_type) + + self.hidden_dropout = config.hidden_dropout + + self.post_attention_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) + + # MLP + ffn_hidden_size = config.ffn_hidden_size + if config.gated_linear_unit: + ffn_hidden_size *= 2 + + self.output_size_c = config.ffn_hidden_size + self.input_size_c = config.hidden_size + self.input_size_r = config.ffn_hidden_size + self.output_size_r = self.input_size_c + + tp_world_size = mpu.get_tensor_model_parallel_world_size() + self.TP_group = mpu.get_tensor_model_parallel_group() + self.output_size_per_partition = self.output_size_c // tp_world_size + self.input_size_per_partition = self.input_size_r // tp_world_size + + self.linear_fc1 = DominoAsyncColumnParallelLinear(self.input_size_c, + self.output_size_per_partition, + mpu.get_tensor_model_parallel_group(), + config=config, + init_method=config.init_method, + bias=config.add_bias_linear) + + self.mlp_activation_func = F.gelu + + self.linear_fc2 = RowParallelLinearNoComm(self.input_size_per_partition, + self.output_size_r, + config=config, + init_method=config.output_layer_init_method, + bias=config.add_bias_linear, + skip_bias_add=True) + + self.bias_dropout_add_func = bias_dropout_add(self.hidden_dropout) + + def forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + + hidden_states0, hidden_states1 = hidden_states + + layernorm_output0 = self.input_layernorm(hidden_states0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + # Micro batch 0: attention + attention_output0, attention_bias0 = self.self_attention(layernorm_output0, + attention_mask, + DominoUtil.BATCH_0, + rotary_pos_emb=rotary_pos_emb) + + fwd_handle0 = dist.all_reduce(attention_output0, group=self.TP_group, async_op=True) + # End of Micro batch 0: attention + + # Micro batch 1: attention + layernorm_output1 = self.input_layernorm(hidden_states1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + + attention_output1, attention_bias1 = self.self_attention(layernorm_output1, + attention_mask, + DominoUtil.BATCH_1, + rotary_pos_emb=rotary_pos_emb) + fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True) + + # Micro batch 0: Residual connection. + fwd_handle0.wait() + if self.apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = hidden_states0 + + layernorm_input0 = self.bias_dropout_add_func(attention_output0, attention_bias0, residual0) + + layernorm_output0 = self.post_attention_layernorm(layernorm_input0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + if self.apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = layernorm_input0 + # End of Micro batch 0: Residual connection. + + # ------------ MLP ------------ + # Micro batch 0: MLP + output0, _ = self.linear_fc1(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + output0 = self.mlp_activation_func(output0) + + # Micro batch 1: Residual connection. + fwd_handle1.wait() + if self.apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = hidden_states1 + + layernorm_input1 = self.bias_dropout_add_func(attention_output1, attention_bias1, residual1) + + layernorm_output1 = self.post_attention_layernorm(layernorm_input1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + + if self.apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = layernorm_input1 + # End of Micro batch 1: Residual connection. + + hidden_states0, last_mlp_bias = self.linear_fc2(output0) + fwd_handle0 = dist.all_reduce(hidden_states0, group=self.TP_group, async_op=True) + # End of Micro batch 0: MLP + + # Micro batch 1: MLP + output1, _ = self.linear_fc1(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + output1 = self.mlp_activation_func(output1) + + hidden_states1, last_mlp_bias = self.linear_fc2(output1) + + fwd_handle1 = dist.all_reduce(hidden_states1, group=self.TP_group, async_op=True) + # End of Micro batch 1: MLP + + # ------------ End of MLP ------------ + + fwd_handle0.wait() + hidden_states0 = self.bias_dropout_add_func(hidden_states0, last_mlp_bias, residual0) + + fwd_handle1.wait() + hidden_states1 = self.bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) + + return hidden_states0, hidden_states1 + + +class DominoTransformer(DominoModule): + """Transformer class.""" + + def __init__(self, + config, + mpu, + apply_rotary_pos_emb, + model_type, + layer_type=LayerType.encoder, + self_attn_mask_type=AttnMaskType.causal, + post_layer_norm=True, + pre_process=True, + post_process=True, + drop_path_rate=0.0): + super(DominoTransformer, self).__init__() + + self.layer_type = layer_type + self.model_type = model_type + self.post_layer_norm = post_layer_norm + self.post_process = post_process + self.input_tensor = None + self.drop_path_rate = drop_path_rate + self.TP_group = mpu.get_tensor_model_parallel_group() + + if not dist.is_initialized(): + dist.init_distributed() + assert dist.is_initialized(), "deepspeed.comm failed to initialize!" + + self.num_layers = config.num_layers + + self.drop_path_rates = [rate.item() for rate in torch.linspace(0, self.drop_path_rate, config.num_layers)] + + def build_layer(layer_number): + + current_layer_type = layer_type + return DominoTransformerLayer(config, + mpu, + apply_rotary_pos_emb, + layer_number, + layer_type=current_layer_type, + self_attn_mask_type=self_attn_mask_type, + drop_path_rate=self.drop_path_rates[layer_number - 1]) + + self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) + + if self.post_process and self.post_layer_norm: + self.final_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) + + self._forward_impl = self.inter_layer_overlap_forward + if config.domino_intra_layer_overlap: + self._forward_impl = self.intra_layer_overlap_forward + + def forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + + return self._forward_impl(hidden_states, attention_mask, rotary_pos_emb) + + def inter_layer_overlap_forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + # hidden_states: [s, b, h] + + hidden_states0, hidden_states1 = torch.chunk(hidden_states, chunks=2, dim=1) + + last_mlp_bias = None + fwd_handle0, fwd_handle1 = None, None + residual0, residual1 = None, None + + layernorm_output0 = self.layers[0].input_layernorm(hidden_states0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + for index in range(self.num_layers): + + # Micro batch 0: attention + attention_output0, _ = self.layers[index].self_attention.query_key_value( + layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + attention_output0 = self.layers[index].self_attention.domino_core_attention_forward( + attention_output0, attention_mask, rotary_pos_emb=rotary_pos_emb) + + # Micro batch 1: Residual connection + if index > 0: + fwd_handle1.wait() + hidden_states1 = self.layers[index - 1].bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) + + layernorm_output1 = self.layers[index].input_layernorm(hidden_states1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + # End of Micro batch 1: Residual connection + + attention_output0, attention_bias0 = self.layers[index].self_attention.dense(attention_output0) + + fwd_handle0 = dist.all_reduce(attention_output0, group=self.TP_group, async_op=True) + # End of Micro batch 0: attention + + # Micro batch 1: attention + attention_output1, _ = self.layers[index].self_attention.query_key_value( + layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + attention_output1 = self.layers[index].self_attention.domino_core_attention_forward( + attention_output1, attention_mask, rotary_pos_emb=rotary_pos_emb) + + # Micro batch 0: Residual connection. + fwd_handle0.wait() + if self.layers[index].apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = hidden_states0 + + layernorm_input0 = self.layers[index].bias_dropout_add_func(attention_output0, attention_bias0, residual0) + + layernorm_output0 = self.layers[index].post_attention_layernorm(layernorm_input0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + if self.layers[index].apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = layernorm_input0 + # End of Micro batch 0: Residual connection. + + attention_output1, attention_bias1 = self.layers[index].self_attention.dense(attention_output1) + fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True) + # End of Micro batch 1: attention + + # ------------ MLP ------------ + # Micro batch 0: MLP + output0, _ = self.layers[index].linear_fc1(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + output0 = self.layers[index].mlp_activation_func(output0) + + # Micro batch 1: Residual connection. + fwd_handle1.wait() + if self.layers[index].apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = hidden_states1 + + layernorm_input1 = self.layers[index].bias_dropout_add_func(attention_output1, attention_bias1, residual1) + + layernorm_output1 = self.layers[index].post_attention_layernorm(layernorm_input1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + + if self.layers[index].apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = layernorm_input1 + # End of Micro batch 1: Residual connection. + + hidden_states0, last_mlp_bias = self.layers[index].linear_fc2(output0) + fwd_handle0 = dist.all_reduce(hidden_states0, group=self.TP_group, async_op=True) + # End of Micro batch 0: MLP + + # Micro batch 1: MLP + output1, _ = self.layers[index].linear_fc1(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + output1 = self.layers[index].mlp_activation_func(output1) + + # Micro batch 0: Residual connection. + fwd_handle0.wait() + hidden_states0 = self.layers[index].bias_dropout_add_func(hidden_states0, last_mlp_bias, residual0) + + if index < self.num_layers - 1: + layernorm_output0 = self.layers[index + 1].input_layernorm(hidden_states0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + # End of Micro batch 0: Residual connection. + + hidden_states1, last_mlp_bias = self.layers[index].linear_fc2(output1) + + fwd_handle1 = dist.all_reduce(hidden_states1, group=self.TP_group, async_op=True) + # End of Micro batch 1: MLP + + # ------------ End of MLP ------------ + + if self.post_process and self.post_layer_norm: + hidden_states0 = self.final_layernorm(hidden_states0) + + index = self.num_layers - 1 + + fwd_handle1.wait() + hidden_states1 = self.layers[index].bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) + + if self.post_process and self.post_layer_norm: + hidden_states1 = self.final_layernorm(hidden_states1) + + hidden_states = torch.cat([hidden_states0, hidden_states1], dim=1) + + return hidden_states + + def intra_layer_overlap_forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + + hidden_states = torch.chunk(hidden_states, chunks=2, dim=1) + + for index in range(self.num_layers): + layer = self.layers[index] + hidden_states = layer(hidden_states, attention_mask, rotary_pos_emb) + + hidden_states0, hidden_states1 = hidden_states + if self.post_process and self.post_layer_norm: + hidden_states0 = self.final_layernorm(hidden_states0) + hidden_states1 = self.final_layernorm(hidden_states1) + + hidden_states = torch.cat([hidden_states0, hidden_states1], dim=1) + return hidden_states diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/eigenvalue.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/eigenvalue.py new file mode 100644 index 0000000000000000000000000000000000000000..a82d8b1d5c7a04d793940fe8c85d23e15be91c69 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/eigenvalue.py @@ -0,0 +1,153 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils import log_dist +import numpy as np +import logging +from deepspeed.utils.torch import required_torch_version + + +class Eigenvalue(object): + + def __init__(self, + verbose=False, + max_iter=100, + tol=1e-2, + stability=0, + gas_boundary_resolution=1, + layer_name='', + layer_num=0): + super().__init__() + + self.verbose = verbose + self.max_iter = max_iter + self.tol = tol + self.stability = stability + self.gas_boundary_resolution = gas_boundary_resolution + self.layer_name = layer_name + self.layer_num = layer_num + + assert len(self.layer_name) > 0 and layer_num > 0 + + log_dist( + f'enabled eigenvalue with verbose={verbose}, max_iter={max_iter}, tol={tol}, stability={stability}, gas_boundary_resolution={gas_boundary_resolution}, layer_name={layer_name}, layer_num={layer_num}', + ranks=[0]) + + # Replace all nan/pos-inf/neg-inf to zero + def nan_to_num(self, x): + if required_torch_version(min_version=1.8): + return torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + else: + # Fallback to numpy based implementation for backwards-compatibility with PyTorch 1.7 or older versions. + device = x.device + x = x.cpu().numpy() + x = np.nan_to_num(x=x, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + return torch.from_numpy(x).to(device) + + def normalize(self, v): + norm_squared = self.inner_product(v, v) + norm = norm_squared**0.5 + self.stability + normalized_vectors = [vector / norm for vector in v] + normalized_vectors = [self.nan_to_num(vector) for vector in normalized_vectors] + return normalized_vectors + + def inner_product(self, xs, ys): + return sum([torch.sum(x * y) for (x, y) in zip(xs, ys)]) + + def get_layers(self, module): + scope_names = self.layer_name.split('.') + assert len(scope_names) > 0 + + m = module + for name in scope_names: + assert hasattr(m, name), "layer_name configuration is invalid." + m = getattr(m, name) + + return m + + def compute_eigenvalue(self, module, device=None, scale=1.0): + block_eigenvalue = [] + param_keys = [] + layers = self.get_layers(module) + + for block in range(self.layer_num): + model_block = layers[block] + + # We found this randn() has obvious accuracy impact in some cases, save/recover random state here. + rng_state = torch.random.get_rng_state() + if device is None: + v = [ + torch.randn(p.size()) for p in model_block.parameters() + if p.grad is not None and p.grad.grad_fn is not None + ] + else: + v = [ + torch.randn(p.size(), device=device) for p in model_block.parameters() + if p.grad is not None and p.grad.grad_fn is not None + ] + torch.random.set_rng_state(rng_state) + + grads = [ + param.grad for param in model_block.parameters() + if param.grad is not None and param.grad.grad_fn is not None + ] + params = [ + param for param in model_block.parameters() + if param.grad is not None and param.grad.grad_fn is not None + ] + + layer_keys = [id(p) for p in model_block.parameters()] + param_keys.append(layer_keys) + + v = self.normalize(v) + + # Disable eigenvalue if the model doesn't support second order gradients computation, + # e.g. when enabling DS transformer kernel. + if len(grads) == 0 or len(params) == 0: + log_dist(f'The model does NOT support eigenvalue computation.', ranks=[0], level=logging.WARNING) + return [] + + i = 0 + eigenvalue_current, eigenvalue_previous = 1., 0. + + while (i < self.max_iter) and abs(eigenvalue_current) > 0 and (abs( + (eigenvalue_current - eigenvalue_previous) / eigenvalue_current) + >= self.tol): # test convergence criteria + eigenvalue_previous = eigenvalue_current + + Hv = torch.autograd.grad(grads, params, grad_outputs=v, only_inputs=True, retain_graph=True) + #Hv = [hv.float() for hv in Hv] + Hv = [self.nan_to_num(hv).float() for hv in Hv] + + eigenvalue_current = self.inner_product(Hv, v).item() + + v = self.normalize(Hv) + v = [x / scale for x in v] + i += 1 + + eigenvalue_current *= scale + block_eigenvalue.append(eigenvalue_current) + + if self.verbose: + log_dist(f'block: {block}, power iteration: {i}, eigenvalue: {eigenvalue_current}', ranks=[0]) + + block_eigenvalue = self.post_process(block_eigenvalue) + + if self.verbose: + log_dist(f'post processed block_eigenvalue: {block_eigenvalue}', ranks=[0]) + + # {param_id: (eigenvalue, layer_id)} + ev_dict = {} + for i, (layer_keys, value) in enumerate(zip(param_keys, block_eigenvalue)): + ev_dict.update(dict.fromkeys(layer_keys, (value, i))) + + return ev_dict + + # 1. Map all eigenvalues to [0, 1.0]. + # 2. Some layers can't generate valid eigenvalues on fp16 precision, use 1.0 instead. + def post_process(self, value_list): + max_value = abs(max(value_list, key=abs)) + return [abs(v) / max_value if v != 0.0 else 1.0 for v in value_list] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..76b83a716ffee91d78c7d2dea44233ba8928cfdf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py @@ -0,0 +1,4074 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import re +import stat +import torch +import hashlib +from collections import defaultdict, OrderedDict, deque +from shutil import copyfile +import gc + +from torch.nn.modules import Module +from torch.nn.parameter import Parameter +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from contextlib import contextmanager + +from typing import Callable, Dict, Union, Iterable, Container, List + +import deepspeed + +from deepspeed import comm as dist +from deepspeed.runtime.utils import see_memory_usage, DummyOptim +from .zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum +from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer +from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus +from deepspeed.runtime.zero.utils import is_zero_supported_optimizer, ZeRORuntimeException +from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload +from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION + +from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer +from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer +from deepspeed.runtime.bf16_optimizer import BF16_Optimizer + +from deepspeed.linear.optimized_linear import LoRAOptimizedLinear +from deepspeed.module_inject.layers import GatherReplacedLayerParams, configure_tensor_parallel_runtime +from deepspeed.runtime.config import DEEPSPEED_OPTIMIZERS, \ + ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, \ + TORCH_ADAM_PARAM, ADAM_W_MODE, ADAM_W_MODE_DEFAULT, ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, \ + MUSGD_OPTIMIZER, LION_OPTIMIZER + +from deepspeed.runtime.model_checkpointing.constants import ValidationMode, \ + CHECKPOINT_TAG_VALIDATION, CHECKPOINT_WRITER, CHECKPOINT_SERIALIZATION + +from deepspeed.runtime.dataloader import DeepSpeedDataLoader +from deepspeed.runtime.constants import \ + ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \ + PLD_THETA, PLD_GAMMA, BFLOAT16, FP16, AMP, GRADIENT_ACCUMULATION_STEPS, \ + DATA_PARALLEL_GROUP, GLOBAL_RANK +from deepspeed.runtime.zero.config import ZeroStageEnum +from deepspeed.compression import compression_scheduler +from deepspeed.compression.constants import \ + WEIGHT_QUANTIZE_IN_FORWARD_ENABLED, \ + WEIGHT_QUANTIZATION, SHARED_PARAMETERS, \ + WEIGHT_QUANTIZE_ENABLED, \ + WEIGHT_QUANTIZE_GROUPS, \ + WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE, \ + WEIGHT_QUANTIZE_CHANGE_RATIO, \ + WEIGHT_QUANTIZE_TYPE, \ + WEIGHT_QUANTIZE_ROUNDING, \ + WEIGHT_QUANTIZE_VERBOSE, \ + WEIGHT_QUANTIZE_KERNEL +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FROZEN_PARAM_FRAGMENTS +from deepspeed.checkpoint.utils import clone_tensors_for_torch_save +from deepspeed.runtime.sparse_tensor import SparseTensor + +from deepspeed.runtime import lr_schedules +from deepspeed.utils import groups +from deepspeed.utils import logger, log_dist, instrument_w_nvtx +from deepspeed.utils.timer import NoopTimer, ThroughputTimer, SynchronizedWallClockTimer, \ + FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, \ + STEP_MICRO_TIMER, \ + FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ + STEP_GLOBAL_TIMER +from deepspeed.utils.debug import debug_extract_module_and_param_names, debug_clear_module_and_param_names +from deepspeed.monitor.monitor import MonitorMaster +from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop +from deepspeed.runtime.utils import clip_grad_norm_, compare_tensors_in_structures +from deepspeed.runtime.eigenvalue import Eigenvalue +from deepspeed.runtime.data_pipeline.constants import DATA_SAMPLING, \ + DATA_ROUTING, DATA_SAMPLING_ENABLED, CURRICULUM_LEARNING, \ + CURRICULUM_LEARNING_ENABLED, DATA_SAMPLING_NUM_WORKERS, RANDOM_LTD, \ + RANDOM_LTD_ENABLED, RANDOM_LTD_LAYER_ID, RANDOM_LTD_LAYER_NUM, \ + RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE, RANDOM_LTD_LAYER_TOKEN_LR_ENABLED, \ + RANDOM_LTD_GLOBAL_BATCH_SIZE, RANDOM_LTD_MICRO_BATCH_SIZE, DATA_EFFICIENCY +from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler +from deepspeed.runtime.checkpoint_engine import (create_checkpoint_engine, TorchCheckpointEngine, CheckpointCommitInfo) + +from deepspeed.runtime.data_pipeline.data_routing.scheduler import RandomLTDScheduler +from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict +from deepspeed.runtime.data_pipeline.data_routing.basic_layer import RandomLayerTokenDrop + +from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint +from deepspeed.runtime.torch_autocast import init_autocast_params, get_default_autocast_lower_precision_modules, validate_nested_autocast + +from .pipe.module import PipelineModule +from .utils import get_ma_status +from .compiler import is_compile_supported +from ..ops.adam import FusedAdam +from ..moe.sharded_moe import TopKGate, MOELayer +from ..moe.layer import MoE +from ..moe.utils import is_moe_param, configure_moe_param_groups +from ..git_version_info import version + +from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler +from deepspeed.utils.logging import print_json_dist, print_configuration + +from deepspeed.accelerator import get_accelerator + +from deepspeed.runtime.config import DtypeEnum + +from deepspeed.compile.util import is_deepcompile_supported, get_deepcompile_handle, deepcompile_backward_prologue +from deepspeed.compile.backend import register_compile_pass, opt_passes +from deepspeed.compile.passes import zero3_compile, prefetch, selective_gather, offload_adam_states +from deepspeed.compile.init_z1 import init_z1 +from deepspeed.compile.init_z3 import init_z3 + +MEMORY_OPT_ALLREDUCE_SIZE = 500000000 + +DeepSpeedOptimizerCallable = \ + Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer] +DeepSpeedSchedulerCallable = Callable[[Optimizer], _LRScheduler] + +try: + import apex + from apex import amp + APEX_INSTALLED = True +except ImportError: + # Fail silently so we don't spam logs unnecessarily if user isn't using amp + APEX_INSTALLED = False + + +def split_half_float_double_sparse(tensors): + device_type = get_accelerator().device_name() + supported_types = get_accelerator().supported_dtypes() + + for t in tensors: + assert t.dtype in supported_types, f"attempting to reduce an unsupported grad type: {t.dtype}" + + sparse_tensor_buckets, dense_tensor_buckets = [], [] + for i, dtype in enumerate(supported_types): + sparse_bucket, dense_bucket = [], [] + for t in tensors: + if t.dtype == dtype: + if isinstance(t, SparseTensor): + sparse_bucket.append(t) + else: + dense_bucket.append(t) + if sparse_bucket: + sparse_tensor_buckets.append((dtype, sparse_bucket)) + if dense_bucket: + dense_tensor_buckets.append((dtype, dense_bucket)) + return sparse_tensor_buckets, dense_tensor_buckets + + +class EngineTimers(object): + r"""Wallclock timers for DeepSpeedEngine""" + + def __init__(self, enable_micro_timers, enable_global_timers): + self.forward_timers = [] + self.backward_timers = [] + self.backward_inner_timers = [] + self.backward_reduce_timers = [] + self.step_timers = [] + self.global_timers = [] + self.micro_timers = [] + + if enable_micro_timers: + self.forward_timers += [FORWARD_MICRO_TIMER] + self.backward_timers += [BACKWARD_MICRO_TIMER] + self.backward_inner_timers += [BACKWARD_INNER_MICRO_TIMER] + self.backward_reduce_timers += [BACKWARD_REDUCE_MICRO_TIMER] + self.step_timers += [STEP_MICRO_TIMER] + self.micro_timers += [ + FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, + STEP_MICRO_TIMER + ] + + if enable_global_timers: + self.forward_timers += [FORWARD_GLOBAL_TIMER] + self.backward_timers += [BACKWARD_GLOBAL_TIMER] + self.backward_inner_timers += [BACKWARD_INNER_GLOBAL_TIMER] + self.backward_reduce_timers += [BACKWARD_REDUCE_GLOBAL_TIMER] + self.step_timers += [STEP_GLOBAL_TIMER] + self.global_timers += [ + FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, + STEP_GLOBAL_TIMER + ] + + +class DeepSpeedEngine(Module): + r"""DeepSpeed engine for training.""" + + def __init__(self, + args, + model, + optimizer=None, + model_parameters=None, + training_data=None, + lr_scheduler=None, + mpu=None, + dist_init_required=None, + collate_fn=None, + config=None, + config_class=None, + mesh_device=None, + dont_change_device=False): + super(DeepSpeedEngine, self).__init__() + self.dont_change_device = dont_change_device + self.client_optimizer = optimizer + self.client_lr_scheduler = lr_scheduler + self.training_data = training_data + self.collate_fn = collate_fn + self.mpu = mpu + self.all_to_all_group = None + self.data_parallel_group = None + self.global_steps = 0 + self.global_samples = 0 + self.micro_steps = 0 + self.skipped_steps = 0 + self.gradient_average = True + self.warn_unscaled_loss = True + self.config = config + self._config = config_class + self.loaded_checkpoint_mp_world_size = None + self.loaded_checkpoint_dp_world_size = None + self.enable_backward_allreduce = True + self.inside_no_sync_ctxt = False + self.progressive_layer_drop = None + self.eigenvalue = None + self.block_eigenvalue = None + self.gas_boundary_ctr = 0 + self.dist_backend = get_accelerator().communication_backend_name() + self.has_moe_layers = False + self.num_experts = [] + self.gate_modules = [] + self.moe_layers = [] + self._step_applied = False + self._global_grad_norm = None + self.use_ds_comm = False # False --> Use torch.dist, True --> Use ds.comm backend. + self.checkpoint_engine = None + self.optimizer = None + self.basic_optimizer = None + self.lr_scheduler = None + + self._is_gradient_accumulation_boundary = None + self.scale_wrt_gas = None + self.losses = None + self.mesh_device = mesh_device + + # for debug purposes - can then debug print: debug_get_module_name(module) + debug_extract_module_and_param_names(model) + + if self.mesh_device: + groups.mesh_device = self.mesh_device + + self._do_args_sanity_check(args) + self._configure_with_arguments(args, mpu) + self._do_sanity_check() + if self.autotp_size() > 1: + self._configure_tensor_parallel(model, self.tensor_parallel_config()) + see_memory_usage(f"DeepSpeed Engine: After args sanity test", force=self.memory_breakdown()) + if mpu is not None: + if self.elasticity_enabled(): + if not self.is_elastic_model_parallel_supported(): + assert not self.elasticity_enabled(), ("Elasticity is not currently supported" + " with model parallelism.") + + self._set_distributed_vars(args) + + dist.configure(self._config) + + self.monitor = MonitorMaster(self._config.monitor_config) + + see_memory_usage( + f"DeepSpeed Engine: Before configure distributed model", + force=self.memory_breakdown(), + ) + + self.pipeline_parallelism = isinstance(model, PipelineModule) + + # Configure distributed model + self._configure_distributed_model(model) + + if not self.is_deepcompile_enabled(): + self.module_forward_pre_hook = self._create_module_forward_pre_hook() + self.module_forward_post_hook = self._create_module_forward_post_hook() + + # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict + self.param_names = {param: name for name, param in model.named_parameters()} + + self._get_model_parameters() + + see_memory_usage(f"DeepSpeed Engine: After configure distributed model") + + # Configure wall clock timers + self.timers = SynchronizedWallClockTimer() + # Throughput timer + self.tput_timer = ThroughputTimer(self._config.timers_config, + batch_size=self.train_batch_size(), + steps_per_output=self.steps_per_print(), + monitor_memory=False) + + log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0]) + + if self.flops_profiler_enabled(): + self.flops_profiler = FlopsProfiler(self.module, self, self.flops_profiler_recompute_fwd_factor()) + + if training_data: + self.training_dataloader = self.deepspeed_io(training_data) + else: + self.training_dataloader = None + + # Configure optimizer and scheduler + has_optimizer = False + + if optimizer or self.optimizer_name(): + has_optimizer = True + # If no parameters given by init default to module parameters + if model_parameters is None: + model_parameters = self.module.parameters() + + # Convert model parameters from generator to list + if not isinstance(model_parameters, list): + model_parameters = list(model_parameters) + + if self.torch_autocast_enabled(): + init_autocast_params(self, self.torch_autocast_dtype(), self.torch_autocast_lower_precision_safe_modules()) + + if has_optimizer: + self._configure_optimizer(optimizer, model_parameters) + self._configure_lr_scheduler() + self._report_progress(0) + elif self.zero_optimization(): + # no optim selected but zero is enabled + self.optimizer = self._configure_zero_optimizer(optimizer=None) + elif self.bfloat16_enabled(): + self.optimizer = self._configure_bf16_optimizer(optimizer=None) + + # Hook optimizer for snip_momentum pruning + if hasattr(model, 'pruners'): + from ..compression.helper import rewrite_optimizer_step + self.optimizer.pruners = model.pruners + rewrite_optimizer_step(self.optimizer) + + # Bookkeeping for sparse support + self.sparse_tensor_module_names = set() + # if self.sparse_gradients_enabled(): + for name, module in self.module.named_modules(): + if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)) and self.sparse_gradients_enabled(): + self.sparse_tensor_module_names.add(name + ".weight") + logger.info("Will convert {} to sparse tensor during training".format(name)) + + self._optimized_linear_offload_setup() + + self.save_non_zero_checkpoint = False + self.save_zero_checkpoint = False + if not isinstance(self.optimizer, DeepSpeedZeRoOffload): + self._configure_checkpointing() + + if self.eigenvalue_enabled(): + self.eigenvalue = self._configure_eigenvalue() + + if self.pld_enabled(): + self.progressive_layer_drop = self._configure_progressive_layer_drop() + + if self.curriculum_enabled_legacy(): + self.curriculum_scheduler_legacy = self._configure_curriculum_scheduler_legacy() + + if self.random_ltd_enabled(): + random_ltd_config = self.random_ltd_config() + random_ltd_config[RANDOM_LTD_GLOBAL_BATCH_SIZE] = self.train_batch_size() + random_ltd_config[RANDOM_LTD_MICRO_BATCH_SIZE] = self.train_micro_batch_size_per_gpu() + self.random_ltd_scheduler = self._configure_random_ltd_scheduler(random_ltd_config) + + # Engine timers + + self.engine_timers = EngineTimers(enable_micro_timers=self.wall_clock_breakdown(), + enable_global_timers=self.wall_clock_breakdown() + or self.flops_profiler_enabled()) + + if self.global_rank == 0: + self._config.print("DeepSpeedEngine configuration") + if self.dump_state(): + print_configuration(self, "DeepSpeedEngine") + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + + self._is_compiled = False + if is_deepcompile_supported(): + # Predefined compile passes + self.register_compile_pass(zero3_compile.NAME, zero3_compile.add_z3_gather_release) + self.register_compile_pass(prefetch.NAME, prefetch.schedule_prefetch) + self.register_compile_pass(selective_gather.NAME, selective_gather.selective_gather) + self.register_compile_pass(offload_adam_states.NAME, offload_adam_states.move_opt_states) + + def _optimized_linear_offload_setup(self): + self.optimized_linear_base_weight_sharding = False + self.optimized_linear_lora_enabled = False + offload_ratio = None + for _, module in self.module.named_modules(): + if isinstance(module, LoRAOptimizedLinear): + self.optimized_linear_lora_enabled = True + offload_ratio = None + if offload_ratio is not None: + assert offload_ratio == module.lora_config.offload_ratio, \ + "all lora_config offload ratios should be the same across the model" + offload_ratio = module.lora_config.offload_ratio + if module.zero_shards > 1: + # set attr so checkpoint saving can handle BWS properly + self.optimized_linear_base_weight_sharding = True + + if offload_ratio is None: + # Nothing enabled, do nothing + return + + total_params = 0 + for _, p in self.module.named_parameters(): + if hasattr(p, 'ds_optim_param'): + total_params += p.numel() + + offload_limit = total_params * offload_ratio + logger.info(f'offloading {offload_ratio*100}% of eligible params, specifically {offload_limit} params') + total_offloaded = 0 + for _, p in self.module.named_parameters(): + if hasattr(p, 'ds_optim_param'): + if total_offloaded < offload_limit: + total_offloaded += p.numel() + p.ds_offload = True + p.offload() + else: + p.ds_offload = False + + def _configure_tensor_parallel(self, model, tp_config): + self._configure_tensor_parallel_states(model) + configure_tensor_parallel_runtime(tp_config) + + def _configure_tensor_parallel_states(self, model): + """ + Configures the tensor parallel states for the model. + This includes setting up the tensor parallel groups, initializing the TP mesh, + and registering a pre-hook to ensure that the Dataloader inputs are consistent across ranks. + """ + self._set_client_model(model) + # sanity check + # currently, the compatibility between 'autotp' and 'zero > 1' has not been validated + assert self.zero_optimization_stage( + ) <= 2, "Currently, the compatibility between 'autotp' and 'zero_stage = 3' has not been validated" + + self.mpu = groups + self.mpu._init_tp_mesh_device(tensor_model_parallel_size=self.autotp_size()) + + self.first_dataloader_check = None + + def check_dataloader_inputs_same_across_ranks(module, args, kwargs): + + def broadcast_and_check(args, bcast_rank, bcast_group): + if isinstance(args, tuple): + args = list(args) + if len(args) > 0: + if self.mpu.get_tensor_model_parallel_rank() == 0: + _src_args = [args] + dist.broadcast_object_list(object_list=_src_args, + src=bcast_rank, + group=bcast_group, + device=get_accelerator().current_device()) + # Rank 0 does not need to compare with itself + is_equal = True + else: + _src_args = [None] + dist.broadcast_object_list(object_list=_src_args, + src=bcast_rank, + group=bcast_group, + device=get_accelerator().current_device()) + + is_equal = compare_tensors_in_structures(args, _src_args[0]) + + equal_tensor = torch.tensor(is_equal, + dtype=self.communication_data_type, + device=get_accelerator().current_device()) + dist.all_reduce(equal_tensor, group=bcast_group) + assert torch.equal( + equal_tensor, + torch.tensor(groups.get_tensor_model_parallel_world_size(), + dtype=self.communication_data_type, + device=get_accelerator().current_device()) + ), "Data inconsistency within the TP group. Please check the Dataloader implementation to ensure consistency." + + bcast_rank = self.mpu.get_tensor_model_parallel_src_rank() + bcast_group = self.mpu.get_tensor_model_parallel_group() + + broadcast_and_check(args, bcast_rank, bcast_group) + broadcast_and_check(kwargs, bcast_rank, bcast_group) + + logger.info(f":The Dataloader has passed the TP group consistency check.") + self.first_dataloader_check.remove() + + self.first_dataloader_check = self.module.register_forward_pre_hook(check_dataloader_inputs_same_across_ranks, + prepend=True, + with_kwargs=True) + + def __del__(self): + self.destroy() + + def destroy(self): + if self.optimizer is not None and hasattr(self.optimizer, 'destroy'): + self.optimizer.destroy() + if self.is_deepcompile_enabled(): + get_deepcompile_handle().cleanup() + debug_clear_module_and_param_names() + + if self.checkpoint_engine is not None and self.checkpoint_engine.is_decoupled(): + self.checkpoint_engine.cleanup() + + def _get_model_parameters(self): + if self.autotuning_profile_model_info(): + self.autotuning_model_info = {} + num_params = 0 + trainable_num_params = 0 + + for p in self.module.parameters(): + # since user code might call deepspeed.zero.Init() before deepspeed.initialize(), need to check the attribute to check if the parameter is partitioned in zero 3 already or not + n = 0 + if hasattr(p, "ds_tensor"): # if the parameter is partitioned in zero 3 + n += p.ds_numel + else: # if the parameter is not partitioned in zero 3 yet + n += p.numel() + num_params += n + if p.requires_grad: + trainable_num_params += n + if self.global_rank == 0: + self.autotuning_model_info["num_params"] = num_params * self.mp_world_size + self.autotuning_model_info["trainable_num_params"] = trainable_num_params * self.mp_world_size + + logger.info(f"model parameter = {num_params}") + + def get_batch_info(self): + """Get all training batch related settings. + Returns: + train_batch_size (int): The effective training batch size. This is the amount of data + samples that leads to one step of model update. + train_micro_batch_size_per_gpu (int): Batch size to be processed by one GPU in one + step (without gradient accumulation). + gradient_accumulation_steps (int): Number of training steps to accumulate gradients + before averaging and applying them. + """ + return ( + self.train_batch_size, + self.train_micro_batch_size_per_gpu, + self.gradient_accumulation_steps, + ) + + def set_train_batch_size(self, train_batch_size): + """Adjust the global batch size by increasing or decreasing the number of + micro-batches (i.e., gradient accumulation steps). The size of each micro-batch + (i.e., ``train_micro_batch_size_per_gpu``) is not changed. + Args: + train_batch_size (int): The new global batch size for training. + Raises: + ValueError: if ``train_batch_size`` is not divisible by the + configured micro-batch size and data parallelism. + """ + if train_batch_size % (self.train_micro_batch_size_per_gpu() * self.dp_world_size) != 0: + #print(f'{train_batch_size=} {self.train_micro_batch_size_per_gpu()=} {self.dp_world_size=}') + raise ValueError(f'Train batch size must be divisible by micro-batch data parallelism') + new_gas = train_batch_size // (self.train_micro_batch_size_per_gpu() * self.dp_world_size) + # overwrite config + self._config.train_batch_size = train_batch_size + self._config.gradient_accumulation_steps = new_gas + + def set_train_micro_batch_size(self, micro_batch_size): + """Adjust the micro batch size(i.e., the micro batch size in every data parallel group), + while keep the gradient accumulation steps the same. + Args: + micro_batch_size (int): The new micro batch size for training. + """ + # overwrite config + new_global_batch_size = micro_batch_size * self._config.gradient_accumulation_steps * self.dp_world_size + self._config.train_batch_size = new_global_batch_size + self._config.train_micro_batch_size_per_gpu = micro_batch_size + + def set_data_post_process_func(self, post_process_func): + if self.training_dataloader is not None: + self.training_dataloader.post_process_func = post_process_func + + def set_custom_curriculum_learning_schedule(self, schedule_func_dict): + if self.training_dataloader is not None and self.curriculum_learning_enabled(): + self.training_dataloader.data_sampler.set_custom_curriculum_learning_schedule(schedule_func_dict) + + def get_global_grad_norm(self) -> float: + """Return the 2-norm of all gradients. If there is model parallelism, + the norm will be global. + The computed norm will be cached and reused until the next step() pass. + .. note:: + In the presence of model parallelism, this is a collective call + and acts as a barrier among ``mpu.get_model_parallel_group()``. + Returns: + float: norm + """ + return self._global_grad_norm + + def __getattr__(self, name): + """ + Pass through attributes defined in the model if they are not overridden by ds-engine. + """ + + _module = {} + if "module" in self.__dict__: + _module = self.__dict__['module'] + if name in dir(self): + return getattr(self, name) + elif name in dir(_module): + return getattr(_module, name) + else: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + + def checkpoint_serialization_enabled(self): + return self._config.checkpoint_config[CHECKPOINT_SERIALIZATION] + + def checkpoint_writer_enabled(self): + return self._config.checkpoint_config[CHECKPOINT_WRITER] is not None + + def checkpoint_tag_validation_enabled(self): + return self._config.checkpoint_config[CHECKPOINT_TAG_VALIDATION] != ValidationMode.IGNORE + + def checkpoint_tag_validation_fail(self): + return self._config.checkpoint_config[CHECKPOINT_TAG_VALIDATION] == ValidationMode.FAIL + + def elasticity_enabled(self): + return self._config.elasticity_enabled + + def is_elastic_model_parallel_supported(self): + if self.elasticity_enabled(): + # Add code for finding number of GPUs per node automatically + if self._config.num_gpus_per_node % self._config.elastic_model_parallel_size == 0: + return True + else: + return False + + def pld_enabled(self): + return self._config.pld_enabled + + def pld_params(self): + return self._config.pld_params + + def pld_theta(self): + return self.pld_params()[PLD_THETA] + + def pld_gamma(self): + return self.pld_params()[PLD_GAMMA] + + def eigenvalue_enabled(self): + return self._config.eigenvalue_enabled + + def eigenvalue_verbose(self): + return self._config.eigenvalue_verbose + + def eigenvalue_max_iter(self): + return self._config.eigenvalue_max_iter + + def eigenvalue_tol(self): + return self._config.eigenvalue_tol + + def eigenvalue_stability(self): + return self._config.eigenvalue_stability + + def eigenvalue_gas_boundary_resolution(self): + return self._config.eigenvalue_gas_boundary_resolution + + def eigenvalue_layer_name(self): + return self._config.eigenvalue_layer_name + + def eigenvalue_layer_num(self): + return self._config.eigenvalue_layer_num + + def curriculum_enabled_legacy(self): + return self._config.curriculum_enabled_legacy + + def curriculum_params_legacy(self): + return self._config.curriculum_params_legacy + + def data_efficiency_enabled(self): + return self._config.data_efficiency_enabled + + def data_efficiency_config(self): + return self._config.data_efficiency_config + + def data_sampling_enabled(self): + return self._config.data_efficiency_config[DATA_SAMPLING][DATA_SAMPLING_ENABLED] + + def data_sampling_config(self): + return self._config.data_efficiency_config[DATA_SAMPLING] + + def curriculum_learning_enabled(self): + return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED] + + def curriculum_learning_config(self): + return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING] + + def random_ltd_enabled(self): + return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD][RANDOM_LTD_ENABLED] + + def random_ltd_config(self): + return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD] + + def random_ltd_initialize(self): + assert self.random_ltd_enabled() + random_ltd_config = self.random_ltd_config() + random_ltd_queue = deque([x for x in sorted(random_ltd_config[RANDOM_LTD_LAYER_ID])]) + count = 0 + for name, layer in self.module.named_modules(): + if isinstance(layer, RandomLayerTokenDrop): + if len(random_ltd_queue) != 0 and str(random_ltd_queue[0]) in name: ###[1,2,3] + layer.init_config(random_ltd_config, self.random_ltd_scheduler, count) + random_ltd_queue.popleft() + count += 1 + + if random_ltd_config[RANDOM_LTD_LAYER_NUM] != count: + raise ValueError(f'random_ltd_layer_num {random_ltd_config[RANDOM_LTD_LAYER_NUM]} must be \ + equivalent to the len of random_ltd_layer_id {count}') + + if random_ltd_config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: + assert self.client_lr_scheduler is None + raise ValueError(f'not yet support') + #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler) + + def get_sequence_parallel_group(self): + return self.seq_parallel_group + + def wall_clock_breakdown(self): + return self._config.wall_clock_breakdown + + def flops_profiler_enabled(self): + return self._config.flops_profiler_config.enabled or self.autotuning_enabled() + + def flops_profiler_recompute_fwd_factor(self): + return self._config.flops_profiler_config.recompute_fwd_factor + + def flops_profiler_profile_step(self): + step = self._config.flops_profiler_config.profile_step + if self._config.autotuning_config.enabled: + step = self.autotuning_start_profile_step() + return step + + def flops_profiler_module_depth(self): + return self._config.flops_profiler_config.module_depth + + def flops_profiler_top_modules(self): + return self._config.flops_profiler_config.top_modules + + def flops_profiler_detailed(self): + if self._config.autotuning_config.enabled: + return False + return self._config.flops_profiler_config.detailed + + def flops_profiler_output_file(self): + return self._config.flops_profiler_config.output_file + + def memory_breakdown(self): + return self._config.memory_breakdown + + def autotuning_enabled(self): + return self._config.autotuning_config.enabled + + def autotuning_start_profile_step(self): + return self._config.autotuning_config.start_profile_step + + def autotuning_end_profile_step(self): + return self._config.autotuning_config.end_profile_step + + def autotuning_metric_path(self): + path = self._config.autotuning_config.metric_path + if not path: + path = os.path.join(os.getcwd(), "autotuning_metric.json") + return path + + def autotuning_model_info_path(self): + path = self._config.autotuning_config.model_info_path + if not path: + path = os.path.join(os.getcwd(), "autotuning_model_info.json") + return path + + def autotuning_metric(self): + return self._config.autotuning_config.metric + + def autotuning_profile_model_info(self): + return self.autotuning_enabled( + ) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get( + "profile", False) + + def sparse_gradients_enabled(self): + return self._config.sparse_gradients_enabled + + def train_batch_size(self): + return self._config.train_batch_size + + def train_micro_batch_size_per_gpu(self): + return self._config.train_micro_batch_size_per_gpu + + def optimizer_name(self): + return (self.client_optimizer.__class__.__name__ if self.client_optimizer else self._config.optimizer_name) + + def optimizer_params(self): + return self._config.optimizer_params + + def optimizer_legacy_fusion(self): + return self._config.optimizer_legacy_fusion + + def scheduler_name(self): + return self._config.scheduler_name + + def scheduler_params(self): + return self._config.scheduler_params + + def quantize_training(self): + return ( + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] + [WEIGHT_QUANTIZE_IN_FORWARD_ENABLED], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ENABLED], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_GROUPS], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] + [WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_CHANGE_RATIO], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_TYPE], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ROUNDING], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_VERBOSE], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_KERNEL], + ) + + def zero_optimization(self): + return self._config.zero_enabled + + def zero_allow_untested_optimizer(self): + return self._config.zero_allow_untested_optimizer + + def zero_force_ds_cpu_optimizer(self): + return self._config.zero_force_ds_cpu_optimizer + + def zero_reduce_scatter(self): + return self._config.zero_config.reduce_scatter + + def zero_overlap_comm(self): + return self._config.zero_config.overlap_comm + + def zero_offload_optimizer(self): + return self._config.zero_config.offload_optimizer + + def zero_offload_param(self): + return self._config.zero_config.offload_param + + def zero_use_cpu_optimizer(self): + if self._config.zero_config.offload_optimizer is not None: + return self._config.zero_config.offload_optimizer.device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme] + return False + + def zero_cpu_offload(self): + if self._config.zero_config.offload_optimizer is not None: + return self._config.zero_config.offload_optimizer.device == OffloadDeviceEnum.cpu + return False + + def zero_partial_offload(self): + return getattr(self._config.zero_config.offload_optimizer, "ratio", 1.0) + + def zero_sub_group_size(self): + return self._config.zero_config.sub_group_size + + def zero_optimization_stage(self): + return self._config.zero_optimization_stage + + def mics_shard_size(self): + return self._config.mics_shard_size + + def zero_reduce_bucket_size(self): + return self._config.zero_config.reduce_bucket_size + + def zero_multi_rank_bucket_allreduce(self): + return self._config.zero_config.use_multi_rank_bucket_allreduce + + def zero_allgather_bucket_size(self): + return self._config.zero_config.allgather_bucket_size + + def zero_optimization_partition_gradients(self): + return self.zero_optimization_stage() >= ZeroStageEnum.gradients + + def zero_optimization_partition_weights(self): + return self.zero_optimization_stage() >= ZeroStageEnum.weights + + def is_first_weights_partition_group(self): + ret = True if self.mics_shard_size() < 0 \ + and self.zero_optimization_partition_weights() else False + if self.mics_shard_size() > 0 and self.global_rank < self.mics_shard_size(): + ret = True + return ret + + def zero_contiguous_gradients(self): + return self._config.zero_config.contiguous_gradients + + def zero_load_from_fp32_weights(self): + return self._config.zero_config.load_from_fp32_weights + + def zero_elastic_checkpoint(self): + return self._config.zero_config.elastic_checkpoint + + def zero_nvme_offload_optimizer(self): + return getattr(self.optimizer, "swap_optimizer", False) + + def zero_max_live_parameters(self): + return self._config.zero_config.max_live_parameters + + def zero_max_reuse_distance(self): + return self._config.zero_config.max_reuse_distance + + def zero_prefetch_bucket_size(self): + return self._config.zero_config.prefetch_bucket_size + + def zero_module_granularity_threshold(self): + return self._config.zero_config.module_granularity_threshold + + def zero_param_persistence_threshold(self): + return self._config.zero_config.param_persistence_threshold + + def zero_model_persistence_threshold(self): + return self._config.zero_config.model_persistence_threshold + + def zero_gather_16bit_weights_on_model_save(self): + return self._config.zero_config.gather_16bit_weights_on_model_save + + def zero_grad_hooks(self): + return self._config.zero_config.grad_hooks + + def zero_legacy_stage1(self): + return self._config.zero_config.legacy_stage1 + + def zero_ignore_unused_parameters(self): + return self._config.zero_config.ignore_unused_parameters + + def tensor_parallel_config(self): + return self._config.tensor_parallel_config + + def autotp_size(self): + return self._config.tensor_parallel_config.autotp_size + + def graph_harvesting(self): + return self._config.graph_harvesting + + def fp16_enabled(self): + return self._config.float16_config.enabled + + def bfloat16_enabled(self): + return self._config.bfloat16_config.enabled + + def fp16_master_weights_and_gradients(self): + return self._config.float16_config.fp16_master_weights_and_grads + + def amp_enabled(self): + return self._config.amp_enabled + + def amp_params(self): + return self._config.amp_params + + def torch_autocast_enabled(self) -> bool: + return self._config.torch_autocast_enabled + + def torch_autocast_dtype(self) -> torch.dtype: + return self._config.torch_autocast_dtype + + def torch_autocast_lower_precision_safe_modules(self) -> List[str]: + module_names = self._config.torch_autocast_lower_precision_safe_modules + return get_default_autocast_lower_precision_modules() if module_names is None else module_names + + def fp16_auto_cast(self): + return self._config.float16_config.auto_cast + + def loss_scale(self): + return self._config.float16_config.loss_scale + + def gradient_accumulation_steps(self): + return self._config.gradient_accumulation_steps + + def use_node_local_storage(self): + return self._config.use_node_local_storage + + def load_universal_checkpoint(self): + return self._config.load_universal_checkpoint + + @property + def communication_data_type(self): + res = self._config.communication_data_type + if res is not None: + return res + + if self.fp16_enabled(): + return torch.float16 + + if self.bfloat16_enabled(): + return torch.bfloat16 + + return torch.float32 + + @communication_data_type.setter + def communication_data_type(self, value): + self._config.communication_data_type = value + + def postscale_gradients(self): + return not self._config.prescale_gradients + + def gradient_predivide_factor(self): + return self._config.gradient_predivide_factor + + def steps_per_print(self): + return self._config.steps_per_print + + def zero_allgather_partitions(self): + return self._config.zero_config.allgather_partitions + + def zero_round_robin_gradients(self): + return self._config.zero_config.round_robin_gradients + + def zero_hpz_partition_size(self): + return self._config.zero_config.zero_hpz_partition_size + + def zero_quantized_weights(self): + return self._config.zero_config.zero_quantized_weights + + def zero_quantized_nontrainable_weights(self): + return self._config.zero_config.zero_quantized_nontrainable_weights + + def zero_quantized_gradients(self): + return self._config.zero_config.zero_quantized_gradients + + def zeropp_loco_param(self): + return self._config.zero_config.zeropp_loco_param + + def zero_log_trace_cache_warnings(self): + return self._config.zero_config.log_trace_cache_warnings + + def dump_state(self): + return self._config.dump_state + + def gradient_clipping(self): + return self._config.gradient_clipping + + def dynamic_loss_scale(self): + return self._config.float16_config.loss_scale == 0 + + def initial_dynamic_scale(self): + return self._config.float16_config.initial_dynamic_scale() + + def dynamic_loss_scale_args(self): + return self._config.float16_config.dynamic_loss_scale_args() + + def swap_tensor_config(self): + return self._config.swap_tensor_config + + def aio_config(self): + return self._config.aio_config + + def get_data_types(self): + model_dtype = torch.float32 + if self.fp16_enabled(): + model_dtype = torch.float16 + elif self.bfloat16_enabled(): + model_dtype = torch.bfloat16 + + if self._config.grad_accum_dtype is None: + if model_dtype == torch.bfloat16 and not self.zero_optimization(): + grad_accum_dtype = torch.float32 + else: + grad_accum_dtype = model_dtype + else: + grad_accum_dtype = DtypeEnum(self._config.grad_accum_dtype).value + + return (model_dtype, grad_accum_dtype) + + def _optimizer_has_ckpt_event_prologue(self): + return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_prologue') + + def _optimizer_has_ckpt_event_epilogue(self): + return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_epilogue') + + def _configure_lr_scheduler(self): + if self.client_lr_scheduler: + if isinstance(self.client_lr_scheduler, Callable): + log_dist('DeepSpeed using client callable to create LR scheduler', ranks=[0]) + self.lr_scheduler = self.client_lr_scheduler(self.basic_optimizer) + else: + log_dist('DeepSpeed using client LR scheduler', ranks=[0]) + self.lr_scheduler = self.client_lr_scheduler + else: + # load lr scheduler from json configuration if lr scheduler is not defined and passed in + lr_scheduler = self._scheduler_from_config(self.optimizer) + log_dist(f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}", ranks=[0]) + self.lr_scheduler = lr_scheduler + + log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) + + def _configure_checkpointing(self): + # Enable optimization to parallelize checkpointing of DP state + optimize_dp_state = not self.zero_optimization_partition_weights() + self.checkpoint_engine = create_checkpoint_engine(config_params=self._config, + groups=groups, + zero_stage=self.zero_optimization_stage(), + has_moe_layers=self.has_moe_layers, + optimize_dp_state=optimize_dp_state) + + dp_rank = groups._get_sequence_data_parallel_rank() + rank = self.local_rank if self.use_node_local_storage() else dp_rank + + # Determine if this data parallel process needs to store the model checkpoint + if self.checkpoint_engine.is_data_parallel_writer(rank) \ + or (self.zero_optimization_partition_weights() and self.is_first_weights_partition_group()): + self.save_non_zero_checkpoint = True + + if self.zero_optimization() or self.bfloat16_enabled(): + param_rank = dist.get_rank(group=self.optimizer.dp_process_group) + + # Only the first parameter parallel process needs to store the + # optimizer state checkpoints for zero + self.save_zero_checkpoint = param_rank == dp_rank + + def _scheduler_from_config(self, optimizer): + scheduler_name = self.scheduler_name() + if scheduler_name is not None: + if hasattr(lr_schedules, scheduler_name): + scheduler = getattr(lr_schedules, scheduler_name) + else: + assert hasattr(torch.optim.lr_scheduler, + scheduler_name), f"DeepSpeed does not recognize LR scheduler {scheduler_name}" + + scheduler = getattr(torch.optim.lr_scheduler, scheduler_name) + + scheduler_params = self.scheduler_params() + instantiated_scheduler = scheduler(optimizer, **scheduler_params) + return instantiated_scheduler + else: + return None + + def _set_distributed_vars(self, args): + device_rank = args.device_rank if args is not None and hasattr(args, 'device_rank') else self.local_rank + if device_rank >= 0: + get_accelerator().set_device(device_rank) + self.device = torch.device(get_accelerator().device_name(device_rank)) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + else: + self.world_size = 1 + self.global_rank = 0 + self.device = get_accelerator().device() + + # Configure based on command line arguments + def _configure_with_arguments(self, args, mpu): + # After the distributed backend is initialized we are guaranteed the LOCAL_RANK + # environment variable is set. We must align args.local_rank to this value for + # backwards compatibility with scripts relying on [args|self].local_rank containing + # the correct local rank info. _do_args_sanity_check will ensure this is the case. + + if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: + ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK") + local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank) + assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({ompi_local_rank}), " \ + "not sure how to proceed as we're seeing conflicting local rank info." + os.environ['LOCAL_RANK'] = local_rank + + self.local_rank = int(os.environ['LOCAL_RANK']) + if hasattr(args, 'local_rank'): + args.local_rank = self.local_rank + + # Validate command line arguments + def _do_args_sanity_check(self, args): + assert "LOCAL_RANK" in os.environ or "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment " \ + "variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch's launcher. If using a " \ + "different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed." + + if hasattr(args, 'local_rank') and args.local_rank is not None: + assert isinstance(args.local_rank, + int), f"args.local_rank of {args.local_rank} is an unknown type {type(args.local_rank)}" + if args.local_rank >= 0: + env_local_rank = int(os.environ.get("LOCAL_RANK")) + assert ( + env_local_rank == args.local_rank + ), f"Mismatch in local rank setting, args.local_rank={args.local_rank} but env['LOCAL_RANK']={env_local_rank}." + + def _is_supported_optimizer(self, optimizer_name): + return (optimizer_name in DEEPSPEED_OPTIMIZERS or getattr(torch.optim, optimizer_name, None) is not None) + + def _supported_optims(self): + FairseqOptimizer = None + try: + from fairseq.optim.fairseq_optimizer import FairseqOptimizer + except ImportError: + pass + + expected_optim_types = [Optimizer] + if FairseqOptimizer: + # fairseq optims are not torch.optim objects + expected_optim_types.append(FairseqOptimizer) + return expected_optim_types + + # Validate configuration based on command line arguments + def _do_sanity_check(self): + if self.fp16_enabled() and not get_accelerator().is_fp16_supported(): + raise ValueError("Type fp16 is not supported on your device.") + + if self.bfloat16_enabled() and not get_accelerator().is_bf16_supported(): + raise ValueError("Type bf16 is not supported on your device.") + + expected_optim_types = self._supported_optims() + expected_optim_types += [type(None), Callable] + assert isinstance(self.client_optimizer, tuple(expected_optim_types)), \ + f'Client Optimizer is of unexpected type {type(self.client_optimizer)}' + + if not self.client_optimizer: + if self.optimizer_name() is not None: + assert self._is_supported_optimizer( + self.optimizer_name()), "{} is not a supported DeepSpeed Optimizer".format(self.optimizer_name()) + + if (self.optimizer_name() == LAMB_OPTIMIZER or self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER): + assert (self.dynamic_loss_scale()), "DeepSpeed {} optimizer requires dynamic loss scaling".format( + self.optimizer_name()) + + # Detect invalid combinations of client optimizer and client scheduler + if isinstance(self.client_lr_scheduler, _LRScheduler): + assert isinstance(self.client_optimizer, Optimizer), \ + f'Client Optimizer (type = {type(self.client_optimizer)} is not instantiated but Client LR Scheduler is instantiated' + + def _broadcast_model(self): + + def is_replicated(p): + if hasattr(p, "ds_status") and p.ds_status is not ZeroParamStatus.AVAILABLE: + return False + elif hasattr(p, 'ds_optim_param'): + # do not broadcast OptimizedLinear parameters, they are unique per base weight shard + return False + return True + + for n, p in self.module.named_parameters(): + # Broadcast the model for different parameters + if is_moe_param(p): + if torch.is_tensor(p) and is_replicated(p): + dist.broadcast(p.data, + groups._get_expert_broadcast_src_rank(p.group_name), + group=self.expert_data_parallel_group[p.group_name]) + else: + if torch.is_tensor(p) and is_replicated(p): + dist.broadcast(p.data, groups._get_broadcast_src_rank(), group=self.seq_data_parallel_group) + + @staticmethod + def __check_params(model: Module, dtype: torch.dtype) -> None: + return + if not all(param.dtype == dtype for param in model.parameters()) and dist.get_rank() == 0: + raise ValueError(f"{dtype} is enabled but the following parameters have dtype that is " + f"not {dtype}: " + f"{[(n, p.dtype) for n, p in model.named_parameters() if p.dtype != dtype]}") + + def _set_client_model(self, model): + # register client model in _modules so that nn.module methods work correctly + modules = self.__dict__.get('_modules') + modules['module'] = model + # register module attribute in engine but avoid getattr + self.__dict__['module'] = model + + def _configure_distributed_model(self, model): + self._set_client_model(model) + is_zero_init_model = self.zero_optimization_partition_weights() and any( + [hasattr(param, "ds_id") for param in self.module.parameters()]) + + if self.fp16_enabled(): + if is_zero_init_model: + self.__check_params(self.module, torch.half) + self.module.half() + elif self.bfloat16_enabled(): + if is_zero_init_model: + self.__check_params(self.module, torch.bfloat16) + self.module.bfloat16() + else: + self.__check_params(self.module, torch.float) + + # zero.Init() handles device placement of model + if not (self.dont_change_device or is_zero_init_model): + self.module.to(self.device) + + # MoE related initialization + for _, module in self.module.named_modules(): + if isinstance(module, MoE): + self.has_moe_layers = True + self.num_experts.append(module.num_experts) + + if self.has_moe_layers: + for _, module in self.module.named_modules(): + if isinstance(module, TopKGate): + self.gate_modules.append(module) + if self.wall_clock_breakdown(): + module.wall_clock_breakdown = True + if isinstance(module, MOELayer): + self.moe_layers.append(module) + if self.wall_clock_breakdown(): + module.wall_clock_breakdown = True + + # Pass the mpu from here to groups. For subsequent use, just query groups + if self.mpu is not None: + groups.mpu = self.mpu + + # Set deepspeed parallelism spec. for the model including expert parallelism + for _, module in self.module.named_modules(): + if hasattr(module, 'set_deepspeed_parallelism'): + module.set_deepspeed_parallelism(self._config.use_data_before_expert_parallel_) + + # Query the groups module to get information about various parallel groups + self.local_all_to_all_group = None + if self.zero_quantized_gradients(): + message = "Using LoCo quantized gradients" if self.zeropp_loco_param() else "Using quantized gradients" + log_dist(message, ranks=[0]) + self.local_all_to_all_group = groups._get_local_all_to_all_group() + self.data_parallel_group = groups._get_data_parallel_group() + self.dp_world_size = groups._get_data_parallel_world_size() + self.seq_data_parallel_group = groups._get_sequence_data_parallel_group() + self.seq_dp_world_size = groups._get_sequence_data_parallel_world_size() + self.mp_world_size = groups._get_model_parallel_world_size() + self.expert_parallel_group = groups._get_expert_parallel_group_dict() + self.expert_data_parallel_group = groups._get_expert_data_parallel_group_dict() + self.sequence_parallel_size = groups._get_sequence_parallel_world_size() + if self.sequence_parallel_size > 1: + self.communication_data_type = self._config.seq_parallel_communication_data_type + self.seq_parallel_group = groups._get_sequence_parallel_group() + + if dist.get_rank() == 0: + summary = "********** distributed groups summary **********\n" + summary += f"\t {self.dp_world_size=}\n" + summary += f"\t {self.mp_world_size=}\n" + summary += f"\t {self.seq_dp_world_size=}\n" + summary += f"\t {self.sequence_parallel_size=}\n" + summary += "***********************************************" + logger.info(summary) + + if not (self.amp_enabled() or is_zero_init_model): + self._broadcast_model() + + # check if parameters are duplicated in optimizer param_groups + def _check_for_duplicates(self, optimizer): + for name, param in self.module.named_parameters(): + param_id = id(param) + + def ids_list(group): + return [id(param) for param in group] + + occurrence = sum([ + ids_list(group['params']).count(param_id) if param_id in ids_list(group['params']) else 0 + for group in optimizer.param_groups + ]) + assert occurrence <= 1, f"Parameter with name: {name} occurs multiple times in optimizer.param_groups. Make sure it only appears once to prevent undefined behavior." + + def _do_optimizer_sanity_check(self, basic_optimizer): + model_dtype, grad_accum_dtype = self.get_data_types() + zero_enabled = self.zero_optimization() + amp_enabled = self.amp_enabled() + # config based assertions + assert ( + not (amp_enabled and zero_enabled) + ), "Amp and ZeRO are not currently compatible, please use (legacy) fp16 mode which performs similar to amp opt_mode=O2" + if zero_enabled: + if not is_zero_supported_optimizer(basic_optimizer): + assert ( + self.zero_allow_untested_optimizer() + ), 'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.' + + if self.global_rank == 0: + logger.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****") + if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage( + ) == 1 and not self.zero_cpu_offload(): + return BFLOAT16 + return ZERO_OPTIMIZATION + elif amp_enabled: + if model_dtype != grad_accum_dtype: + raise NotImplementedError( + "Model data type and gradient accumulation data type must be equal to use Amp") + if model_dtype == torch.bfloat16 or model_dtype == torch.float16: + raise NotImplementedError("Cannot enable both amp with (legacy) fp16 or bfloat16 mode") + try: + logger.info("Initializing Apex amp from: {}".format(amp.__path__)) + except NameError: + # If apex/amp is available it will be imported above + raise RuntimeError("Unable to import apex/amp, please make sure it is installed") + return AMP + # data type checks + elif model_dtype == grad_accum_dtype: + if model_dtype == torch.bfloat16: + if self.pipeline_parallelism: + logger.warning( + "**** BF16 gradient accumulation is not safe numerically with large number of accumulation steps, proceed with caution *****" + ) + return BFLOAT16 + else: + raise NotImplementedError( + "Bfloat16 wrapper must use a gradient accumulation type of fp32, enable ZeRO to use Bfloat16 gradient accumulation" + ) + if model_dtype == torch.float16: + return FP16 + # else optimizer_wrapper = None + elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32: + return BFLOAT16 + else: + raise NotImplementedError("unsupported mix of model dtype and gradient accumulation type") + + return None + + # Configure optimizer + def _configure_optimizer(self, client_optimizer, model_parameters): + if client_optimizer is None: + if self.has_moe_layers: + model_parameters = configure_moe_param_groups(model_parameters) + basic_optimizer = self._configure_basic_optimizer(model_parameters) + log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0]) + else: + if isinstance(client_optimizer, tuple(self._supported_optims())): + basic_optimizer = client_optimizer + log_dist('Using client Optimizer as basic optimizer', ranks=[0]) + else: + basic_optimizer = client_optimizer(model_parameters) + log_dist('Using client callable to create basic optimizer', ranks=[0]) + + if self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam): + if self.zero_force_ds_cpu_optimizer(): + msg = f'You are using ZeRO-Offload with a client provided optimizer ({type(basic_optimizer)}) which in most cases will yield poor performance. Please either use deepspeed.ops.adam.DeepSpeedCPUAdam or set an optimizer in your ds-config (https://www.deepspeed.ai/docs/config-json/#optimizer-parameters). If you really want to use a custom optimizer w. ZeRO-Offload and understand the performance impacts you can also set <"zero_force_ds_cpu_optimizer": false> in your configuration file.' + raise ZeRORuntimeException(msg) + + basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0] + log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0]) + + self._check_for_duplicates(basic_optimizer) + + self.basic_optimizer = basic_optimizer + log_dist(f"DeepSpeed Basic Optimizer = {basic_optimizer.__class__.__name__}", ranks=[0]) + + optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) + + if optimizer_wrapper == ZERO_OPTIMIZATION: + self.optimizer = self._configure_zero_optimizer(basic_optimizer) + elif optimizer_wrapper == AMP: + amp_params = self.amp_params() + log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) + model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params) + self._set_client_model(model) + self._broadcast_model() + # TODO: maybe need to broadcast experts differently? + elif optimizer_wrapper == FP16: + self.optimizer = self._configure_fp16_optimizer(basic_optimizer) + elif optimizer_wrapper == BFLOAT16: + self.optimizer = self._configure_bf16_optimizer(basic_optimizer) + else: + self.optimizer = basic_optimizer + + log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer.__class__.__name__), ranks=[0]) + + self.compression_scheduler = self._configure_compression_scheduler() + self.quantizer = self._configure_quantization() + + def _configure_basic_optimizer(self, model_parameters): + optimizer_parameters = self.optimizer_params() + if optimizer_parameters is None: + optimizer_parameters = {} + # print(optimizer_parameters.keys()) + if "max_grad_norm" in optimizer_parameters.keys(): + raise ValueError( + "'max_grad_norm' is not supported as an optimizer parameter, please switch to using the deepspeed parameter 'gradient_clipping' see: https://www.deepspeed.ai/docs/config-json/#gradient-clipping for more details" + ) + + if self.optimizer_name() in [ADAM_OPTIMIZER, ADAMW_OPTIMIZER]: + torch_adam = optimizer_parameters.pop(TORCH_ADAM_PARAM, False) + adam_w_mode = optimizer_parameters.pop(ADAM_W_MODE, ADAM_W_MODE_DEFAULT) + + # Optimizer name of Adam forces AdamW logic unless adam_w_mode is explicitly set + effective_adam_w_mode = self.optimizer_name() == ADAMW_OPTIMIZER or adam_w_mode + + if torch_adam: + if not effective_adam_w_mode: + optimizer = torch.optim.Adam(model_parameters, **optimizer_parameters) + else: + optimizer = torch.optim.AdamW(model_parameters, **optimizer_parameters) + else: + if self.zero_use_cpu_optimizer(): + from deepspeed.ops.adam import DeepSpeedCPUAdam + optimizer = DeepSpeedCPUAdam(model_parameters, + **optimizer_parameters, + adamw_mode=effective_adam_w_mode) + else: + from deepspeed.ops.adam import FusedAdam + + optimizer = FusedAdam( + model_parameters, + **optimizer_parameters, + adam_w_mode=effective_adam_w_mode, + ) + + elif self.optimizer_name() == ADAGRAD_OPTIMIZER: + if self.zero_use_cpu_optimizer(): + from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad + optimizer = DeepSpeedCPUAdagrad(model_parameters, **optimizer_parameters) + else: + optimizer = torch.optim.Adagrad(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == LAMB_OPTIMIZER: + from deepspeed.ops.lamb import FusedLamb + + optimizer = FusedLamb(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER: + assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO" + from deepspeed.runtime.fp16.onebit.adam import OnebitAdam + + optimizer = OnebitAdam(model_parameters, self, **optimizer_parameters) + if not self.fp16_enabled(): + logger.warning(f"Currently the convergence of 1-bit Adam is only verified under FP16") + elif self.optimizer_name() == ZERO_ONE_ADAM_OPTIMIZER: + assert not self.zero_optimization(), "0/1 Adam is not compatible with ZeRO" + from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam + + optimizer = ZeroOneAdam(model_parameters, self, **optimizer_parameters) + if not self.fp16_enabled(): + logger.warning(f'Currently the convergence of 0/1 Adam is only verified under FP16') + elif self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER: + assert not self.zero_optimization(), "1bit-Lamb is not compatible with ZeRO" + from deepspeed.runtime.fp16.onebit.lamb import OnebitLamb + + optimizer = OnebitLamb(model_parameters, self, **optimizer_parameters) + if not self.fp16_enabled(): + logger.warning(f"Currently the convergence of 1-bit Lamb is only verified under FP16") + elif self.optimizer_name() == LION_OPTIMIZER: + if self.zero_use_cpu_optimizer(): + from deepspeed.ops.lion import DeepSpeedCPULion + optimizer = DeepSpeedCPULion(model_parameters, **optimizer_parameters) + else: + from deepspeed.ops.lion import FusedLion + optimizer = FusedLion(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == MUADAM_OPTIMIZER: + try: + from mup import MuAdam + except ImportError: + logger.error(f"Install mup to use MuAdam optimizer") + optimizer = MuAdam(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == MUADAMW_OPTIMIZER: + try: + from mup import MuAdamW + except ImportError: + logger.error(f"Install mup to use MuAdamW optimizer") + optimizer = MuAdamW(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == MUSGD_OPTIMIZER: + try: + from mup import MuSGD + except ImportError: + logger.error(f"Install mup to use MuSGD optimizer") + optimizer = MuSGD(model_parameters, **optimizer_parameters) + else: + torch_optimizer = getattr(torch.optim, self.optimizer_name()) + optimizer = torch_optimizer(model_parameters, **optimizer_parameters) + return optimizer + + def _configure_compression_scheduler(self): + return compression_scheduler(self.module, self._config.compression_config) + + def _configure_random_ltd_scheduler(self, configs): + return RandomLTDScheduler(configs) + + def _configure_quantization(self): + ( + quantize_weight_in_forward, + quantize_enabled, + q_groups, + q_mixed_fp16, + q_change_ratio, + q_type, + q_rounding, + q_verbose, + use_quantizer_kernel, + ) = self.quantize_training() + if quantize_enabled and not quantize_weight_in_forward: + assert self.fp16_enabled( + ), "MoQ (quantize in optimization step) weight quantization is only supported for FP16" + quantizer = None + if quantize_enabled and not quantize_weight_in_forward: + from deepspeed.runtime.quantize import Quantizer + + quantizer = Quantizer( + q_groups, + q_mixed_fp16, + q_change_ratio, + q_type, + q_rounding, + q_verbose, + self.eigenvalue_enabled(), + use_quantizer_kernel, + self.eigenvalue_layer_num() if self.eigenvalue_enabled() else 0, + ) + return quantizer + + def _configure_fp16_optimizer(self, optimizer): + initial_dynamic_scale = self.initial_dynamic_scale() + dynamic_loss_args = self.dynamic_loss_scale_args() + clip_grad = self.gradient_clipping() + + if APEX_INSTALLED: + fused_opts = (apex.optimizers.FusedAdam, FusedAdam) + else: + fused_opts = FusedAdam + + if isinstance(optimizer, fused_opts) \ + or self.optimizer_name() in [ONEBIT_ADAM_OPTIMIZER, ZERO_ONE_ADAM_OPTIMIZER]: + if self.dynamic_loss_scale(): + log_dist(f'Creating fp16 optimizer with dynamic loss scale', ranks=[0]) + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + optimizer = FP16_Optimizer( + optimizer, + deepspeed=self, + dynamic_loss_scale=True, + initial_dynamic_scale=initial_dynamic_scale, + dynamic_loss_args=dynamic_loss_args, + mpu=self.mpu, + clip_grad=clip_grad, + fused_adam_legacy=self.optimizer_legacy_fusion(), + timers=timers, + has_moe_layers=self.has_moe_layers, + ) + else: + log_dist(f'Creating fp16 optimizer with static loss scale: {self.loss_scale()}', ranks=[0]) + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + optimizer = FP16_Optimizer( + optimizer, + deepspeed=self, + static_loss_scale=self.loss_scale(), + mpu=self.mpu, + clip_grad=clip_grad, + fused_adam_legacy=self.optimizer_legacy_fusion(), + timers=timers, + has_moe_layers=self.has_moe_layers, + ) + else: + log_dist(f'Creating fp16 unfused optimizer with dynamic loss scale', ranks=[0]) + optimizer = FP16_UnfusedOptimizer( + optimizer, + deepspeed=self, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=dynamic_loss_args, + mpu=self.mpu, + clip_grad=clip_grad, + fused_lamb_legacy=self.optimizer_name() == LAMB_OPTIMIZER, + ) + + return optimizer + + def _configure_bf16_optimizer(self, optimizer): + clip_grad = self.gradient_clipping() + + if optimizer is None: + optimizer = DummyOptim(list(self.module.parameters())) + + log_dist('Creating BF16 optimizer', ranks=[0]) + + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + optimizer = BF16_Optimizer(optimizer, + self.param_names, + bfloat16_config=self._config.bfloat16_config, + mpu=self.mpu, + clip_grad=clip_grad, + allgather_bucket_size=self.zero_allgather_bucket_size(), + dp_process_group=self.seq_data_parallel_group, + timers=timers, + grad_acc_dtype=self.get_data_types()[1], + graph_harvesting=self.graph_harvesting(), + has_moe_layers=self.has_moe_layers) + + return optimizer + + def _configure_zero_optimizer(self, optimizer): + zero_stage = self.zero_optimization_stage() + + mics_shard_size = self.mics_shard_size() + model_dtype, gradient_accumulation_dtype = self.get_data_types() + + if self.bfloat16_enabled(): + check_grad_overflow = self._config.bfloat16_config.check_grad_overflow + elif self.fp16_enabled(): + check_grad_overflow = True + else: + check_grad_overflow = False + + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + + if optimizer is None: + optimizer = DummyOptim(list(self.module.parameters())) + + if self.zero_legacy_stage1(): + raise Exception( + "The deprecated version of ZeRO Stage 1 is not supported in deepspeed >= 0.5.9. Please downgrade to a version less than 0.5.9 if you need to use this deprecated version of ZeRO." + ) + + if zero_stage <= ZeroStageEnum.gradients: + overlap_comm = self.zero_overlap_comm() + contiguous_gradients = self.zero_contiguous_gradients() + round_robin_gradients = self.zero_round_robin_gradients() + assert not isinstance(optimizer, DummyOptim), "zero stage {} requires an optimizer".format(zero_stage) + + log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) + + if isinstance(self.module, PipelineModule): + if overlap_comm: + logger.warning("Pipeline parallelism does not support overlapped communication, will be disabled.") + overlap_comm = False + optimizer = DeepSpeedZeroOptimizer( + optimizer, + self.param_names, + timers=timers, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=contiguous_gradients, + reduce_bucket_size=self.zero_reduce_bucket_size(), + use_multi_rank_bucket_allreduce=self.zero_multi_rank_bucket_allreduce(), + allgather_bucket_size=self.zero_allgather_bucket_size(), + dp_process_group=self.seq_data_parallel_group, + expert_parallel_group=self.expert_parallel_group if self.has_moe_layers else None, + expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=overlap_comm, + offload_optimizer_config=self.zero_offload_optimizer(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + ignore_unused_parameters=self.zero_ignore_unused_parameters(), + partition_grads=zero_stage == ZeroStageEnum.gradients, + round_robin_gradients=round_robin_gradients, + has_moe_layers=self.has_moe_layers, + fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type, + elastic_checkpoint=self.zero_elastic_checkpoint(), + check_grad_overflow=check_grad_overflow) + + elif zero_stage == ZeroStageEnum.weights: + assert not self.has_moe_layers, "MoE not supported with Stage 3" + if isinstance(optimizer, DummyOptim): + log_dist("Creating ZeRO Offload", ranks=[0]) + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + if self.zero_hpz_partition_size() > 1 and zero_param_parallel_group is None: + self._set_zero_group_parallelism() + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + optimizer = DeepSpeedZeRoOffload( + self.module, + timers=timers, + ds_config=self.config, + overlap_comm=self.zero_overlap_comm(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + offload_param_config=self.zero_offload_param(), + mpu=self.mpu, + zero_param_parallel_group=zero_param_parallel_group, + zero_quantized_weights=self.zero_quantized_weights(), + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), + zero_module_granularity_threshold=self.zero_module_granularity_threshold(), + log_trace_cache_warnings=self.zero_log_trace_cache_warnings(), + ) + else: + log_dist( + f'Creating fp16 ZeRO stage {zero_stage} optimizer,' + f' MiCS is enabled {mics_shard_size>0},' + f' Hierarchical params gather {self._config.mics_hierarchial_params_gather}', + ranks=[0]) + if mics_shard_size > 0: + return self._return_mics_optimizer(optimizer, timers) + + log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) + from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 + optimizer = DeepSpeedZeroOptimizer_Stage3( + self.module, + optimizer, + timers=timers, + ds_config=self.config, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=self.zero_contiguous_gradients(), + reduce_bucket_size=self.zero_reduce_bucket_size(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + dp_process_group=self.seq_data_parallel_group, + all2all_process_group=self.local_all_to_all_group, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=self.zero_overlap_comm(), + offload_optimizer_config=self.zero_offload_optimizer(), + offload_param_config=self.zero_offload_param(), + sub_group_size=self.zero_sub_group_size(), + offload_ratio=self.zero_partial_offload(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + aio_config=self.aio_config(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type, + zero_hpz_partition_size=self.zero_hpz_partition_size(), + zero_quantized_weights=self.zero_quantized_weights(), + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), + zero_module_granularity_threshold=self.zero_module_granularity_threshold(), + zeropp_loco_param=self.zeropp_loco_param(), + log_trace_cache_warnings=self.zero_log_trace_cache_warnings(), + ) + + else: + raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage)) + + return optimizer + + def _return_mics_optimizer(self, basic_optimizer, timers): + from deepspeed.runtime.zero.mics import MiCS_Optimizer + model_dtype, gradient_accumulation_dtype = self.get_data_types() + optimizer = MiCS_Optimizer(self.module, + basic_optimizer, + timers=timers, + ds_config=self.config, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=self.zero_contiguous_gradients(), + reduce_bucket_size=self.zero_reduce_bucket_size(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + dp_process_group=self.seq_data_parallel_group, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=self.zero_overlap_comm(), + offload_optimizer_config=self.zero_offload_optimizer(), + offload_param_config=self.zero_offload_param(), + sub_group_size=self.zero_sub_group_size(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + aio_config=self.aio_config(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type) + return optimizer + + def _configure_eigenvalue(self): + eigenvalue = Eigenvalue( + verbose=self.eigenvalue_verbose(), + max_iter=self.eigenvalue_max_iter(), + tol=self.eigenvalue_tol(), + stability=self.eigenvalue_stability(), + gas_boundary_resolution=self.eigenvalue_gas_boundary_resolution(), + layer_name=self.eigenvalue_layer_name(), + layer_num=self.eigenvalue_layer_num(), + ) + + return eigenvalue + + def _configure_progressive_layer_drop(self): + pld = ProgressiveLayerDrop(theta=self.pld_theta(), gamma=self.pld_gamma()) + + return pld + + def _configure_curriculum_scheduler_legacy(self): + scheduler = CurriculumScheduler(self.curriculum_params_legacy()) + return scheduler + + @staticmethod + def is_map_style_dataset(obj): + return hasattr(obj, "__getitem__") and hasattr(obj, "__len__") + + @staticmethod + def is_iterable_style_dataset(obj): + return isinstance(obj, torch.utils.data.IterableDataset) # hasattr(obj, "__iter__") should work as well + + def dataloader_drop_last(self): + return self._config.dataloader_drop_last + + def was_step_applied(self) -> bool: + """Returns True if the latest ``step()`` produced in parameter updates. + Note that a ``False`` return is not an error condition. Steps are frequently + no-ops, such as between gradient accumulation boundaries or when overflows + occur. + Returns: + bool: Whether the latest ``step()`` modified model parameters. + """ + return self._step_applied + + def deepspeed_io(self, + dataset, + batch_size=None, + route=ROUTE_TRAIN, + pin_memory=True, + data_sampler=None, + collate_fn=None, + num_local_io_workers=None): + if not (self.is_map_style_dataset(dataset) or self.is_iterable_style_dataset(dataset)): + raise ValueError("Training data must be a torch Dataset") + + if batch_size is None: + batch_size = self.train_micro_batch_size_per_gpu() + + if collate_fn is None: + collate_fn = self.collate_fn + + # Currently we only use timer in train route + deepspeed_io_timer = None + if route == ROUTE_TRAIN: + deepspeed_io_timer = self.tput_timer + + # If mpu is provided, forward world size and parallel rank to sampler. + data_parallel_world_size = self.dp_world_size + data_parallel_rank = self.global_rank + if self.mpu is not None: + data_parallel_world_size = self.mpu.get_data_parallel_world_size() + data_parallel_rank = self.mpu.get_data_parallel_rank() + + if data_sampler is None and (route == ROUTE_PREDICT or route == ROUTE_EVAL): + data_sampler = torch.utils.data.DistributedSampler( + dataset, + num_replicas=data_parallel_world_size, + rank=data_parallel_rank, + shuffle=False, + ) + + deepspeed_dataloader_config = {} + if self.curriculum_learning_enabled(): + deepspeed_dataloader_config = { + CURRICULUM_LEARNING: self.curriculum_learning_enabled(), + DATA_EFFICIENCY: self.data_efficiency_config(), + DATA_PARALLEL_GROUP: self.data_parallel_group, + GRADIENT_ACCUMULATION_STEPS: self.gradient_accumulation_steps(), + GLOBAL_RANK: self.global_rank, + DATA_SAMPLING_NUM_WORKERS: self.data_sampling_config()[DATA_SAMPLING_NUM_WORKERS] + } + return DeepSpeedDataLoader(dataset=dataset, + batch_size=batch_size, + pin_memory=pin_memory, + collate_fn=collate_fn, + local_rank=self.local_rank, + tput_timer=deepspeed_io_timer, + num_local_io_workers=num_local_io_workers, + data_sampler=data_sampler, + data_parallel_world_size=data_parallel_world_size, + data_parallel_rank=data_parallel_rank, + dataloader_drop_last=self.dataloader_drop_last(), + deepspeed_dataloader_config=deepspeed_dataloader_config) + + def train(self, mode=True): + r"""""" + + self.warn_unscaled_loss = True + self.module.train(mode) + + def eval(self): + r"""""" + + self.warn_unscaled_loss = True + self.module.train(False) + + def _scale_loss_by_gas(self, prescaled_loss, eval_micro_batches=None): + # In pipeline evaluation, there is an option to use different micro-bs, which creates different number of + # micro batches, thus the training gas, is not valid in this case. need to use the number of eval_micro_batches + scaling_factor = self.gradient_accumulation_steps() if eval_micro_batches is None else eval_micro_batches + if isinstance(prescaled_loss, torch.Tensor): + scaled_loss = prescaled_loss / scaling_factor + elif isinstance(prescaled_loss, tuple) or isinstance(prescaled_loss, list): + scaled_loss = [] + for l in prescaled_loss: + if isinstance(l, torch.Tensor): + scaled_loss.append(l / scaling_factor) + else: + scaled_loss.append(l) + else: + scaled_loss = prescaled_loss + if self.warn_unscaled_loss: + logger.warning(f"DeepSpeed unable to scale loss because of type: {type(prescaled_loss)}") + self.warn_unscaled_loss = False + + return scaled_loss + + def _create_module_forward_pre_hook(self): + + def _module_forward_pre_hook(module, inputs, kwargs): + return self._forward_prologue(inputs, kwargs) + + return self.module.register_forward_pre_hook(_module_forward_pre_hook, prepend=False, with_kwargs=True) + + def _create_module_forward_post_hook(self): + + def _module_forward_post_hook(module, input, output): + self._forward_epilogue() + + return self.module.register_forward_hook(_module_forward_post_hook) + + def _forward_prologue(self, inputs, kwargs): + return_modified = False + + if not self.autotuning_profile_model_info(): + see_memory_usage("Engine before forward", force=self.memory_breakdown()) + + flops_profiler_active = (self.flops_profiler_enabled() + and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) + + # used to check quantization happens at step 0! + if self.global_steps == 0 and hasattr(self, "compression_scheduler"): + self.compression_scheduler.step(step_zero_check=True) + if self.quantizer: + tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( + ) == 2 else self.optimizer.fp16_groups + if self.compression_scheduler.weight_quantization_enabled: + self.quantizer.quantize( + tensor_to_quantize, + (self.optimizer.overflow if self.fp16_enabled() else False), + self.eigenvalue_enabled(), + None, + ) + return_modified = True + + if flops_profiler_active: + self.flops_profiler.start_profile(ignore_list=None) + + if kwargs is not None: + if self.module.training: + if self.progressive_layer_drop: + kwargs.update(self.progressive_layer_drop.get_state()) + + if self.__class__.__name__ != "PipelineEngine": + # TODO: The above if condition is a HACK since for PipelineEngine + # it's difficult to inject argument in forward pass. + if self.module.training and self.curriculum_enabled_legacy(): + self.curriculum_scheduler_legacy.update_difficulty(self.global_steps + 1) + if self.curriculum_params_legacy()["curriculum_type"] == "seqlen": + kwargs.update({"curriculum_seqlen": self.curriculum_scheduler_legacy.get_current_difficulty()}) + return_modified = True + + if self.module.training and self.random_ltd_enabled(): + self.random_ltd_scheduler.update_seq(self.global_steps) + + if self.training_dataloader is None: + self.tput_timer.start() + + self._start_timers(self.engine_timers.forward_timers) + + if self.zero_optimization_partition_weights(): + # Enable automated discovery of external parameters by indicating that + # we are in a forward pass. + for module in self.module.modules(): + module._parameters._in_forward = True + + if self.fp16_auto_cast(): + inputs = self._cast_inputs_half(inputs) + return_modified = True + + if return_modified: + return inputs, kwargs + + def _forward_epilogue(self): + if self.zero_optimization_partition_weights(): + # Disable automated discovery of external parameters + for module in self.module.modules(): + module._parameters._in_forward = False + + self._stop_timers(self.engine_timers.forward_timers) + + flops_profiler_active = (self.flops_profiler_enabled() + and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) + + if flops_profiler_active: + self.flops_profiler.stop_profile() + + if not self.autotuning_profile_model_info(): + see_memory_usage("Engine after forward", force=self.memory_breakdown()) + + @instrument_w_nvtx + def forward(self, *inputs, **kwargs): + r"""Execute forward propagation + Arguments: + *inputs: Variable length input list + **kwargs: variable length keyword arguments + """ + if self.autotuning_profile_model_info(): + ma = get_ma_status() + + if self.is_deepcompile_enabled() and hasattr(self, "launch_compile_passes"): + # We can't have this in forward prologue as the compiler compiles hooks including the forward prologue. + self.launch_compile_passes(self.global_steps) + + validate_nested_autocast(self) + with torch.autocast(device_type=get_accelerator().device_name(), + dtype=self.torch_autocast_dtype(), + enabled=self.torch_autocast_enabled()): + loss = self.module(*inputs, **kwargs) + + if self.autotuning_profile_model_info(): + activation_mem = get_ma_status() - ma + self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem + print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path()) + exit() + + return loss + + def _cast_inputs_half(self, inputs): + if isinstance(inputs, (list, tuple)): + new_inputs = [] + for v in inputs: + new_inputs.append(self._cast_inputs_half(v)) + return inputs.__class__(new_inputs) + elif isinstance(inputs, dict): + new_inputs = {} + for k, v in inputs.items(): + new_inputs[k] = self._cast_inputs_half(v) + return new_inputs + elif hasattr(inputs, 'half') and inputs.is_floating_point(): + return inputs.half() + else: + return inputs + + def print_forward_breakdown(self, fwd_time): + gate_time = 0.0 + moe_time = 0.0 + falltoall = 0.0 + salltoall = 0.0 + + for gate in self.gate_modules: + #logger.info(f"Individual TopK gate time: {gate.gate_time:.2f} ms") + gate_time += gate.gate_time + + for l in self.moe_layers: + #logger.info(f"MoE layer; total: {l.time_moe:.2f} ms, first alltoall: {l.time_falltoall:.2f}, second alltoall: {l.time_salltoall:.2f}") + moe_time += l.time_moe + falltoall += l.time_falltoall + salltoall += l.time_salltoall + + # TODO: Allreduce/average them across ranks for more accurate timing. + + # if deepspeed.comm.get_rank() == 0: + log_dist( + f"time (ms) | fwd: {fwd_time:.2f} (fwd_moe: {moe_time:.2f}, 1st_a2a: {falltoall:.2f}, 2nd_a2a: {salltoall:.2f}, top_k: {gate_time:.2f})", + ranks=[0]) + + @instrument_w_nvtx + def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): + # Pass (PP) gas boundary flag to optimizer (required for zero) + self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() + # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well + if self.zero_optimization_partition_gradients(): + self.optimizer.overlapping_partition_gradients_reduce_epilogue() + + # Communicate only at gradient accumulation boundaries + elif self.is_gradient_accumulation_boundary(): + if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr( + self.optimizer, 'reduce_gradients'): + self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism) + else: + grads = None + self.buffered_allreduce_fallback(grads=grads, elements_per_buffer=bucket_size) + + def _backward_prologue(self, loss, scale_wrt_gas=True): + see_memory_usage("Engine before backward", force=self.memory_breakdown()) + if self.scale_wrt_gas is not None: + scale_wrt_gas = self.scale_wrt_gas + + # scale loss w.r.t. gradient accumulation if reduction is not disabled + do_gradient_reduction = self.enable_backward_allreduce and not self.inside_no_sync_ctxt and not self.is_deepcompile_enabled( + ) + if do_gradient_reduction and self.gradient_accumulation_steps() > 1 and scale_wrt_gas: + loss = self._scale_loss_by_gas(loss.float()) + + # Log training loss + mean_loss = loss.mean().detach() + self.losses = mean_loss if self.losses is None else self.losses + mean_loss + if self.monitor.enabled: + if self.is_gradient_accumulation_boundary(): + if self.global_rank == 0: + self.summary_events = [( + f"Train/Samples/train_loss", + self.losses.item(), + self.global_samples, + )] + self.monitor.write_events(self.summary_events) + + if self.is_deepcompile_enabled(): + deepcompile_backward_prologue(self.is_gradient_accumulation_boundary()) + + return loss + + def _backward_epilogue(self): + self._start_timers(self.engine_timers.backward_reduce_timers) + if self.enable_backward_allreduce and not self.inside_no_sync_ctxt: + # Traditional code path that allreduces the module parameter grads + self.allreduce_gradients() + + self._stop_timers(self.engine_timers.backward_reduce_timers) + see_memory_usage("Engine after backward", force=self.memory_breakdown()) + + def _do_optimizer_backward(self, loss, retain_graph): + self._start_timers(self.engine_timers.backward_inner_timers) + if self.zero_optimization(): + self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() + self.optimizer.backward(loss, retain_graph=retain_graph) + elif self.amp_enabled(): + # AMP requires delaying unscale when inside gradient accumulation boundaries + # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations + delay_unscale = not self.is_gradient_accumulation_boundary() + with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss: + scaled_loss.backward(retain_graph=retain_graph) + elif self.fp16_enabled(): + if self.eigenvalue_enabled(): + self.optimizer.backward(loss, create_graph=True, retain_graph=True) + else: + self.optimizer.backward(loss, retain_graph=retain_graph) + elif self.bfloat16_enabled(): + self.optimizer.backward(loss, retain_graph=retain_graph) + else: + if self.eigenvalue_enabled(): + loss.backward(create_graph=True, retain_graph=True) + else: + loss.backward(retain_graph=retain_graph) + self._stop_timers(self.engine_timers.backward_inner_timers) + + @contextmanager + def no_sync(self): + r""" + Context manager to disable gradient reduction during backward pass. + This context manager has the following effects on other DeepSpeed features: + 1. Incompatible with ZeRO stage 2/3 which rely on reduction for gradient partitioning. + 2. It is illegal to call engine.step() within the context manager. + 3. Tracking of gradient accumulation steps is disabled. + """ + assert not self.zero_optimization_partition_gradients(), \ + f"no_sync context manager is incompatible with gradient partitioning logic of ZeRO stage {self.zero_optimization_stage()}" + + assert not self.inside_no_sync_ctxt, f"no_sync context manager reentry is unsupported" + + self.inside_no_sync_ctxt = True + try: + yield + finally: + self.inside_no_sync_ctxt = False + + @instrument_w_nvtx + def backward(self, loss, retain_graph=False, scale_wrt_gas=True): + r"""Execute backward pass on the loss + Arguments: + loss: Torch tensor on which to execute backward propagation + retain_graph: bool, default: false + forward on user defined choice of retain_graph + """ + assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ + "must provide optimizer during init in order to use backward" + + self._start_timers(self.engine_timers.backward_timers) + loss = self._backward_prologue(loss, scale_wrt_gas) + self._do_optimizer_backward(loss, retain_graph) + self._backward_epilogue() + self._stop_timers(self.engine_timers.backward_timers) + + return loss + + def is_gradient_accumulation_boundary(self): + """ + Query whether the current micro-batch is at the boundary of + gradient accumulation, and thus will trigger gradient reductions and + an optimizer step. + + Returns: + bool: if the current step is a gradient accumulation boundary. + + """ + if self._is_gradient_accumulation_boundary is None: + return (self.micro_steps + 1) % \ + self.gradient_accumulation_steps() == 0 + else: + return self._is_gradient_accumulation_boundary + + def set_gradient_accumulation_boundary(self, is_boundary): + """ + Manually overrides the DeepSpeed engine's gradient accumulation boundary state, this is an optional + feature and should be used with care. The state should be set before to the intended + value before each forward/backward. The final forward/backward should have the + boundary state set to True. This style allows client code to only call engine.step() once after all + the gradient accumulation passes are complete. See example below: + .. code-block:: python + engine.set_gradient_accumulation_boundary(False) + for _ in range(gradient_accumulation_steps - 1): + micro_batch = next(data_loader) + loss = engine(micro_batch) + engine.backward(loss) + engine.set_gradient_accumulation_boundary(True) + micro_batch = next(data_loader) + loss = engine(micro_batch) + engine.backward(loss) + engine.step() + Arguments: + is_boundary (bool): are we at a gradient accumulation boundary or not? + """ + self._is_gradient_accumulation_boundary = is_boundary + self.optimizer.is_gradient_accumulation_boundary = is_boundary + + def zero_grad(self): + """ + Zero parameter grads. + """ + for param_name, param in self.module.named_parameters(): + param.grad = None + + def clip_fp32_gradients(self): + clip_grad_norm_(parameters=self.module.parameters(), max_norm=self.gradient_clipping(), mpu=self.mpu) + + def _take_model_step(self, lr_kwargs, block_eigenvalue={}): + if self.gradient_clipping() > 0.0: + if not (self.fp16_enabled() or self.bfloat16_enabled() or self.amp_enabled() or self.zero_optimization()): + self.clip_fp32_gradients() + elif self.amp_enabled(): + # AMP's recommended way of doing clipping + # https://nvidia.github.io/apex/advanced.html#gradient-clipping + master_params = amp.master_params(self.optimizer) + clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clipping(), mpu=self.mpu) + self.optimizer.step() + + if hasattr(self.optimizer, '_global_grad_norm'): + self._global_grad_norm = self.optimizer._global_grad_norm + + # Quantize the updated parameter if there is no overflow + if self.quantizer: + tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( + ) == 2 else self.optimizer.fp16_groups + if self.compression_scheduler.weight_quantization_enabled: + self.quantizer.quantize( + tensor_to_quantize, + (self.optimizer.overflow if self.fp16_enabled() else False), + self.eigenvalue_enabled(), + block_eigenvalue, + ) + # zero grad in basic optimizer could be unreliable and may not exhibit + # the behavior that we want + if self.bfloat16_enabled(): + # TODO: Temporary until bf16_optimizer and zero_optimizer are integrated + if self.zero_optimization() and hasattr(self.optimizer, "zero_grad"): + self.optimizer.zero_grad() + else: + pass + elif self.zero_optimization() or self.fp16_enabled() or self.amp_enabled(): + self.optimizer.zero_grad() + else: + self.zero_grad() + + # Check overflow here since in DS fp16 optimizer, the overflow is updated in above step() function. + overflow = False + if hasattr(self.optimizer, "overflow"): + overflow = self.optimizer.overflow + self._step_applied = not overflow + + if overflow: + self.skipped_steps += 1 + else: + self.compression_scheduler.step() + if self.lr_scheduler is not None: + try: + self.lr_scheduler.step(**(lr_kwargs or {})) + except TypeError: + # XXX Hack to work with Megatron 2.0 and DeepSpeed pipelines. + # We don't currently have a way to specify lr_kwargs from + # pipe_engine.train_batch() + self.lr_scheduler.step(self.train_batch_size()) + + if self.steps_per_print() is not None: + report_progress = self.global_rank == 0 if self.global_rank else True + if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0: + self._report_progress(self.global_steps + 1) + + self.losses = None + self.global_steps += 1 + self.global_samples += self.train_batch_size() + + def step(self, lr_kwargs=None): + r"""Execute the weight update step after forward and backward propagation + on effective_train_batch. + """ + assert not self.inside_no_sync_ctxt, \ + "It is illegal to call Engine.step() inside no_sync context manager" + + see_memory_usage("Engine before step", force=self.memory_breakdown()) + + # Check early because self.global_steps is incremented at some point here. + # TODO: Delay self.global_steps increment until very end of this function. + flops_profiler_active = self.flops_profiler_enabled( + ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 + + self._start_timers(self.engine_timers.step_timers) + + assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ + "must provide optimizer during init in order to use step" + + report_progress = False + + self._step_applied = False # assume False, will flip to True + + # Update the model when we reach gradient accumulation boundaries + if self.is_gradient_accumulation_boundary(): + self.gas_boundary_ctr += 1 + + if self.checkpoint_engine.is_decoupled(): + self._commit_decoupled_checkpoint() + + if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) + and self.quantizer.any_precision_switch()): + log_dist(f"computing eigenvalue...", ranks=[0]) + self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, + self.optimizer.cur_scale) + + if self.progressive_layer_drop: + self.progressive_layer_drop.update_state(self.global_steps) + + if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() + and self.quantizer.any_precision_switch()): + self._take_model_step(lr_kwargs, self.block_eigenvalue) + else: + self._take_model_step(lr_kwargs) + + report_progress = self.global_rank == 0 if self.global_rank else True + + self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) + + self._stop_timers(self.engine_timers.step_timers) + + # Log learning rate + if self.monitor.enabled: + if self.is_gradient_accumulation_boundary(): + if self.global_rank == 0: + self.summary_events = [(f"Train/Samples/lr", self.get_lr()[0], self.global_samples)] + + if self.fp16_enabled() and hasattr(self.optimizer, "cur_scale"): + self.summary_events.append(( + f"Train/Samples/loss_scale", + self.optimizer.cur_scale, + self.global_samples, + )) + + if (self.eigenvalue_enabled() + and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()): + ev_values = self.block_eigenvalue.values() + for i in range(len(ev_values)): + self.summary_events.append(( + f"Train/Eigenvalues/ModelBlockParam_{i}", + self.ev_values[i][0], + self.global_samples, + )) + self.monitor.write_events(self.summary_events) + + # Check flops profiling + if flops_profiler_active: + if self.autotuning_enabled(): + self.flops = self.flops_profiler.get_total_flops() * 3 + self.fwd_duration = self.flops_profiler.get_total_duration() + else: + self.flops_profiler.print_model_profile( + profile_step=self.global_steps, + module_depth=self.flops_profiler_module_depth(), + top_modules=self.flops_profiler_top_modules(), + detailed=self.flops_profiler_detailed(), + output_file=self.flops_profiler_output_file(), + ) + self.flops_profiler.end_profile() + + if self.autotuning_enabled() and self.global_steps == (self.autotuning_end_profile_step() + 1): + self._autotuning_exit() + + if self.wall_clock_breakdown(): + # Log micro timing and reset + self.timers.log(names=self.engine_timers.micro_timers, memory_breakdown=self.memory_breakdown()) + + if self.wall_clock_breakdown() or self.flops_profiler_enabled(): + # Log global timing and reset + if self.is_gradient_accumulation_boundary(): + if self.monitor.enabled: + self._write_monitor() + + if self.has_moe_layers: + fwd_time = self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False) + self.print_forward_breakdown(fwd_time=fwd_time) + + self.timers.log(self.engine_timers.global_timers) + + self.micro_steps += 1 + see_memory_usage("Engine after step", force=self.memory_breakdown()) + + def _start_timers(self, timer_names): + for name in timer_names: + self.timers(name).start() + + def _stop_timers(self, timer_names): + record = self.is_gradient_accumulation_boundary() and \ + self.flops_profiler_enabled() and \ + (self.global_steps >= self.flops_profiler_profile_step()) + for name in timer_names: + self.timers(name).stop(record=record) + + def _autotuning_exit(self): + if self.global_rank == 0: + msg = self.timers.get_mean([ + FORWARD_GLOBAL_TIMER, + BACKWARD_GLOBAL_TIMER, + STEP_GLOBAL_TIMER, + ], reset=False) + titer = 0.0 + titer += msg[FORWARD_GLOBAL_TIMER] if FORWARD_GLOBAL_TIMER in msg else 0 + titer += msg[BACKWARD_GLOBAL_TIMER] if BACKWARD_GLOBAL_TIMER in msg else 0 + titer += msg[STEP_GLOBAL_TIMER] if STEP_GLOBAL_TIMER in msg else 0 + titer *= self.gradient_accumulation_steps() + msg["latency"] = titer + msg["FLOPS_per_gpu"] = self.flops * 1_000_000 * self.gradient_accumulation_steps() / titer + msg["throughput"] = self.train_batch_size() * 1_000_000 / \ + msg["latency"] + print_json_dist(msg, [0], path=self.autotuning_metric_path()) + log_dist( + f"Wrote metrics to {self.autotuning_metric_path()}, {os.path.abspath(self.autotuning_metric_path())}", + ranks=[0]) + import atexit + atexit.register(print, "Autotuning: done with running current ds config.") + exit() + + def _write_monitor(self): + if self.global_rank == 0: + self.summary_events = [ + ( + f"Train/Samples/elapsed_time_ms_forward", + self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_backward", + self.timers(BACKWARD_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_backward_inner", + self.timers(BACKWARD_INNER_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_backward_allreduce", + self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_step", + self.timers(STEP_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ] + self.monitor.write_events(self.summary_events) + + def _get_optimizer_param(self, param_name): + result = [] + if not self.optimizer: + return result + for group in self.optimizer.param_groups: + if param_name in group: + result.append(group[param_name]) + else: + result.append(0.0) + return result + + def get_lr(self): + return self._get_optimizer_param("lr") + + def get_type(self): + return self._get_optimizer_param("type") + + def get_mom(self): + if self.optimizer_name() in ["SGD", "RMSprop"]: + return self._get_optimizer_param("momentum") + else: + return self._get_optimizer_param("betas") + + def get_pld_theta(self): + if self.progressive_layer_drop: + return self.progressive_layer_drop.get_theta() + else: + return None + + def _report_progress(self, step): + lr = self.get_lr() + mom = self.get_mom() + log_dist(f"step={step}, skipped={self.skipped_steps}, lr={lr}, mom={mom}", ranks=[0]) + + def allreduce_bucket(self, bucket, dp_group, dp_world_size=None): + tensor = self.flatten(bucket) + + tensor_to_allreduce = tensor + + if self.communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(self.communication_data_type) + + if dp_world_size is None: + dp_world_size = dist.get_world_size(group=dp_group) + if self.postscale_gradients(): + if self.gradient_predivide_factor() != 1.0: + tensor_to_allreduce.mul_(1.0 / self.gradient_predivide_factor()) + + dist.all_reduce(tensor_to_allreduce, group=dp_group) + if self.gradient_average: + if self.gradient_predivide_factor() != dp_world_size: + tensor_to_allreduce.mul_(self.gradient_predivide_factor() / dp_world_size) + else: + tensor_to_allreduce.mul_(1. / dp_world_size) + dist.all_reduce(tensor_to_allreduce, group=dp_group) + + if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + tensor.copy_(tensor_to_allreduce) + + return tensor + + def allreduce_and_copy(self, small_bucket, dp_group, dp_world_size=None): + allreduced = self.allreduce_bucket(small_bucket, dp_group, dp_world_size) + for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + buf.copy_(synced) + + def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000, dp_world_size=None): + small_bucket = [] + numel = 0 + for tensor in bucket: + small_bucket.append(tensor) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy(small_bucket, dp_group, dp_world_size) + small_bucket = [] + numel = 0 + if len(small_bucket) > 0: + self.allreduce_and_copy(small_bucket, dp_group, dp_world_size) + + def _get_gradients_for_reduction(self): + non_expert_grads = [] + expert_grads = {} + if self.has_moe_layers: + for key in self.expert_data_parallel_group.keys(): + expert_grads[key] = [] + + for param_name, param in self.module.named_parameters(): + if not param.requires_grad: + continue + + if param.grad is None: + # In cases where there is an imbalance of empty grads across + # ranks we must create empty grads, this will ensure that every + # rank is reducing the same size. In some cases it may make + # sense in the future to support the ability to average not + # w.r.t. world size but with a different value. + param.grad = torch.zeros(param.size(), dtype=param.dtype, device=param.device) + + grad_data = param.grad.data + if param_name in self.sparse_tensor_module_names or grad_data.is_sparse: + # Call param.grad without data to avoid problem with setting of updated grads + grad_data = SparseTensor(param.grad) + + if is_moe_param(param): + expert_grads[param.group_name].append(grad_data) + else: + non_expert_grads.append(grad_data) + + return non_expert_grads, expert_grads + + def _reduce_non_expert_gradients(self, grads, elements_per_buffer): + split_sparse_tensor_buckets, split_dense_tensor_buckets = split_half_float_double_sparse(grads) + if self.pipeline_parallelism: + dp_group = self.mpu.get_data_parallel_group() + dp_world_size = dist.get_world_size(dp_group) + else: + dp_group = groups._get_sequence_data_parallel_group() + dp_world_size = dist.get_world_size(dp_group) / float(self.sequence_parallel_size) + for _, sparse_bucket_tuple in enumerate(split_sparse_tensor_buckets): + if sparse_bucket_tuple: + bucket_type, sparse_bucket = sparse_bucket_tuple + self.sparse_allreduce_no_retain(sparse_bucket, dp_group=dp_group, dp_world_size=dp_world_size) + + for _, dense_bucket_tuple in enumerate(split_dense_tensor_buckets): + if dense_bucket_tuple: + bucket_type, dense_bucket = dense_bucket_tuple + self.allreduce_no_retain(dense_bucket, + dp_group=dp_group, + numel_per_bucket=elements_per_buffer, + dp_world_size=dp_world_size) + + def _reduce_expert_gradients(self, expert_grads, elements_per_buffer): + # to maintain the gradients value unaffected by ep_size setting, + # utilize dp_world_size for allreduce average + dp_world_size = dist.get_world_size(groups._get_data_parallel_group()) + for ep_name, expert_grads_group in expert_grads.items(): + ep_dp_group = groups._get_expert_data_parallel_group(ep_name) + split_sparse_tensor_buckets, split_dense_tensor_buckets = split_half_float_double_sparse( + expert_grads_group) + + for _, sparse_bucket_tuple in enumerate(split_sparse_tensor_buckets): + if sparse_bucket_tuple: + bucket_type, sparse_bucket = sparse_bucket_tuple + self.sparse_allreduce_no_retain(sparse_bucket, dp_group=ep_dp_group, dp_world_size=dp_world_size) + + for _, dense_bucket_tuple in enumerate(split_dense_tensor_buckets): + if dense_bucket_tuple: + bucket_type, dense_bucket = dense_bucket_tuple + # Separate between diff groups + self.allreduce_no_retain(dense_bucket, + dp_group=ep_dp_group, + numel_per_bucket=elements_per_buffer, + dp_world_size=dp_world_size) + + def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000): + if grads is None: + if hasattr(self.optimizer, "get_grads_for_reduction"): + # This is currently for BF16 optimizer + non_expert_grads, expert_grads = self.optimizer.get_grads_for_reduction() + else: + non_expert_grads, expert_grads = self._get_gradients_for_reduction() + else: + assert not self.has_moe_layers, "attempting to reduce grads in unsupported way w.r.t. MoE" + non_expert_grads = grads + + self._reduce_non_expert_gradients(non_expert_grads, elements_per_buffer) + + if self.has_moe_layers: + self._reduce_expert_gradients(expert_grads, elements_per_buffer) + + def sparse_allreduce_no_retain(self, bucket, dp_group, dp_world_size=None): + allreduced_sparses = self.sparse_allreduce_bucket(bucket, dp_group, dp_world_size) + # Densify sparse tensor and copy back to original location + for tensor in allreduced_sparses: + if tensor.is_sparse: + tensor.orig_dense_tensor.data = tensor.to_coo_tensor() + else: + tensor.orig_dense_tensor.copy_(tensor.to_dense()) + + def sparse_allreduce_bucket(self, bucket, dp_group, dp_world_size=None): + sparse_list = [] + for sparse in bucket: + sparse_list.append(self.sparse_allreduce(sparse, dp_group, dp_world_size)) + return sparse_list + + def sparse_allreduce(self, sparse, dp_group, dp_world_size=None): + original_data_type = sparse.values.dtype + if self.communication_data_type != sparse.values.dtype: + if self.communication_data_type in (torch.float16, torch.bfloat16): + indices = sparse.indices.to(torch.int32) + else: + indices = sparse.indices + values = sparse.values.to(self.communication_data_type) + else: + indices = sparse.indices + values = sparse.values + + if dp_world_size is None: + dp_world_size = dist.get_world_size(group=dp_group) + if self.postscale_gradients(): + if self.gradient_average: + values.mul_(self.gradient_predivide_factor() / (dp_world_size)) + else: + values.mul_(1. / (dp_world_size)) + + indices_device_list = self.sparse_all_gather(indices, dp_group) + values_device_list = self.sparse_all_gather(values, dp_group) + + sparse.indices = torch.cat(indices_device_list).to(torch.long) + sparse.values = torch.cat(values_device_list).to(original_data_type) + return sparse + + def sparse_all_gather(self, value, dp_group): + my_size = torch.LongTensor([value.size()[0]]).to(self.device) + all_sizes = self.all_gather_scalar(my_size, dp_group) + max_size = torch.cat(all_sizes).max() + fill_size = max_size - my_size + + assert value.dim() in [1, 2] + if value.dim() == 1: + if fill_size > 0: + value = torch.cat([value, value.new_empty(fill_size)]) + tensor_list = [value.new_empty(max_size) for _ in range(dist.get_world_size(group=dp_group))] + else: + if fill_size > 0: + value = torch.cat([value, value.new_empty(fill_size, value.size()[1])]) + tensor_list = [ + value.new_empty(max_size, + value.size()[1]) for _ in range(dist.get_world_size(group=dp_group)) + ] + + dist.all_gather(tensor_list, value, group=dp_group) + tensors = [] + for dev_idx, t in enumerate(tensor_list): + size = all_sizes[dev_idx][0] + tensors.append(t.index_select(0, torch.arange(size, dtype=torch.long, device=self.device))) + + return tensors + + def all_gather_scalar(self, value, dp_group): + tensor_list = [value.new_zeros(value.size()) for _ in range(dist.get_world_size(group=dp_group))] + dist.all_gather(tensor_list, value, group=dp_group) + return tensor_list + + def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclude_frozen_parameters=False): + sd = self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) + + # Remove frozen parameter weights from state_dict if specified + if exclude_frozen_parameters: + for n, p in self.module.named_parameters(): + if not p.requires_grad and n in sd: + del sd[n] + + if self.random_ltd_enabled(): + sd = remove_random_ltd_state_dict(sd) + return sd + + @staticmethod + def load_moe_state_dict(checkpoint_path, + tag, + state_dict, + old_moe_load, + model=None, + mpu=None, + num_experts=1, + checkpoint_engine=TorchCheckpointEngine()): + if old_moe_load: + expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name()) + + num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size( + groups._get_max_expert_size_name()) + for local_expert_id in range(num_local_experts): + global_expert_id = expp_rank * num_local_experts + local_expert_id + expert_state_dict = checkpoint_engine.load( + DeepSpeedEngine._get_expert_ckpt_name( + checkpoint_path, + -1, # -1 means ignore layer_id + global_expert_id, + tag, + mpu), + map_location=torch.device('cpu')) + + # Updating global -> local expert ids + moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' + for key in list(expert_state_dict.keys()): + local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', + f'{moe_str_prefix}{local_expert_id}') + expert_state_dict[local_key] = expert_state_dict.pop(key) + state_dict.update(expert_state_dict) + + else: + moe_layer_id = 0 + for n_module, module in model.named_modules(): + if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: + group_name = module.expert_group_name + num_local_experts = module.num_local_experts + expp_rank = groups._get_expert_parallel_rank(group_name) + # loop all local_experts + for local_expert_id in range(num_local_experts): + global_expert_id = expp_rank * num_local_experts + local_expert_id + expert_state_dict = checkpoint_engine.load(DeepSpeedEngine._get_expert_ckpt_name( + checkpoint_path, moe_layer_id, global_expert_id, tag, mpu), + map_location=torch.device('cpu')) + # print(expert_state_dict.keys()) + # Updating global -> local expert ids + moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' + for key in list(expert_state_dict.keys()): + local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', + f'{moe_str_prefix}{local_expert_id}') + expert_state_dict[local_key] = expert_state_dict.pop(key) + state_dict.update(expert_state_dict) + moe_layer_id += 1 + + def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): + if fetch_z3_params: + params_to_fetch = [ + p for p in self.module.parameters() + if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE + ] + else: + params_to_fetch = [] + + with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0): + module_state_dict = checkpoint['module'] + if custom_load_fn: + custom_load_fn(src=module_state_dict, dst=self.module) + else: + self.module.load_state_dict( + module_state_dict, # TODO + strict=strict) + + if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: + saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] + for param in self.module.parameters(): + if param.requires_grad: + continue + if param not in self.param_names: + raise ValueError(f"failed to find frozen {param} in named params") + name = self.param_names[param] + if hasattr(param, 'ds_id'): + param.ds_tensor.data.copy_(saved_frozen_params[name].data) + else: + param.data.copy_(saved_frozen_params[name].data) + + def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode): + return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}' + + def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode): + file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode) + zero_ckpt_name = os.path.join( + checkpoints_path, + str(tag), + f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt", + ) + return zero_ckpt_name + + def _get_zero_ckpt_name(self, checkpoints_path, tag): + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + pp_rank = dist.get_rank(group=self.optimizer.dp_process_group) + bf16_mode = self.bfloat16_enabled() + return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode) + + def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None): + if mp_placeholder is not None: + mp_rank_str = mp_placeholder + else: + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + mp_rank_str = f"{mp_rank:02d}" + + if self.zero_optimization_partition_weights(): + if self.load_universal_checkpoint(): + filename = "zero_pp_rank_0" + else: + filename = "zero_pp_rank_{}".format(dist.get_rank(group=self.optimizer.dp_process_group)) + ckpt_name = os.path.join( + checkpoints_path, + str(tag), + f"{filename}_mp_rank_{mp_rank_str}_model_states.pt", + ) + else: + ckpt_name = os.path.join( + checkpoints_path, + str(tag), + "mp_rank_" + mp_rank_str + "_model_states.pt", + ) + return ckpt_name + + def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank): + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + ckpt_name = os.path.join(checkpoints_path, str(tag), + f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt') + return ckpt_name + + @staticmethod + def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None): + mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank() + if layer_id <= -1: + # Used to support old checkpoint loading + ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), + f'expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') + else: + # Used to support new checkpoint loading + ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), + f'layer_{layer_id}_expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') + return ckpt_name + + def _get_all_ckpt_names(self, checkpoints_path, tag): + # It is required that (checkpoints_path, tag) are consistent among all ranks. + ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*") + import glob + + ckpt_files = glob.glob(ckpt_file_pattern) + ckpt_files.sort() + return ckpt_files + + def load_checkpoint(self, + load_dir, + tag=None, + load_module_strict=True, + load_optimizer_states=True, + load_lr_scheduler_states=True, + load_module_only=False, + custom_load_fn=None): + """ + Load training checkpoint + + Arguments: + load_dir: Required. Directory to load the checkpoint from + tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file + load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. + load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance + load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. + load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. + custom_load_fn: Optional. Custom model load function. + + Returns: + A tuple of ``load_path`` and ``client_state``. + *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. + *``client_state``: State dictionary used for loading required training states in the client code. + + Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right + after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and + ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine + before ``load_checkpoint()``. + + """ + + if tag is None: + latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" + latest_path = os.path.join(load_dir, latest_tag) + if os.path.isfile(latest_path): + with open(latest_path, "r") as fd: + tag = fd.read().strip() + else: + if self.load_universal_checkpoint(): + raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') + else: + logger.warning( + f"Unable to find latest file at {latest_path}, if trying to load latest " + "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." + ) + return None, None + + if self._optimizer_has_ckpt_event_prologue(): + # Prepare for checkpoint load by ensuring all parameters are partitioned + self.optimizer.checkpoint_event_prologue() + + load_path, client_states = self._load_checkpoint(load_dir, + tag, + load_module_strict=load_module_strict, + load_optimizer_states=load_optimizer_states, + load_lr_scheduler_states=load_lr_scheduler_states, + load_module_only=load_module_only, + custom_load_fn=custom_load_fn) + + load_zero_checkpoint = load_path is not None and (self.zero_optimization() or self.bfloat16_enabled()) + if load_zero_checkpoint: + if (load_optimizer_states and not load_module_only) or self.load_universal_checkpoint(): + success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states) + else: + success = False + if not success: + self.optimizer._restore_from_bit16_weights() + + if self.zero_nvme_offload_optimizer(): + from shutil import copytree, disk_usage + offload_dir = self.optimizer.optimizer_swapper.swap_folder + offload_ckpt_dir = os.path.join(load_dir, tag, "offloaded_tensors") + _, _, free = disk_usage(offload_dir) + logger.info( + f"Copying NVMe offload checkpoint from {offload_ckpt_dir} to {offload_dir}, {free / 1e9:,.2f} GB free on target filesystem..." + ) + copytree(offload_ckpt_dir, offload_dir, dirs_exist_ok=True) + _, _, free = disk_usage(offload_dir) + logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem") + self.optimizer.reset_swap_buffers() + + if self._optimizer_has_ckpt_event_epilogue(): + self.optimizer.checkpoint_event_epilogue() + + if self.load_universal_checkpoint() and not self.zero_optimization_partition_weights(): + self.optimizer.update_lp_params() + + return load_path, client_states + + def _load_checkpoint(self, + load_dir, + tag, + load_module_strict=True, + load_optimizer_states=True, + load_lr_scheduler_states=True, + load_module_only=False, + custom_load_fn=None): + + from deepspeed.runtime.state_dict_factory import SDLoaderFactory + + ckpt_list = self._get_all_ckpt_names(load_dir, tag) + sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine) + + is_pipe_parallel = isinstance(self.module, PipelineModule) + + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel) + + if checkpoint is None: + return None, None + + fetch_z3_params = False + if self.zero_optimization_partition_weights() and not load_optimizer_states: + checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) + fetch_z3_params = True + + if is_pipe_parallel: + # Pipeline parallelism uses this to load its own checkpoint files. + self._curr_ckpt_path = os.path.join(load_dir, tag) + + if self.has_moe_layers: + # print(checkpoint.keys()) + old_moe_load = False + if not isinstance(checkpoint['num_experts'], list): + old_moe_load = True + DeepSpeedEngine.load_moe_state_dict(load_dir, + tag, + state_dict=checkpoint['module'], + old_moe_load=old_moe_load, + model=self.module, + mpu=self.mpu, + num_experts=self.num_experts, + checkpoint_engine=self.checkpoint_engine) + if not self.load_universal_checkpoint(): + self.load_module_state_dict(checkpoint=checkpoint, + strict=load_module_strict, + custom_load_fn=custom_load_fn, + fetch_z3_params=fetch_z3_params) + + self.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size'] + + optim_checkpoint = None + if load_module_only: + deepspeed_states = ['module'] + if self.optimizer is not None and hasattr(self.optimizer, 'refresh_fp32_params'): + self.optimizer.refresh_fp32_params() + else: + has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() + if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state: + if self.has_moe_layers: + largest_group_name = groups._get_max_expert_size_name() + expp_rank = groups._get_expert_parallel_rank(largest_group_name) + optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank) + optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu')) + else: + optim_checkpoint = checkpoint + + if self.fp16_enabled() or self.bfloat16_enabled(): + self.optimizer.load_state_dict(optim_checkpoint['optimizer'], + load_optimizer_states=load_optimizer_states) + else: + optim_checkpoint = checkpoint + + self.optimizer.load_state_dict(optim_checkpoint['optimizer']) + + if load_lr_scheduler_states and self.lr_scheduler is not None: + self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + + if self.random_ltd_enabled() and self.random_ltd_scheduler is not None and 'random_ltd' in checkpoint: + self.random_ltd_scheduler.load_state_dict(checkpoint['random_ltd']) + + if self.training_dataloader is not None and self.curriculum_learning_enabled( + ) and 'data_sampler' in checkpoint: + self.training_dataloader.data_sampler.load_state_dict(checkpoint['data_sampler']) + + def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters, loaded_parameters): + result = set() + + for name in original_set: + if name in loaded_parameters and name not in loaded_set: + continue # parameter existed in previous model and was not sparse + result.add(name) + + for name in loaded_set: + if name in original_parameters: + result.add(name) # parameter exists in both configs and it was sparse + + return result + + if 'sparse_tensor_module_names' in checkpoint: + sparse_tensor_module_names = checkpoint['sparse_tensor_module_names'] + elif 'csr_tensor_module_names' in checkpoint: + sparse_tensor_module_names = checkpoint['csr_tensor_module_names'] + else: + sparse_tensor_module_names = None + if sparse_tensor_module_names is not None: + if load_module_strict: + self.sparse_tensor_module_names = sparse_tensor_module_names + else: + self.sparse_tensor_module_names = get_sparse_tensor_module_names( + self.sparse_tensor_module_names, sparse_tensor_module_names, + dict(self.module.named_parameters()), checkpoint["module"]) + + self.global_steps = checkpoint['global_steps'] + self.global_samples = checkpoint.get('global_samples', self.global_steps * self.train_batch_size()) + self.skipped_steps = checkpoint['skipped_steps'] + self.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size'] + deepspeed_states = [ + 'module', 'sparse_tensor_module_names', 'skipped_steps', 'global_steps', 'dp_world_size', + 'mp_world_size', 'data_sampler', 'random_ltd' + ] + client_state = {} + + if load_lr_scheduler_states: + deepspeed_states.append('lr_scheduler') + if load_optimizer_states: + deepspeed_states.append('optimizer') + + client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states} + + if optim_checkpoint is not None: + client_state['optimizer'] = optim_checkpoint['optimizer'] + + return load_path, client_state + + def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): + + load_serial = None + # When use loading checkpoint serial, checkpoint loading start from local rank 0, + # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. + if self._config.zero_config.pipeline_loading_checkpoint: + assert self.zero_optimization_stage( + ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" + load_serial = torch.zeros(1).to(self.device) + if dist.get_local_rank() != 0: + dist.recv(tensor=load_serial, src=dist.get_rank() - 1) + if self.load_universal_checkpoint(): + zero_sd_list = None + checkpoint_folder = f'{os.path.join(load_dir, tag)}' + else: + if load_optimizer_states and self.seq_dp_world_size != self.loaded_checkpoint_dp_world_size: + raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ + f"world size of {self.loaded_checkpoint_dp_world_size} but the " \ + f"current world size is {self.seq_dp_world_size}. Automatic adjustment " \ + "of ZeRO's optimizer state partitioning with a new world size is not " \ + "currently supported.") + checkpoint_folder = None + zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) + if zero_sd_list is None: + return False + + param_shapes = self._get_zero_param_shapes() + self.optimizer.load_state_dict(state_dict_list=zero_sd_list, + load_optimizer_states=load_optimizer_states, + load_from_fp32_weights=self.zero_load_from_fp32_weights(), + checkpoint_folder=checkpoint_folder, + load_serial=load_serial, + param_shapes=param_shapes) + + if self.load_universal_checkpoint(): + logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') + else: + logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") + return True + + def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): + zero_ckpt_names = [] + for dp_rank in range(dp_world_size): + ckpt_name = self._get_rank_zero_ckpt_name(checkpoints_path=load_dir, + tag=tag, + mp_rank=mp_rank, + dp_rank=dp_rank, + bf16_mode=bf16_mode) + zero_ckpt_names.append(ckpt_name) + + return zero_ckpt_names + + def _get_all_zero_checkpoint_names(self, load_dir, tag, bf16_mode): + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + zero_ckpt_names = self._get_mp_rank_zero_checkpoint_names(load_dir=load_dir, + tag=tag, + mp_rank=mp_rank, + dp_world_size=self.loaded_checkpoint_dp_world_size, + bf16_mode=bf16_mode) + for i, ckpt_name in enumerate(zero_ckpt_names): + if not os.path.exists(ckpt_name): + # transparently handle the old file pattern for optim_states + if "optim_states.pt" in ckpt_name: + ckpt_name_try = ckpt_name.replace("_optim_states.pt", "optim_states.pt") + if os.path.exists(ckpt_name_try): + zero_ckpt_names[i] = ckpt_name_try + continue + + return zero_ckpt_names + + def _get_all_zero_checkpoint_state_dicts(self, zero_ckpt_names): + zero_sd_list = [] + for i, ckpt_name in enumerate(zero_ckpt_names): + _state = None + if ckpt_name is None: + _state = {OPTIMIZER_STATE_DICT: None} + # Fully load state for current rank + elif self.zero_elastic_checkpoint() or dist.get_rank(group=self.optimizer.dp_process_group) == i: + _state = self.checkpoint_engine.load( + ckpt_name, + map_location='cpu', + ) + else: + _state = {OPTIMIZER_STATE_DICT: None} + zero_sd_list.append(_state) + + zero_optimizer_sd = [sd[OPTIMIZER_STATE_DICT] for sd in zero_sd_list] + logger.info(f"successfully read {len(zero_optimizer_sd)} ZeRO state_dicts for rank {self.global_rank}") + return zero_optimizer_sd + + def _get_all_zero_checkpoints(self, load_dir, tag): + for bf16_mode in [self.bfloat16_enabled(), not self.bfloat16_enabled()]: + zero_ckpt_names = self._get_all_zero_checkpoint_names(load_dir, tag, bf16_mode) + if zero_ckpt_names is not None: + # Warn if loading checkpoint of different bit16 type + if bf16_mode is not self.bfloat16_enabled(): + checkpoint_bit16 = BFLOAT16 if bf16_mode else FP16 + engine_bit16 = BFLOAT16 if self.bfloat16_enabled() else FP16 + logger.warning(f'Loading {checkpoint_bit16} zero checkpoints into {engine_bit16} training engine') + return self._get_all_zero_checkpoint_state_dicts(zero_ckpt_names) + + return None + + def _checkpoint_tag_validation(self, tag): + if self.checkpoint_tag_validation_enabled(): + s_hash = hashlib.sha1(tag.encode()) + bhash = torch.ByteTensor([s_hash.digest()]).flatten().to(self.device) + max_bhash = bhash.clone() + min_bhash = bhash.clone() + dist.all_reduce(max_bhash, op=dist.ReduceOp.MAX) + dist.all_reduce(min_bhash, op=dist.ReduceOp.MIN) + valid = all(min_bhash == bhash) and all(max_bhash == bhash) + msg = (f"[rank={dist.get_rank()}] The checkpoint tag name '{tag}' is not consistent across " + "all ranks. Including rank unique information in checkpoint tag could cause issues when " + "restoring with different world sizes.") + if self.checkpoint_tag_validation_fail(): + assert valid, msg + elif not valid: + logger.warning(msg) + + def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False): + """Save training checkpoint + + Arguments: + save_dir: Required. Directory for saving the checkpoint + tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is + used if not provided. Tag name must be the same across all ranks. + client_state: Optional. State dictionary used for saving required training states in the client code. + save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint. + exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. + Important: all processes must call this method and not just the process with rank 0. It is + because each process needs to save its master weights and scheduler+optimizer states. This + method will hang waiting to synchronize with other processes if it's called just for the + process with rank 0. + + """ + if self._optimizer_has_ckpt_event_prologue(): + # Custom preparation for checkpoint save, if applicable + self.optimizer.checkpoint_event_prologue() + + rank = self.local_rank if self.use_node_local_storage() else self.global_rank + + # This is to make sure the checkpoint names are created without collision + # There seems to be issue creating them in parallel + + # Ensure save_dir directory exists + if rank == 0: + self.checkpoint_engine.makedirs(save_dir, exist_ok=True) + dist.barrier() + + if tag is None: + tag = f"global_step{self.global_steps}" + + # Ensure tag is a string + tag = str(tag) + commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=save_latest) + + self.checkpoint_engine.create(commit_info) + + # Ensure checkpoint tag is consistent across ranks + self._checkpoint_tag_validation(tag) + + if self.has_moe_layers: + self.save_non_zero_checkpoint = False + self._create_checkpoint_file(save_dir, tag, False) + self._save_moe_checkpoint(save_dir, + tag, + client_state=client_state, + exclude_frozen_parameters=exclude_frozen_parameters) + + # We distribute the task of saving layer checkpoint files among + # data parallel instances, so all procs should call _save_checkpoint. + # All procs then call module_state_dict(), but only procs of data + # parallel rank 0 save the general model params. + if not self.has_moe_layers: + self._create_checkpoint_file(save_dir, tag, False) + self._save_checkpoint(save_dir, + tag, + client_state=client_state, + exclude_frozen_parameters=exclude_frozen_parameters) + + if self.save_zero_checkpoint: + self._create_zero_checkpoint_files(save_dir, tag) + self._save_zero_checkpoint(save_dir, tag) + + if self.zero_nvme_offload_optimizer(): + from shutil import copytree, disk_usage + offload_dir = self.optimizer.optimizer_swapper.swap_folder + offload_ckpt_dir = os.path.join(save_dir, tag, "offloaded_tensors") + _, _, free = disk_usage(save_dir) + logger.info( + f"Copying NVMe offload files from {offload_dir} to {offload_ckpt_dir}, {free / 1e9:,.2f} GB free on target filesystem..." + ) + copytree(offload_dir, + offload_ckpt_dir, + ignore=lambda _, dir_list: list(filter(lambda x: 'gradient' in x, dir_list)), + dirs_exist_ok=False) + _, _, free = disk_usage(save_dir) + logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem") + + if self._optimizer_has_ckpt_event_epilogue(): + self.optimizer.checkpoint_event_epilogue() + + # Save latest checkpoint tag + if not self.checkpoint_engine.is_decoupled(): + self.checkpoint_engine.commit(tag) + if save_latest and self.global_rank == 0: + with open(os.path.join(save_dir, 'latest'), 'w') as fd: + fd.write(tag) + + dist.barrier() + + return True + + def _commit_decoupled_checkpoint(self): + assert self.checkpoint_engine.is_decoupled(), \ + f'{self.checkpoint_engine} is not a Decoupled Checkpoint Engine' + + commit_info = self.checkpoint_engine.get_commit_info() + if commit_info is None: + return + + self.checkpoint_engine.commit(commit_info) + + if self.global_rank == 0 and commit_info.save_latest: + with open(os.path.join(commit_info.save_dir, 'latest'), 'w') as fd: + fd.write(commit_info.tag) + + dist.barrier() + + def _get_non_moe_state_dict(self, full_state_dict): + """ + Get the state dict of the non-moe layers + """ + for key in list(full_state_dict.keys()): + if 'expert' in key and 'moe.gate.wg.weight' not in key: + full_state_dict.pop(key) + + return full_state_dict + + def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): + save_path = self._get_ckpt_name(save_dir, tag) + + # A hack to save the checkpointing directory. Pipeline parallelism overrides + # module_state_dict() and uses this path to save the model. module_state_dict() + # then instead just returns None. + + # Using layer_#_export_# to save the model's expert state_dict + moe_layer_id = 0 + for n_module, module in self.module.named_modules(): + if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: + group_name = module.expert_group_name + num_local_experts = module.num_local_experts + expp_rank = groups._get_expert_parallel_rank(group_name) + exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) + # print(expp_rank, exp_dp_rank) + # if exp_dp_rank != 0: + if not self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank): + moe_layer_id += 1 + continue + + # get all moe parameters + moe_state_dict = {} + for n, p in module.state_dict().items(): + if 'expert' in n and 'moe.gate.wg.weight' not in n: + moe_state_dict[n_module + '.' + n] = p + moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' + # print(moe_state_dict.keys()) # until now, everything is fine. So the bug happens at next few lines + # Reorder the moe name rank, so that each checkpoint only has one expert + experts_state_dict = defaultdict(dict) + for key in list(moe_state_dict.keys()): + m = re.match(f".*{moe_str_prefix}([0-9]+).*", key) + + local_expert_id = None + if not m: + logger.warning(f'No expert found in key {key}.') + else: + local_expert_id = m.group(1) + + global_expert_id = expp_rank * \ + num_local_experts + int(local_expert_id) + expert_key = key.replace(f'{moe_str_prefix}{local_expert_id}', + f'{moe_str_prefix}{global_expert_id}') + # truncating extra tensor (shared) storage + truncated = moe_state_dict.pop(key).clone().detach() + experts_state_dict[str(global_expert_id)][expert_key] = truncated + + # let save the moe parameters + for global_expert_id, expert_state_dict in experts_state_dict.items(): + # save the moe parameters + moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu) + if self.random_ltd_enabled(): + expert_state_dict = remove_random_ltd_state_dict(expert_state_dict) + saveable_state_dict = clone_tensors_for_torch_save(expert_state_dict) + self.checkpoint_engine.save(saveable_state_dict, moe_save_path) + moe_layer_id += 1 + + self._curr_ckpt_path = os.path.join(save_dir, tag) + + largest_group_name = groups._get_max_expert_size_name() + expp_rank = groups._get_expert_parallel_rank(largest_group_name) + exp_dp_rank = groups._get_expert_data_parallel_rank(largest_group_name) + + # In the case of E + D parallelism, only the + # first expert parallel group should save the expert weights + # since each expert parallel group is a copy of the model's experts + if not self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank): + return + + # Save optimizer states. They are different across each exp parallel rank. + optimizer_state = { + 'optimizer': self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None + } + # TODO: why use BufferedWriter not the path + file_path = self._get_optimizer_ckpt_name(save_dir, tag, expp_rank) + saveable_state_dict = clone_tensors_for_torch_save(optimizer_state) + self.checkpoint_engine.save(saveable_state_dict, file_path) + + # Load flow uses below saved file for model parameters, RNG and more + if groups._get_data_parallel_rank() == 0: + # Get non-moe parameters + # Classes DeepSpeedEngine and PipelineEngine have different behavior for method module_state_dict. + # DeepSpeedEngine returns the state dict, where PipelineEngine saves the state dict and returns None. + # We need to get the state dict, therefore, call to DeepSpeedEngine (base class for PipelineEngine) + model_state_dict = self._get_non_moe_state_dict( + DeepSpeedEngine.module_state_dict(self, exclude_frozen_parameters=exclude_frozen_parameters)) + + # TODO: update num experts info,.. in checkpoint + state = { + 'module': + model_state_dict, + 'lr_scheduler': + self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, + 'data_sampler': + self.training_dataloader.data_sampler.state_dict() if + (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, + 'random_ltd': + self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, + 'sparse_tensor_module_names': + self.sparse_tensor_module_names, + 'skipped_steps': + self.skipped_steps, + 'global_steps': + self.global_steps, + 'global_samples': + self.global_samples, + 'dp_world_size': + self.dp_world_size, + 'mp_world_size': + self.mp_world_size, + 'num_experts': + self.num_experts + } + state.update(client_state) + logger.info(f'Saving model checkpoint: {save_path}') + saveable_state_dict = clone_tensors_for_torch_save(state) + self.checkpoint_engine.save(saveable_state_dict, save_path) + + def _create_checkpoint_file(self, save_dir, tag, zero_checkpoint): + name_function = (self._get_zero_ckpt_name if zero_checkpoint else self._get_ckpt_name) + try: + checkpoint_name = name_function(save_dir, tag) + path = os.path.dirname(checkpoint_name) + self.checkpoint_engine.makedirs(path, exist_ok=True) + except: + logger.error(f"Failed saving model checkpoint to {save_dir} with tag {tag}") + return False + + return True + + def _create_zero_checkpoint_files(self, save_dir, tag): + success = True + # zero checkpoint files are created sequentially + for rank in range(dist.get_world_size(self.optimizer.dp_process_group)): + if rank == self.global_rank: + success = self._create_checkpoint_file(save_dir, tag, True) + + return success + + def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): + + save_path = self._get_ckpt_name(save_dir, tag) + + zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() + + save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters + + # A hack to save the checkpointing directory. Pipeline parallelism overrides + # module_state_dict() and uses this path to save the model. module_state_dict() + # then instead just returns None. The module_state_dict() implementation in + # PipelineEngine expects the save path to be set in self._curr_ckpt_path. + self._curr_ckpt_path = os.path.join(save_dir, tag) + module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) + self._curr_ckpt_path = None + + state = dict(module=module, + buffer_names=self._get_buffer_names(), + optimizer=self.optimizer.state_dict() if self.optimizer and not zero_optimizer_state else None, + param_shapes=self._get_zero_param_shapes() if self.optimizer and zero_optimizer_state else None, + frozen_param_shapes=self._get_zero_frozen_param_attributes(self._get_param_shape_func) + if save_frozen_param else None, + shared_params=self._get_shared_params() if self.optimizer and zero_optimizer_state else None, + frozen_param_fragments=self._get_zero_frozen_param_attributes(self._get_param_fragment_func) + if save_frozen_param else None, + lr_scheduler=self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, + data_sampler=self.training_dataloader.data_sampler.state_dict() if + (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, + random_ltd=self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, + sparse_tensor_module_names=self.sparse_tensor_module_names, + skipped_steps=self.skipped_steps, + global_steps=self.global_steps, + global_samples=self.global_samples, + dp_world_size=self.seq_dp_world_size, + mp_world_size=self.mp_world_size, + ds_config=self.config, + ds_version=version) + state.update(client_state) + log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0]) + + if self.save_non_zero_checkpoint: + self.checkpoint_engine.save(state_dict=state, path=save_path) + + def _get_buffer_names(self): + buffer_names = [] + + # we save buffer names so that we could extract later the real buffers from the saved + # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they + # are intermixed with param placeholders + + # have to traverse the tree to be able to skip non-persistent buffers + def get_layer_named_buffers(module, prefix=""): + for name, buf in module.named_buffers(recurse=False): + if buf is not None and name not in module._non_persistent_buffers_set: + buffer_names.append(prefix + name) + + for name, child in module.named_children(): + if child is not None: + get_layer_named_buffers(child, prefix + name + ".") + + get_layer_named_buffers(self.module, prefix="") + + return buffer_names + + def _get_param_shape_func(self, param): + return param.ds_shape if hasattr(param, 'ds_id') else param.shape + + def _get_param_fragment_func(self, param): + return param.ds_tensor.detach().cpu() if hasattr(param, 'ds_id') else param.detach().cpu() + + def _get_zero_frozen_param_attributes(self, attr_func): + frozen_param_fragments = OrderedDict() + + for param in self.module.parameters(): + if param.requires_grad: + continue + if param not in self.param_names: + raise ValueError(f"failed to find frozen {param} in named params") + name = self.param_names[param] + frozen_param_fragments[name] = attr_func(param) + + return frozen_param_fragments + + def _get_zero_param_shapes(self): + """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the + optimizer. the names are exactly as in state_dict. The order is absolutely important, since + the saved data is just flattened data with no identifiers and requires reconstruction in the + same order it was saved. + We can't rely on self.module.named_parameters() to get the saved tensors, as some params + will be missing and others unsaved and then it'd be impossible to reconstruct state_dict + from the flattened weights. + optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions. + """ + param_group_shapes = [] + cnt = 0 + numel = 0 + + # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups - + # if we don't use it, we get parameters ordered incorrectly + if hasattr(self.optimizer, "round_robin_bit16_groups"): + bit16_groups = self.optimizer.round_robin_bit16_groups + elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"): + bit16_groups = self.optimizer.bf16_groups + else: + bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage( + ) == 2 else self.optimizer.fp16_groups + + for bit16_group in bit16_groups: + param_shapes = OrderedDict() + for param in bit16_group: + cnt += 1 + numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel() + shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape + if param not in self.param_names: + raise ValueError(f"failed to find optimizer param in named params") + name = self.param_names[param] + param_shapes[name] = shape + + # uncomment to debug zero_to_fp32.py problems + # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})") + param_group_shapes.append(param_shapes) + # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params") + + return param_group_shapes + + def _get_shared_params(self): + """ + Returns a dict of shared params, which can later be used to reconstruct the original state dict, + e.g. in `zero_to_fp32`. Each dict entry is a pair of param names, where the key is the name + of the variable that isn't stored and the value is the actual param holding data. + """ + shared_index = {} + shared_params_by_full_name = {} + + is_zero3_model = (self.zero_optimization_partition_weights() + and any(hasattr(param, "ds_id") for param in self.module.parameters())) + + def get_layer_state_dict(module, prefix=""): + # handle params + for name, param in module.named_parameters(recurse=False): + if param is None or (is_zero3_model and not hasattr(param, "ds_id")): + continue + key = prefix + name + + # When weights are manged by stage 3, we can't rely on param.data_ptr() as it will be reused + # as weights get gathered and reduced, but param.ds_id is unique across all zero weights + # (and shared params will have the same param.ds_id) + param_id = param.ds_id if is_zero3_model else param.data_ptr() + + if param_id in shared_index: + # shared weights + #print(f"`{key}` is shared with `{shared_index[param_id]}`") + shared_params_by_full_name[key] = shared_index[param_id] + else: + shared_index[param_id] = key + + for name, child in module.named_children(): + if child is not None: + get_layer_state_dict(child, prefix + name + ".") + + if dist.get_rank() == 0: + get_layer_state_dict(self.module, prefix="") + + return shared_params_by_full_name + + def _copy_recovery_script(self, save_path): + base_dir = os.path.dirname(os.path.dirname(__file__)) + script = "zero_to_fp32.py" + src = os.path.join(base_dir, "utils", script) + dst = os.path.join(save_path, script) + #logger.info(f"creating recovery script {dst}") + copyfile(src, dst) + self._change_recovery_script_permissions(dst) + + def _change_recovery_script_permissions(self, dst): + # make executable (safeguard for file shares - Azure as example) + try: + os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC) + except (FileNotFoundError, PermissionError) as e: + #this message is used in unit test TestZeRONonDistributed + logger.info( + f'Warning: Could not change permissions for {dst} due to error: {e}. Continuing without changing permissions.' + ) + + def _save_zero_checkpoint(self, save_path, tag): + zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag) + zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version) + self.checkpoint_engine.save(zero_sd, zero_checkpoint_name) + + if self.global_rank == 0: + self._copy_recovery_script(save_path) + ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero' + #logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}') + + def _replace_module_consolidated_state_dict(self): + """ + Get a full non-partitioned state_dict with fp16 weights on cpu. + Important: this function must be called on all ranks and not just rank 0. + This is similar to nn.Module.state_dict (modelled after _save_to_state_dict) + This method is used for tensor parallel training. + + Returns: + OrderedDict: The consolidated state dictionary if the current process rank is 0, otherwise None. + """ + #TODO: If we use both Zero3 and tensor parallel simultaneously + # we need to consolidate the gather mechanisms of both. + state_dict = OrderedDict() if dist.get_rank() == 0 else None + + def get_layer_state_dict(module, prefix=""): + with GatherReplacedLayerParams(list(module.parameters(recurse=False)), module, enabled=True): + for name, param in module.named_parameters(recurse=False): + if param is None: + continue + key = prefix + name + if (dist.get_rank() == 0): + state_dict[key] = param.detach().cpu() + # print(key,module, param.detach().cpu().shape) + + for name, child in module.named_children(): + if child is not None: + get_layer_state_dict(child, prefix + name + ".") + + get_layer_state_dict(self.module, prefix="") + + # ensure that all GPU communication tasks are completed before the process exits + get_accelerator().synchronize() + return state_dict + + def _consolidated_16bit_state_dict(self, exclude_frozen_parameters=False): + """ + Consolidate the 16-bit state dictionary. + """ + if self.zero_optimization_stage() == ZeroStageEnum.weights: + return self._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters) + elif self.autotp_size() > 1: + return self._replace_module_consolidated_state_dict() + + raise ValueError("consolidated_16bit_state_dict is only applicable to cases where weights are partitioned, " + "including Zero Stage 3 and tensor parallelism.") + + def _zero3_consolidated_16bit_state_dict(self, exclude_frozen_parameters=False): + """ + Get a full non-partitioned state_dict with fp16 weights on cpu. + Important: this function must be called on all ranks and not just rank 0. + This is similar to nn.Module.state_dict (modelled after _save_to_state_dict), but: + 1. consolidates the weights from different partitions on gpu0 + 2. works on one layer at a time to require as little gpu0 memory as possible, by + moving the already consolidated weights to cpu + 3. takes care to keep the shared params shared when gradually copying the params to cpu + Returns: + a consolidated fp16 ``state_dict`` on cpu on rank 0, ``None`` on other ranks + """ + if not self.zero_optimization_partition_weights(): + raise ValueError("this function requires ZeRO-3 mode") + + state_dict = OrderedDict() if dist.get_rank() == 0 else None + shared_params = {} + + def get_layer_state_dict(module, prefix=""): + # gather one layer at a time to be memory-efficient + # must use modifier_rank=0 to release GPU memory after each layer gathered + #see_memory_usage("before GatheredParameters", force=True) + with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): + if dist.get_rank() == 0: + # handle params + for name, param in module.named_parameters(recurse=False): + if param is None or (exclude_frozen_parameters and not param.requires_grad): + continue + key = prefix + name + # can't rely on param.data_ptr() as it will be reused as weights gets + # gathered and reduced, but param.ds_id is unique across all zero weights + # (and shared params will have the same param.ds_id) + if param.ds_id in shared_params: + # shared weights + #print(f"`{key}` is shared with `{shared_params[param.ds_id]}`") + state_dict[key] = state_dict[shared_params[param.ds_id]] + else: + state_dict[key] = param.detach().cpu() + shared_params[param.ds_id] = key + #print(f"param {param.ds_id} {param.shape} {key} ") + + # now buffers - not sure if need to take care of potentially shared weights here + for name, buf in module.named_buffers(recurse=False): + if (buf is not None and name not in module._non_persistent_buffers_set): + state_dict[prefix + name] = buf.detach().cpu() + #see_memory_usage("after GatheredParameters", force=True) + + for name, child in module.named_children(): + if child is not None: + get_layer_state_dict(child, prefix + name + ".") + + # Prepare for checkpoint save by ensuring all parameters are partitioned + if self._optimizer_has_ckpt_event_prologue(): + self.optimizer.checkpoint_event_prologue() + + see_memory_usage("before get_layer_state_dict", force=False) + get_layer_state_dict(self.module, prefix="") + see_memory_usage("after get_layer_state_dict", force=False) + + if self._optimizer_has_ckpt_event_epilogue(): + self.optimizer.checkpoint_event_epilogue() + + return state_dict + + def save_fp16_model(self, save_dir, save_filename="pytorch_model.bin"): + """has been renamed to save_16bit_model, keeping this around for backwards + compatibility""" + return self.save_16bit_model(save_dir, save_filename) + + def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False): + """ + Save 16bit model weights + + This method saves the 16bit model weights at the desired destination. + + Arguments: + save_dir: Required. Directory for saving the model + save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin`` + exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. + + Returns: + ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if + stage3_gather_16bit_weights_on_model_save is ``False``. + + Important: all processes must call this method and not just the process with rank 0. It is + because the processes need to work in sync to gather the weights. This method will hang + waiting to synchronize with other processes if it's called just for the process with rank 0. + + """ + + path = os.path.join(save_dir, save_filename) + + if self.zero_optimization_partition_weights(): + if self.zero_gather_16bit_weights_on_model_save(): + # consolidation is expensive in time and memory and therefore isn't a default + state_dict = self._zero3_consolidated_16bit_state_dict( + exclude_frozen_parameters=exclude_frozen_parameters) + else: + # the model will be bogus if not consolidated so don't confuse the user by saving it + logger.info( + f"Did not save the model {path} because stage3_gather_16bit_weights_on_model_save is False") + return False + else: + state_dict = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) + + tag = f"global_step{self.global_steps}" + tag = str(tag) + commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=False) + self.checkpoint_engine.create(commit_info) + + 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={}, schedule=None) -> None: + """Compile the module using the specified backend and kwargs. + If a compiler_fn is set, it will be used instead of torch.compile(). + """ + # Avoid graph breaks + deepspeed.utils.nvtx.enable_nvtx = False + + if not is_compile_supported(): + raise RuntimeError("compile is not supported in your version of PyTorch.") + + if self.is_compiled: + return + + if 'backend' in compile_kwargs: + logger.warning("The `backend` in `compile_kwargs` will be overridden. Use the `backend` argument instead.") + + logger.info(f"Compiling deepcompile={self.is_deepcompile_enabled()} backend={backend}") + + enable_deepcompile = self.is_deepcompile_enabled() + if enable_deepcompile and self.zero_optimization_stage() != ZeroStageEnum.optimizer_states \ + and self.zero_optimization_stage() != ZeroStageEnum.weights \ + and self.zero_optimization_stage() != ZeroStageEnum.gradients: + logger.info( + f"Currently DeepCompile supports ZeRO stage 1, 2, or 3 only, but ZeRO stage is set to {self.zero_optimization_stage()}. Falling back to the torch compiler." + ) + enable_deepcompile = False + + if enable_deepcompile: + + if schedule is not None: + + def passes_name_to_fn(passes): + for p in passes: + assert callable(p) or p in opt_passes, f"Unknown pass {p}" + return [p if callable(p) else opt_passes[p] for p in passes] + + schedule = [(step, passes_name_to_fn(passes)) for step, passes in schedule] + + assert backend in ['inductor', 'eager'], f"Backend {backend} is not supported for DeepCompile." + + compile_config = self._config.compile_config + if (("zero_optimization" in self.config and "offload_optimizer" in self.config["zero_optimization"] + and "offload_param" in self.config["zero_optimization"]) + and self._config.zero_config.offload_param.device == "cpu" + and self._config.zero_config.offload_optimizer.device == "cpu"): + compile_config.offload_parameters = True + if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states: + backend = init_z1(self, backend, compile_config, compile_kwargs, schedule) + elif self.zero_optimization_stage() == ZeroStageEnum.gradients: + backend = init_z1(self, backend, compile_config, compile_kwargs, schedule, use_z2=True) + elif self.zero_optimization_stage() == ZeroStageEnum.weights: + backend = init_z3(self, backend, compile_config, compile_kwargs, schedule) + + # create new dict to avoid modifying original dict + self.module.compile(**{**compile_kwargs, 'backend': backend}) + + self._is_compiled = True + + def get_compile_time(self): + from deepspeed.compile.backend import opt_pass_times + return opt_pass_times + + def register_compile_pass(self, pass_name: str, pass_fn: Callable) -> None: + register_compile_pass(pass_name, pass_fn) + + def is_deepcompile_enabled(self): + return self._config.compile_config.deepcompile + + @property + def is_compiled(self) -> bool: + return self._is_compiled + + def offload_states(self, + include: Container[OffloadStateTypeEnum] = None, + device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, + pin_memory: bool = True, + non_blocking: bool = False) -> None: + """Offload the engine's states to the specified device. + + Arguments: + include: Optional. The set of states to offload. If not provided, all states are offloaded. + device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported. + pin_memory: Optional. Whether to pin the memory of the offloaded states. + non_blocking: Optional. Whether to offload the states asynchronously. + """ + assert self.zero_optimization_stage( + ) == ZeroStageEnum.weights, "Moving buffers across devices is supported only for ZeRO stage 3." + + opt_offload_config = self.zero_offload_optimizer() + assert opt_offload_config is None or opt_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded optimizer states." + param_offload_config = self.zero_offload_param() + assert param_offload_config is None or param_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded parameters." + + assert not isinstance( + self.optimizer, + DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." + + if device == OffloadDeviceEnum.none: + logger.warning("No device specified for offloading states.") + return + + if device == OffloadDeviceEnum.nvme: + raise ValueError("NVMe offload is not supported for offloading states.") + + self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking) + + def reload_states(self, non_blocking: bool = False) -> None: + """Reload the engine states to the original device. + + Arguments: + non_blocking: Optional. Whether to offload the states asynchronously. + """ + assert self.zero_optimization_stage( + ) == ZeroStageEnum.weights, "Moving buffers back is supported only for ZeRO stage 3." + + assert not isinstance( + self.optimizer, + DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." + + self.optimizer.reload_states(non_blocking=non_blocking) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87a13af244e12fd50badbe51120fa582c4cb4438 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e31230449e156aa69cd4b49d307a392c34f51d16 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a3385ed6644c5c3dca42ac1c2218f62f1b36494 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d05f83a5b6907d9c0ef4d8cf2ad928fdd04993ec Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/fused_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/fused_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..706fb6ec78e38f9750f62f0addc0b104b1af91be --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/fused_optimizer.py @@ -0,0 +1,526 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from FP16_Optimizer in NVIDIA/apex +""" + +import torch +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer +from deepspeed.runtime.utils import get_global_norm, get_flattened_grad_norm, CheckOverflow, get_weight_norm, get_norm_with_moe_layers, is_model_parallel_parameter +from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE +from deepspeed.utils import logger, log_dist +from deepspeed.utils.torch import required_torch_version +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD +from deepspeed.accelerator import get_accelerator +from deepspeed.moe.utils import is_moe_param_group +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank + +OVERFLOW_CHECK_TIMER = 'overflow_check' +COMPUTE_NORM_TIMER = 'compute_norm' +UNSCALE_AND_CLIP_TIMER = 'unscale_and_clip' +BASIC_STEP_TIMER = 'basic_step' +UPDATE_FP16_TIMER = 'update_fp16' + +OVERFLOW_TIMERS = [COMPUTE_NORM_TIMER, OVERFLOW_CHECK_TIMER] +STEP_TIMERS = OVERFLOW_TIMERS + [UNSCALE_AND_CLIP_TIMER, BASIC_STEP_TIMER, UPDATE_FP16_TIMER] + + +class FP16_Optimizer(DeepSpeedOptimizer): + """ + FP16 Optimizer for training fp16 models. Handles loss scaling. + + For usage example please see, TODO: DeepSpeed V2 Tutorial + """ + + def __init__(self, + init_optimizer, + deepspeed=None, + static_loss_scale=1.0, + dynamic_loss_scale=False, + initial_dynamic_scale=2**32, + dynamic_loss_args=None, + verbose=True, + mpu=None, + clip_grad=0.0, + fused_adam_legacy=False, + has_moe_layers=False, + timers=None): + + self.fused_adam_legacy = fused_adam_legacy + self.timers = timers + self.deepspeed = deepspeed + self.has_moe_layers = has_moe_layers + self.using_pipeline = getattr(self.deepspeed, 'pipeline_parallelism', False) + if not get_accelerator().is_available(): + raise SystemError("Cannot use fp16 without accelerator.") + self.optimizer = init_optimizer + + # param flattened by groups + self.fp16_groups = [] + self.fp16_groups_flat = [] + self.fp32_groups_flat = [] + + self.flatten_grad_norm_mask_list = [] + self.has_executed_step = False + self._global_grad_norm = 0. + + # loop to deal with groups + for i, param_group in enumerate(self.optimizer.param_groups): + # push this group to list before modify + self.fp16_groups.append(param_group['params']) + # init fp16 weight buffer, flattened + self.fp16_groups_flat.append(_flatten_dense_tensors([p.clone().detach() for p in self.fp16_groups[i]])) + # set model fp16 weight to slices of flattened buffer + updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) + for p, q in zip(self.fp16_groups[i], updated_params): + p.data = q.data + # init master weight, flattened + self.fp32_groups_flat.append(self.fp16_groups_flat[i].clone().float().detach()) + # modify optimizer of have flat master weight + self.fp32_groups_flat[i].requires_grad = True # keep this in case internal optimizer uses it + param_group['params'] = [self.fp32_groups_flat[i]] + + # we may have a way of fusing dynamic scale. Do not support for now + if dynamic_loss_scale: + self.dynamic_loss_scale = True + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = 2 + + if dynamic_loss_args is None: + self.cur_scale = initial_dynamic_scale + self.scale_window = 1000 + self.min_loss_scale = 1 + else: + self.cur_scale = dynamic_loss_args[INITIAL_LOSS_SCALE] + self.scale_window = dynamic_loss_args[SCALE_WINDOW] + self.min_loss_scale = dynamic_loss_args[MIN_LOSS_SCALE] + else: + self.dynamic_loss_scale = False + self.cur_iter = 0 + self.cur_scale = static_loss_scale + self.verbose = verbose + + self.custom_loss_scaler = False + self.external_loss_scale = None + + self.clip_grad = clip_grad + self.norm_type = 2 + + if required_torch_version(max_version=0.4): + self.clip_grad_norm = torch.nn.utils.clip_grad_norm + else: + self.clip_grad_norm = torch.nn.utils.clip_grad_norm_ + + #model parallel object + self.mpu = mpu + + self.overflow = False + self.overflow_checker = CheckOverflow(self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed) + self.initialize_optimizer_states() + + def initialize_optimizer_states(self): + for i, group in enumerate(self.fp16_groups): + self.fp32_groups_flat[i].grad = torch.zeros(self.fp32_groups_flat[i].size(), + device=self.fp32_groups_flat[i].device) + + self.optimizer.step() + + for i, group in enumerate(self.fp16_groups): + self.fp32_groups_flat[i].grad = None + + return + + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if set_to_none: + p.grad = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def step_fused_adam(self, closure=None): + """ + Not supporting closure. + """ + + # First compute norm for all group so we know if there is overflow + grads_groups_flat = [] + norm_groups = [] + for i, group in enumerate(self.fp16_groups): + grads_groups_flat.append( + _flatten_dense_tensors([ + torch.zeros(p.size(), dtype=p.dtype, device=p.device) if p.grad is None else p.grad for p in group + ])) + norm_groups.append(get_weight_norm(grads_groups_flat[i], mpu=self.mpu)) + + self.overflow = self.overflow_checker.check_using_norm(norm_groups) + prev_scale = self.cur_scale + self._update_scale(self.overflow) + + if self.overflow: + if self.verbose: + logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " + "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) + return self.overflow + + scaled_grad_norm = get_global_norm(norm_list=norm_groups) + + combined_scale = self.unscale_and_clip_grads(grads_groups_flat, scaled_grad_norm, apply_scale=False) + + # Stash unscaled gradient norm + self._global_grad_norm = scaled_grad_norm / self.cur_scale + + # norm is in fact norm*cur_scale + self.optimizer.step(grads=[[g] for g in grads_groups_flat], + output_params=[[p] for p in self.fp16_groups_flat], + scale=combined_scale, + grad_norms=norm_groups) + # TODO: we probably don't need this? just to be safe + for i in range(len(norm_groups)): + updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) + for p, q in zip(self.fp16_groups[i], updated_params): + p.data = q.data + return self.overflow + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + def _require_avoid_recompute_norm(self, p, tensor_model_parallel_rank): + # for filtering replicated tensors from tensor + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + return True + if (tensor_model_parallel_rank > 0) and not is_model_parallel_parameter(p): + return True + + def _get_norm_mask_idx(self, group): + """The function preserves the parallel information for norm + from unflattened gradients. + + Args: + group (Iterable[Tensor] ): params group + + Returns: + torch.Tensor: A 2D tensor containing index ranges for each group, + where each row represents a [start index, end index]. + """ + group_mask_idx_list = [] + grad_flat_st_idx = 0 + grad_flat_en_idx = 0 + + for p in group: + grad_flat_en_idx = grad_flat_st_idx + p.numel() + if p.grad is not None and self._require_avoid_recompute_norm(p, bwc_tensor_model_parallel_rank(self.mpu)): + # merge range + if len(group_mask_idx_list) > 0 and grad_flat_st_idx == group_mask_idx_list[-1][-1]: + group_mask_idx_list[-1][-1] = grad_flat_en_idx + else: + group_mask_idx_list.append([grad_flat_st_idx, grad_flat_en_idx]) + grad_flat_st_idx = grad_flat_en_idx + + return torch.tensor(group_mask_idx_list, device=get_accelerator().current_device_name()) + + def step(self, closure=None): + """ + Not supporting closure. + """ + + if self.fused_adam_legacy: + return self.step_fused_adam() + + # First determine if there is overflow. + if self.timers: + self.timers(OVERFLOW_CHECK_TIMER).start() + fp16_params = [] + for i, group in enumerate(self.fp16_groups): + fp16_params.extend([p for p in group if p.grad is not None]) + self.overflow = self.overflow_checker.has_overflow(fp16_params) + if self.timers: + self.timers(OVERFLOW_CHECK_TIMER).stop() + prev_scale = self.cur_scale + self._update_scale(self.overflow) + if self.overflow: + if self.verbose: + log_dist( + "Overflow detected. Skipping step. Attempted loss " + f"scale: {prev_scale}, reducing to {self.cur_scale}", + ranks=[0]) + # Clear gradients + for i, group in enumerate(self.fp16_groups): + for p in group: + p.grad = None + + if self.timers: + self.timers.log(OVERFLOW_TIMERS) + return self.overflow + + grads_groups_flat = [] + non_experts_grads_for_norm = [] + expert_grads_for_norm = {} + assert len(self.fp16_groups) == len(self.optimizer.param_groups) + + for i, group in enumerate(self.fp16_groups): + data_type = self.fp32_groups_flat[i].dtype + + grads_groups_flat.append( + _flatten_dense_tensors([ + torch.zeros(p.size(), dtype=data_type, device=p.device) if p.grad is None else p.grad.to(data_type) + for p in group + ])) + + self.fp32_groups_flat[i].grad = grads_groups_flat[i] + param_group = self.optimizer.param_groups[i] + + # split expert and non_expert grads for norm + if self.has_moe_layers and is_moe_param_group(param_group): + if param_group['name'] not in expert_grads_for_norm: + expert_grads_for_norm[param_group['name']] = [] + + expert_grads_for_norm[param_group['name']].append(self.fp32_groups_flat[i]) + else: + # retrieves the required mask for calculating the norm of flat_grad + # perform this collect operation only once + if not self.has_executed_step: + cur_flat_grad_norm_mask = self._get_norm_mask_idx(group) + self.flatten_grad_norm_mask_list.append(cur_flat_grad_norm_mask) + + non_experts_grads_for_norm.append(self.fp32_groups_flat[i]) + + for p in group: + p.grad = None + + if self.timers: + self.timers(COMPUTE_NORM_TIMER).start() + + all_groups_norm = get_flattened_grad_norm(non_experts_grads_for_norm, + mpu=self.mpu, + grad_norm_mask=self.flatten_grad_norm_mask_list) + + if self.has_moe_layers: + all_groups_norm = get_norm_with_moe_layers(all_groups_norm, + mpu=self.mpu, + expert_tensors=expert_grads_for_norm, + norm_type=self.norm_type) + + scaled_global_grad_norm = get_global_norm(norm_list=[all_groups_norm]) + if self.timers: + self.timers(COMPUTE_NORM_TIMER).stop() + + # Stash unscaled gradient norm + self._global_grad_norm = scaled_global_grad_norm / self.cur_scale + + if self.timers: + self.timers(UNSCALE_AND_CLIP_TIMER).start() + self.unscale_and_clip_grads(grads_groups_flat, scaled_global_grad_norm) + if self.timers: + self.timers(UNSCALE_AND_CLIP_TIMER).stop() + + if self.timers: + self.timers(BASIC_STEP_TIMER).start() + self.optimizer.step() + if self.timers: + self.timers(BASIC_STEP_TIMER).stop() + + #get rid of the fp32 gradients. Not needed anymore + for group in self.fp32_groups_flat: + group.grad = None + + if self.timers: + self.timers(UPDATE_FP16_TIMER).start() + + for i in range(len(self.fp16_groups)): + updated_params = _unflatten_dense_tensors(self.fp32_groups_flat[i], self.fp16_groups[i]) + for p, q in zip(self.fp16_groups[i], updated_params): + p.data.copy_(q.data) + self.has_executed_step = True + if self.timers: + self.timers(UPDATE_FP16_TIMER).stop() + + if self.timers: + self.timers.log(STEP_TIMERS) + + return self.overflow + + def unscale_and_clip_grads(self, grad_groups_flat, total_norm, apply_scale=True): + # compute combined scale factor for this group + combined_scale = self.cur_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad + if clip > 1: + combined_scale = clip * self.cur_scale + + if apply_scale: + for grad in grad_groups_flat: + grad.data.mul_(1. / combined_scale) + + return combined_scale + + def backward(self, loss, create_graph=False, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + else: + scaled_loss = (loss.float()) * self.cur_scale + scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph) + + def _update_scale(self, skip): + if self.dynamic_loss_scale: + prev_scale = self.cur_scale + if skip: + self.cur_scale = max(self.cur_scale / self.scale_factor, self.min_loss_scale) + self.last_overflow_iter = self.cur_iter + if self.verbose: + logger.info(f"\nGrad overflow on iteration {self.cur_iter}") + logger.info(f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + # Ensure self.scale_window updates since last overflow + stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 + if (stable_interval > 0) and (stable_interval % self.scale_window == 0): + self.cur_scale *= self.scale_factor + if self.verbose: + logger.info(f"No Grad overflow for {self.scale_window} iterations") + logger.info(f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + if skip: + logger.info("Grad overflow on iteration: %s", self.cur_iter) + logger.info("Using static loss scale of: %s", self.cur_scale) + self.cur_iter += 1 + return + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + + param_groups = property(_get_param_groups, _set_param_groups) + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + state_dict = {} + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['cur_scale'] = self.cur_scale + state_dict['cur_iter'] = self.cur_iter + if state_dict['dynamic_loss_scale']: + state_dict['last_overflow_iter'] = self.last_overflow_iter + state_dict['scale_factor'] = self.scale_factor + state_dict['scale_window'] = self.scale_window + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + state_dict['fp32_groups_flat'] = self.fp32_groups_flat + state_dict[CLIP_GRAD] = self.clip_grad + return state_dict + + # Refresh fp32 master params from fp16 copies + def refresh_fp32_params(self): + for current, saved in zip(self.fp32_groups_flat, self.fp16_groups_flat): + current.data.copy_(saved.data) + + def load_state_dict(self, state_dict, load_optimizer_states=True): + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + # I think it should actually be ok to reload the optimizer before the model. + self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] + self.cur_scale = state_dict['cur_scale'] + self.cur_iter = state_dict['cur_iter'] + if state_dict['dynamic_loss_scale']: + self.last_overflow_iter = state_dict['last_overflow_iter'] + self.scale_factor = state_dict['scale_factor'] + self.scale_window = state_dict['scale_window'] + if load_optimizer_states: + self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) + self.clip_grad = state_dict[CLIP_GRAD] + # At this point, the optimizer's references to the model's fp32 parameters are up to date. + # The optimizer's hyperparameters and internal buffers are also up to date. + # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still + # out of date. There are two options. + # 1: Refresh the master params from the model's fp16 params. + # This requires less storage but incurs precision loss. + # 2: Save and restore the fp32 master copies separately. + # We choose option 2. + # + # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device + # of their associated parameters, because it's possible those buffers might not exist yet in + # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been + # constructed in the same way as the one whose state_dict we are loading, the same master params + # are guaranteed to exist, so we can just copy_() from the saved master params. + for current, saved in zip(self.fp32_groups_flat, state_dict['fp32_groups_flat']): + current.data.copy_(saved.data) + + def __repr__(self): + return repr(self.optimizer) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py new file mode 100644 index 0000000000000000000000000000000000000000..3f5ed019fc78fdfe0a7a5e610fdce6f3b8498e60 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py @@ -0,0 +1,272 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright (c) 2019, 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. +Taken and modified for DeepSpeed from: + https://github.com/NVIDIA/Megatron-LM/blob/master/fp16/loss_scaler.py +Commit: 93ab4bea59dc5cbf97c079d313741866af4deac9 +""" + +import torch +from deepspeed.runtime.config_utils import DeepSpeedConfigObject +from deepspeed import comm as dist +from deepspeed.utils import logger + +INITIAL_LOSS_SCALE = 'init_scale' +SCALE_WINDOW = 'scale_window' +DELAYED_SHIFT = 'delayed_shift' +CONSECUTIVE_HYSTERESIS = 'consecutive_hysteresis' +MIN_LOSS_SCALE = 'min_scale' + + +# item() is a recent addition, so this helps with backward compatibility. +def to_python_float(t): + if hasattr(t, 'item'): + return t.item() + return t[0] + + +class LossScalerBase(DeepSpeedConfigObject): + """LossScalarBase + Base class for a loss scaler + """ + + def __init__(self, cur_scale): + super(LossScalerBase, self).__init__() + self.cur_scale = cur_scale + self.dynamic = False + + @property + def loss_scale(self): + return self.cur_scale + + def scale_gradient(self, module, grad_in, grad_out): + return tuple(self.loss_scale * g for g in grad_in) + + def update_scale(self, overflow): + pass + + def backward(self, loss, retain_graph=False): + scaled_loss = loss * self.loss_scale + scaled_loss.backward(retain_graph=retain_graph) + # print(f'LossScalerBackward: {scaled_loss=}') + + +class LossScaler(LossScalerBase): + """ + Class that manages a static loss scale. This class is intended to interact with + :class:`FP16_Optimizer`, and should not be directly manipulated by the user. + + Use of :class:`LossScaler` is enabled via the ``static_loss_scale`` argument to + :class:`FP16_Optimizer`'s constructor. + + Args: + scale (float, optional, default=1.0): The loss scale. + """ + + def __init__(self, scale=1): + super(LossScaler, self).__init__(scale) + + # `params` is a list / generator of torch.Variable + def has_overflow(self, params): + return False + + # `x` is a torch.Tensor + def _has_inf_or_nan(x): + return False + + +class DynamicLossScaler(LossScalerBase): + """ + Class that manages dynamic loss scaling. It is recommended to use :class:`DynamicLossScaler` + indirectly, by supplying ``dynamic_loss_scale=True`` to the constructor of + :class:`FP16_Optimizer`. However, it's important to understand how :class:`DynamicLossScaler` + operates, because the default options can be changed using the + the ``dynamic_loss_args`` argument to :class:`FP16_Optimizer`'s constructor. + + Loss scaling is designed to combat the problem of underflowing gradients encountered at long + times when training fp16 networks. Dynamic loss scaling begins by attempting a very high loss + scale. Ironically, this may result in OVERflowing gradients. If overflowing gradients are + encountered, :class:`DynamicLossScaler` informs :class:`FP16_Optimizer` that an overflow has + occurred. + :class:`FP16_Optimizer` then skips the update step for this particular iteration/minibatch, + and :class:`DynamicLossScaler` adjusts the loss scale to a lower value. + If a certain number of iterations occur without overflowing gradients detected, + :class:`DynamicLossScaler` increases the loss scale once more. + In this way :class:`DynamicLossScaler` attempts to "ride the edge" of + always using the highest loss scale possible without incurring overflow. + + Args: + init_scale (float, optional, default=2**32): Initial loss scale attempted by :class:`DynamicLossScaler.` + scale_factor (float, optional, default=2.0): Factor used when adjusting the loss scale. If an overflow is encountered, the loss scale is readjusted to loss scale/``scale_factor``. If ``scale_window`` consecutive iterations take place without an overflow, the loss scale is readjusted to loss_scale*``scale_factor``. + scale_window (int, optional, default=1000): Number of consecutive iterations without an overflow to wait before increasing the loss scale. + consecutive_hysteresis (bool, optional, default=False): Whether to refill hysteresis if we reach an iteration that doesn't overflow + """ + + def __init__(self, + init_scale, + scale_window, + min_scale, + delayed_shift, + consecutive_hysteresis, + raise_error_at_min_scale=True, + dtype=torch.half): + super(DynamicLossScaler, self).__init__(init_scale) + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = 2.0 + self.scale_window = scale_window + self.min_scale = min_scale + self.delayed_shift = delayed_shift + self.cur_hysteresis = delayed_shift + self.consecutive_hysteresis = consecutive_hysteresis + self.raise_error_at_min_scale = raise_error_at_min_scale + self.dynamic = True + self.dtype = dtype + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params): + for p in params: + if p.grad is not None and self._has_inf_or_nan(p.grad.data): + return True + + return False + + # `x` is a torch.Tensor + def _has_inf_or_nan(x): + try: + # if x is half, the .float() incurs an additional deep copy, but it's necessary if + # Pytorch's .sum() creates a one-element tensor of the same type as x + # (which is true for some recent version of pytorch). + cpu_sum = float(x.float().sum()) + # More efficient version that can be used if .sum() returns a Python scalar + # cpu_sum = float(x.sum()) + except RuntimeError as instance: + # We want to check if inst is actually an overflow exception. + # RuntimeError could come from a different error. + # If so, we still want the exception to propagate. + if "value cannot be converted" not in instance.args[0]: + raise + return True + else: + if cpu_sum in [float('inf'), -float('inf')] or cpu_sum != cpu_sum: + return True + return False + + # `overflow` is boolean indicating whether the gradient overflowed + def update_scale(self, overflow): + if overflow: + # self.cur_scale /= self.scale_factor + if self.delayed_shift == 1 or self.cur_hysteresis == 1: + if (self.cur_scale == self.min_scale) and self.raise_error_at_min_scale: + raise Exception( + "Current loss scale already at minimum - cannot decrease scale anymore. Exiting run.") + else: + next_scale = max(self.cur_scale / self.scale_factor, self.min_scale) + if dist.get_rank() == 0: + overflow_msg = f"[deepspeed] OVERFLOW! Rank {dist.get_rank()} Skipping step." + if self.dtype == torch.half: + overflow_msg += f" Attempted loss scale: {int(self.cur_scale)}, reducing to {int(next_scale)}" + logger.info(overflow_msg) + self.cur_scale = next_scale + else: + if dist.get_rank() == 0: + overflow_msg = f"[deepspeed] OVERFLOW! Rank {dist.get_rank()} Skipping step." + if self.dtype == torch.half: + overflow_msg += f" Attempted loss scale: {int(self.cur_scale)}, but hysteresis is {self.cur_hysteresis}. Reducing hysteresis to {self.cur_hysteresis-1}" + logger.info(overflow_msg) + self.cur_hysteresis -= 1 + self.last_overflow_iter = self.cur_iter + else: + if self.consecutive_hysteresis: + if dist.get_rank() == 0: + hysteresis_msg = f"Consecutive hysteresis is enabled. Restoring hysteresis to {self.delayed_shift}" + logger.info(hysteresis_msg) + self.cur_hysteresis = self.delayed_shift + + stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 + if (stable_interval > 0) and (stable_interval % self.scale_window == 0): + if not self.consecutive_hysteresis: + self.cur_hysteresis = self.delayed_shift + self.cur_scale *= self.scale_factor + self.cur_iter += 1 + + +# Although loss scaling is only defined for fp16, yet for backwards compatibility +# we still create a scaler for other dtypes (fp32, bf16) which does not perform any scaling. +def CreateLossScaler(dtype, static_loss_scale, dynamic_scaling, dynamic_loss_args): + if dtype == torch.half and dynamic_scaling: + assert dynamic_loss_args is not None, f"Dynamic loss scaling parameters must be defined." + return DynamicLossScaler(dtype=dtype, **dynamic_loss_args) + + loss_scale_value = static_loss_scale if dtype == torch.half else 1.0 + return LossScaler(scale=loss_scale_value) + + +############################################################## +# Example usage below here -- assuming it's in a separate file +############################################################## +""" +TO-DO separate out into an example. +if __name__ == "__main__": + import torch + from torch.autograd import Variable + from dynamic_loss_scaler import DynamicLossScaler + + # N is batch size; D_in is input dimension; + # H is hidden dimension; D_out is output dimension. + N, D_in, H, D_out = 64, 1000, 100, 10 + + # Create random Tensors to hold inputs and outputs, and wrap them in Variables. + x = Variable(torch.randn(N, D_in), requires_grad=False) + y = Variable(torch.randn(N, D_out), requires_grad=False) + + w1 = Variable(torch.randn(D_in, H), requires_grad=True) + w2 = Variable(torch.randn(H, D_out), requires_grad=True) + parameters = [w1, w2] + + learning_rate = 1e-6 + optimizer = torch.optim.SGD(parameters, lr=learning_rate) + loss_scaler = DynamicLossScaler() + + for t in range(500): + y_pred = x.mm(w1).clamp(min=0).mm(w2) + loss = (y_pred - y).pow(2).sum() * loss_scaler.loss_scale + print('Iter {} loss scale: {}'.format(t, loss_scaler.loss_scale)) + print('Iter {} scaled loss: {}'.format(t, loss.data[0])) + print('Iter {} unscaled loss: {}'.format(t, loss.data[0] / loss_scaler.loss_scale)) + + # Run backprop + optimizer.zero_grad() + loss.backward() + + # Check for overflow + has_overflow = DynamicLossScaler.has_overflow(parameters) + + # If no overflow, unscale grad and update as usual + if not has_overflow: + for param in parameters: + param.grad.data.mul_(1. / loss_scaler.loss_scale) + optimizer.step() + # Otherwise, don't do anything -- ie, skip iteration + else: + print('fp16 dynamic loss scale overflow!') + + # Update loss scale for next iteration + loss_scaler.update_scale(has_overflow) + +""" diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6f7e9784ce60f6c1b4a9134b73b4e415337641 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .adam import OnebitAdam +from .lamb import OnebitLamb +from .zoadam import ZeroOneAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1d97208a48f40cb56e8695e7f7abc747a982489 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6ffe01b1b105ac2db2116434034fb73131c732e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad36fd135f8a75f0c170cd3e8241797c5994870d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..728edd6a08e7a28f3d959ac8acae041eaf0fa28c Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/adam.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/adam.py new file mode 100644 index 0000000000000000000000000000000000000000..fa817573f734801070de5391d161695389351ec3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/adam.py @@ -0,0 +1,310 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +import torch +import numpy as np +from deepspeed.accelerator import get_accelerator +from deepspeed.utils.torch import required_torch_version +from deepspeed import comm as dist + + +class OnebitAdam(torch.optim.Optimizer): + """Implements the 1-bit Adam algorithm. Currently GPU-only. + For usage example please see https://www.deepspeed.ai/tutorials/onebit-adam/ + For technical details please read https://arxiv.org/abs/2102.02888 + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + freeze_step (int, optional): Number of steps for warmup (uncompressed) + stage before we start using compressed communication. (default 100000) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in 1-bit Adam! + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + cuda_aware (boolean, required): Set True if the underlying MPI implementation + supports CUDA-Aware communication. (default: False) + comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') + .. _Adam\\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + deepspeed=None, + lr=1e-3, + freeze_step=100000, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + amsgrad=False, + cuda_aware=False, + comm_backend_name='nccl'): + + if amsgrad: + raise RuntimeError('1-bit Adam does not support the AMSGrad variant.') + + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm) + + super(OnebitAdam, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.comm_time = 0.0 + self.step_time = 0.0 + self.ave_step = 1 + self.bk_time = 0.0 + + self.deepspeed = deepspeed + self.adam_freeze_key = False + self.initialize = False + self.freeze_step = freeze_step + self.cuda_aware = cuda_aware + self.using_pipeline = False + + self.comm_backend_name = comm_backend_name + + assert dist.is_initialized(), "Please initialize the torch distributed backend." + # Empty initializer. Set handle based on the comm backend as follows. + self.comm_backend_handle = None + if self.comm_backend_name == 'nccl': + assert ( + required_torch_version(min_version=1.8) + ), "Please use torch 1.8 or greater to enable NCCL backend in 1-bit Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" + from deepspeed.runtime.comm.nccl import NcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'mpi': + from deepspeed.runtime.comm.mpi import MpiBackend + self.comm_backend_handle = MpiBackend(cuda_aware) + elif self.comm_backend_name == 'hccl': + from deepspeed.runtime.comm.hccl import HcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'compressed': + from deepspeed.runtime.comm.compressed import CompressedBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) + self.size = self.comm_backend_handle.size + + self.divider = int(self.size * 8 / np.gcd(self.size, 8)) + + def step(self, closure=None, grads=None): + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + output params (list of tensors, optional): A reduced precision copy + of the updated weights written out in addition to the regular + updated weights. Have to be of same type as gradients. (default: None) + scale (float, optional): factor to divide gradient tensor values + by before applying to weights. (default: 1) + """ + loss = None + if closure is not None: + loss = closure() + + gather_time = 0 + allgather_time = 0 + all_time = 0 + + if self.adam_freeze_key is False: + v_diff_buffer = 0.0 + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + for group, grads_this_group in zip(self.param_groups, grads_group): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad in zip(group['params'], grads_this_group): + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('1-bit Adam does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + + if not self.initialize or (self.adam_freeze_key and 'worker_error' not in state.keys()): + state['tensor_size'] = torch.numel(p.data) + state['corrected_tensor_size'] = state['tensor_size'] + + if state['tensor_size'] % (self.size * self.divider) != 0: + state['corrected_tensor_size'] += ((self.size * self.divider) - (state['tensor_size'] % + (self.size * self.divider))) + state['server_chunk_size'] = state['corrected_tensor_size'] // self.size + get_accelerator().empty_cache() + state['worker_error'] = torch.zeros(state['corrected_tensor_size'], device=p.device) + state['server_error'] = torch.zeros(state['server_chunk_size'], device=p.device) + get_accelerator().empty_cache() + self.adam_freeze_key = True + if not self.initialize and dist.get_rank() == 0: + print("Cupy Buffers Initialized Successfully.") + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + state['step'] += 1 + + if self.adam_freeze_key is False: + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + grad = None + if self.initialize: + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + + else: + if 'non_freeze' in group.keys() and group['non_freeze'] is True: + dist.all_reduce(grad) + grad.mul_(1 / dist.get_world_size()) + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + grad = None + else: + if self.initialize is True: + exp_avg.mul_(beta1).add_(1 - beta1, grad) + grad = None + + if self.size > 1: + exp_avg.set_( + self.comm_backend_handle.compressed_allreduce(exp_avg, state['worker_error'], + state['server_error'], + self.deepspeed.local_rank)) + # Because 1-bit compression cannot represent exact zero, it is required to + # provide a momentum mask for those params that have constant exact zeros in their + # momentums, otherwise the compression error would keep accumulating. + # For example, for BERT pre-training seq 128, bert.embeddings.position_embeddings.weight + # always have exact zeros in its momentum for row 129 to 512, because it only + # learns up to seq length 128 while the model supports up to 512 seq length. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + if 'exp_avg_mask' in group: + if exp_avg.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=exp_avg.device) + exp_avg.mul_(group['exp_avg_mask']) + + if self.initialize: + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + + if self.initialize: + if group['weight_decay'] > 0.0: + update += group['weight_decay'] * p.data + with torch.no_grad(): + p.add_(-group['lr'] * update) + + if not self.initialize: + print('Pop out errors', flush=True) + state.pop('worker_error') + state.pop('server_error') + + if not self.initialize: + self.adam_freeze_key = False + self.initialize = True + print(f"Finished the initialization step at rank {dist.get_rank()}") + return loss + + if self.adam_freeze_key is False: + if state['step'] >= self.freeze_step: + print('OnebitAdam - starting compressed communication') + self.adam_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + + return loss + + def load_state_dict(self, state_dict): + """ + Overrides load_state_dict() to add special handling when loading checkpoints + """ + # Because at different stage exp_avg_mask may change (e.g., + # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask + # in checkpoints but always use the one user provided in training script. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + # Thus here we keep the exp_avg_mask unchanged when loading checkpoint + for i, group in enumerate(self.param_groups): + if 'exp_avg_mask' in group: + state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] + elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: + state_dict['param_groups'][i].pop('exp_avg_mask') + super().load_state_dict(state_dict) + if self.state[self.param_groups[0]['params'][0]]['step'] < self.freeze_step: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitAdam warmup stage starts/continues.") + if self.adam_freeze_key is True: + self.adam_freeze_key = False + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + else: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitAdam compression stage starts/continues.") + if self.adam_freeze_key is False: + self.adam_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + # We reset the compression errors when loading checkpoints for 3 reasons: + # 1) The worker and server error at each GPU are distinct, so in current implementation + # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors. + # If we want to save them correctly we need O(num_gpu*model_size) memory in order to + # gather all the error, which is a very large memory requirement. It's possible to save + # them in a distributed way, but it will make the checkpoint saving/loading much more complicated. + # 2) Even if we are able to save the compression errors correctly, you need to have the + # exact same number of GPUs in order to load them correctly. + # 3) We verified on BERT pre-training that occasionally resetting the compression error + # at checkpoint loading does not affect the convergence. + # However, please avoid frequent checkpoint loading which could break the error + # compensation mechanism thus affect the convergence. + for group in self.param_groups: + for p in group['params']: + if 'worker_error' in self.state[p]: + self.state[p].pop('worker_error') + if 'server_error' in self.state[p]: + self.state[p].pop('server_error') diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/lamb.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/lamb.py new file mode 100644 index 0000000000000000000000000000000000000000..9e7bae816ecdee5aa7e5aa2349a62f1b43ea5e10 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/lamb.py @@ -0,0 +1,447 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +import torch +import numpy as np +from deepspeed import comm as dist +from deepspeed.utils.torch import required_torch_version +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed.accelerator import get_accelerator + + +class OnebitLamb(torch.optim.Optimizer): + """Implements the 1-bit Lamb algorithm. Currently GPU-only. + For usage example please see https://www.deepspeed.ai/tutorials/onebit-lamb/ + For technical details please see our paper https://arxiv.org/abs/2104.06069. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + freeze_step (int, optional): Number of steps for warmup (uncompressed) + stage before we start using compressed communication. (default 100000) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + max_coeff(float, optional): maximum value of the lamb coefficient (default: 10.0) + min_coeff(float, optional): minimum value of the lamb coefficient (default: 0.01) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in 1-bit Lamb! + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + cuda_aware (boolean, required): Set True if the underlying MPI implementation + supports CUDA-Aware communication. (default: False) + comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') + coeff_beta (float, optional): coefficient used for computing + running averages of lamb coefficient (default: 0.9) note that you may want to + increase or decrease this beta depending on the freeze_step you choose, as + 1/(1 - coeff_beta) should be smaller than or equal to freeze_step + factor_max (float, optional): maximum value of scaling factor to the frozen lamb + coefficient during compression stage (default: 4.0) + factor_min (float, optional): minimum value of scaling factor to the frozen lamb + coefficient during compression stage (default: 0.5) + factor_threshold (float, optional): threshold of how much the scaling factor can + fluctuate between steps (default: 0.1) + .. _Large Batch Optimization for Deep Learning\\: Training BERT in 76 minutes: + https://arxiv.org/abs/1904.00962 + .. _Adam\\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + deepspeed=None, + lr=1e-3, + freeze_step=100000, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + max_coeff=10.0, + min_coeff=0.01, + amsgrad=False, + cuda_aware=False, + comm_backend_name='nccl', + coeff_beta=0.9, + factor_max=4.0, + factor_min=0.5, + factor_threshold=0.1): + + if amsgrad: + raise RuntimeError('1-bit Lamb does not support the AMSGrad variant.') + + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm, + max_coeff=max_coeff, + min_coeff=min_coeff) + + super(OnebitLamb, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.deepspeed = deepspeed + self.lamb_freeze_key = False + self.initialize = False + self.freeze_step = freeze_step + self.cuda_aware = cuda_aware + self.coeff_beta = coeff_beta + self.factor_max = factor_max + self.factor_min = factor_min + self.factor_threshold = factor_threshold + self.using_pipeline = False + + self.comm_backend_name = comm_backend_name + + assert dist.is_initialized(), "Please initialize the torch distributed backend." + # Empty initializer. Set handle based on the comm backend as follows. + self.comm_backend_handle = None + if self.comm_backend_name == 'nccl': + assert ( + required_torch_version(min_version=1.8) + ), "Please use torch 1.8 or greater to enable NCCL backend in 1-bit Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" + from deepspeed.runtime.comm.nccl import NcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'mpi': + from deepspeed.runtime.comm.mpi import MpiBackend + self.comm_backend_handle = MpiBackend(cuda_aware) + elif self.comm_backend_name == 'hccl': + from deepspeed.runtime.comm.hccl import HcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'compressed': + from deepspeed.runtime.comm.compressed import CompressedBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) + + self.size = self.comm_backend_handle.size + + self.divider = int(self.size * 8 / np.gcd(self.size, 8)) + + self.exp_avg_flat = [] + self.dummy_exp_avg = {} + self.corrected_tensor_sizes = [] + self.server_chunk_sizes = [] + self.worker_errors = [] + self.server_errors = [] + + self.lamb_coeffs = [] + + def step(self, closure=None, grads=None): + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + """ + loss = None + if closure is not None: + loss = closure() + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + # remove the previous stats + del self.lamb_coeffs[:] + + if self.lamb_freeze_key: + exp_avg_last_step = [] + for group in self.param_groups: + exp_avg_last_step.append([self.state[p]['exp_avg'].detach().clone() for p in group['params']]) + if 'scaling_coeff' not in self.state[self.param_groups[0]['params'][0]]: + # Compute the scaling_coeff for each momentum at the end of warmup stage. + # This is used to reduce compression error during compression stage. + momentum_scales = [] + for group in self.param_groups: + momentum_scales.append([(torch.linalg.vector_norm(self.state[p]['exp_avg']) / + np.sqrt(torch.numel(self.state[p]['exp_avg']))).item() + for p in group['params']]) + united_scale = sum([sum(x) for x in momentum_scales]) / sum([len(x) for x in momentum_scales]) + for i, group in enumerate(self.param_groups): + for j, p in enumerate(group['params']): + self.state[p]['scaling_coeff'] = united_scale / momentum_scales[i][j] + + for group, grads_this_group in zip(self.param_groups, grads_group): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad in zip(group['params'], grads_this_group): + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('1-bit Lamb does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0 or (len(state) == 1 and 'scaling_coeff' in state.keys()): + state['step'] = 0 + state['lamb_coeff_freeze'] = 0.0 + state['last_factor'] = 1.0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + state['exp_avg_sq_fresh'] = torch.zeros_like(p.data) + + if not self.initialize: + self.lamb_freeze_key = True + + exp_avg, exp_avg_sq, exp_avg_sq_fresh = state['exp_avg'], state['exp_avg_sq'], state[ + 'exp_avg_sq_fresh'] + beta1, beta2 = group['betas'] + max_coeff = group['max_coeff'] + min_coeff = group['min_coeff'] + + state['step'] += 1 + + if self.lamb_freeze_key is False: + # warmup stage, baseline Lamb optimization + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + if state['step'] == self.freeze_step: + exp_avg_sq_fresh.data = exp_avg_sq.detach().clone() + grad = None + if self.initialize: + weight_norm = p.data.pow(2).sum().sqrt() + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + if group['weight_decay'] > 0.0: + update += group['weight_decay'] * p.data + update_norm = update.pow(2).sum().sqrt() + lamb_coeff = 1.0 + if weight_norm != 0 and update_norm != 0: + lamb_coeff = (weight_norm / update_norm).item() + if lamb_coeff > max_coeff: + lamb_coeff = max_coeff + if lamb_coeff < min_coeff: + lamb_coeff = min_coeff + if lamb_coeff != 1.0: + state['lamb_coeff_freeze'] = self.coeff_beta * state['lamb_coeff_freeze'] + ( + 1 - self.coeff_beta) * lamb_coeff + self.lamb_coeffs.append(lamb_coeff) + with torch.no_grad(): + p.add_(-group['lr'] * lamb_coeff * update) + else: + # compression stage, update each momentum locally, then + # communicate based on the compressed_allreduce below + if self.initialize: + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg.mul_(self.state[p]['scaling_coeff']) + grad = None + + # init fused momentum + if len(self.exp_avg_flat) == 0: + momentum_groups = [] + tensor_size = 0 + for group in self.param_groups: + for p in group['params']: + momentum_groups.append(self.state[p]['exp_avg']) + tensor_size += torch.numel(p.data) + corrected_tensor_size = tensor_size + if tensor_size % (self.size * self.divider) != 0: + difference = ((self.size * self.divider) - (tensor_size % (self.size * self.divider))) + corrected_tensor_size += difference + self.dummy_exp_avg[0] = torch.zeros(difference, device=momentum_groups[0].data.device) + momentum_groups.append(self.dummy_exp_avg[0]) + self.corrected_tensor_sizes.append(corrected_tensor_size) + self.server_chunk_sizes.append(corrected_tensor_size // self.size) + + self.exp_avg_flat.append(_flatten_dense_tensors([p.detach().clone() for p in momentum_groups])) + updated_params = _unflatten_dense_tensors(self.exp_avg_flat[0], momentum_groups) + for p, q in zip(momentum_groups, updated_params): + p.data = q.data + + if self.initialize and len(self.worker_errors) == 0: + get_accelerator().empty_cache() + for i in range(len(self.exp_avg_flat)): + self.worker_errors.append( + torch.zeros(self.corrected_tensor_sizes[i], device=self.exp_avg_flat[i].device)) + self.server_errors.append(torch.zeros(self.server_chunk_sizes[i], device=self.exp_avg_flat[i].device)) + get_accelerator().empty_cache() + + if self.lamb_freeze_key: + if self.size > 1: + for i in range(len(self.exp_avg_flat)): + if not self.initialize: + get_accelerator().empty_cache() + self.worker_errors.append( + torch.zeros(self.corrected_tensor_sizes[i], device=self.exp_avg_flat[i].device)) + self.server_errors.append( + torch.zeros(self.server_chunk_sizes[i], device=self.exp_avg_flat[i].device)) + get_accelerator().empty_cache() + if dist.get_rank() == 0: + print("Cupy Buffers Initialized Successfully.") + + self.comm_backend_handle.compressed_allreduce(self.exp_avg_flat[i], self.worker_errors[0], + self.server_errors[0], self.deepspeed.local_rank) + + if dist.get_rank() == 0: + print('Pop out errors', flush=True) + del self.worker_errors[:] + del self.server_errors[:] + else: + self.comm_backend_handle.compressed_allreduce(self.exp_avg_flat[i], self.worker_errors[i], + self.server_errors[i], self.deepspeed.local_rank) + + if self.lamb_freeze_key and self.initialize: + for i, group in enumerate(self.param_groups): + bias_correction = 1 if group['bias_correction'] else 0 + + for j, p in enumerate(group['params']): + state = self.state[p] + exp_avg, exp_avg_sq, exp_avg_sq_fresh = state['exp_avg'], state['exp_avg_sq'], state[ + 'exp_avg_sq_fresh'] + beta1, beta2 = group['betas'] + exp_avg.div_(self.state[p]['scaling_coeff']) + # Because 1-bit compression cannot represent exact zero, it is required to + # provide a momentum mask for those params that have constant exact zeros in their + # momentums, otherwise the compression error would keep accumulating. + # For example, for BERT pre-training seq 128, bert.embeddings.position_embeddings.weight + # always have exact zeros in its momentum for row 129 to 512, because it only + # learns up to seq length 128 while the model supports up to 512 seq length. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py about how + # to add this exp_avg_mask for BERT pre-training.) + if 'exp_avg_mask' in group: + if exp_avg.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=exp_avg.device) + exp_avg.mul_(group['exp_avg_mask']) + + grad_reconstruct = ((exp_avg - exp_avg_last_step[i][j] * beta1) / (1 - beta1)) + exp_avg_sq_fresh.mul_(beta2).addcmul_(1 - beta2, grad_reconstruct, grad_reconstruct) + denom = exp_avg_sq.sqrt() + group['eps'] + update_prelim = exp_avg / denom + + if group['weight_decay'] > 0.0: + update = update_prelim + group['weight_decay'] * p.data + else: + update = update_prelim + + lamb_coeff = 1.0 + update_norm = update.pow(2).sum().sqrt() + denom_real = exp_avg_sq_fresh.sqrt() + group['eps'] + factor = (denom / denom_real).max().item() + if group['weight_decay'] > 0.0: + update_ratio = min(1.0, (update_prelim.pow(2).sum().sqrt() / update_norm).item()) + factor = factor * update_ratio + (1.0 - update_ratio) + if factor > self.factor_max: + factor = self.factor_max + if factor < self.factor_min: + factor = self.factor_min + if factor > state['last_factor'] * (1.0 + self.factor_threshold): + factor = state['last_factor'] * (1.0 + self.factor_threshold) + if factor < state['last_factor'] * (1.0 - self.factor_threshold): + factor = state['last_factor'] * (1.0 - self.factor_threshold) + state['last_factor'] = factor + lamb_coeff = state['lamb_coeff_freeze'] * factor + self.lamb_coeffs.append(lamb_coeff) + with torch.no_grad(): + p.add_(-group['lr'] * lamb_coeff * update) + del exp_avg_last_step[:] + exp_avg_last_step = None + + if not self.initialize: + self.lamb_freeze_key = False + self.initialize = True + print(f"Finished the initialization step at rank {dist.get_rank()}") + return loss + + if self.lamb_freeze_key is False: + if state['step'] >= self.freeze_step: + print('OnebitLamb - starting compressed communication') + self.lamb_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + + return loss + + def load_state_dict(self, state_dict): + """ + Overrides load_state_dict() to add special handling when loading checkpoints + """ + # Because at different stage exp_avg_mask may change (e.g., + # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask + # in checkpoints but always use the one user provided in training script. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + # Thus here we keep the exp_avg_mask unchanged when loading checkpoint + for i, group in enumerate(self.param_groups): + if 'exp_avg_mask' in group: + state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] + elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: + state_dict['param_groups'][i].pop('exp_avg_mask') + super().load_state_dict(state_dict) + # need to reset the fused momentum since loading states will break the linking + del self.exp_avg_flat[:] + self.dummy_exp_avg.clear() + del self.corrected_tensor_sizes[:] + del self.server_chunk_sizes[:] + if self.state[self.param_groups[0]['params'][0]]['step'] < self.freeze_step: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitLamb warmup stage starts/continues.") + if self.lamb_freeze_key is True: + self.lamb_freeze_key = False + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + for group in self.param_groups: + for p in group['params']: + self.state[p]['lamb_coeff_freeze'] = 0.0 + self.state[p]['last_factor'] = 1.0 + if 'scaling_coeff' in self.state[p]: + self.state[p].pop('scaling_coeff') + else: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitLamb compression stage starts/continues.") + if self.lamb_freeze_key is False: + self.lamb_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + # We reset the compression errors when loading checkpoints for 3 reasons: + # 1) The worker and server error at each GPU are distinct, so in current implementation + # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors. + # If we want to save them correctly we need O(num_gpu*model_size) memory in order to + # gather all the error, which is a very large memory requirement. It's possible to save + # them in a distributed way, but it will make the checkpoint saving/loading much more complicated. + # 2) Even if we are able to save the compression errors correctly, you need to have the + # exact same number of GPUs in order to load them correctly. + # 3) We verified on BERT pre-training that occasionally resetting the compression error + # at checkpoint loading does not affect the convergence. + # However, please avoid frequent checkpoint loading which could break the error + # compensation mechanism thus affect the convergence. + del self.worker_errors[:] + del self.server_errors[:] + + def get_lamb_coeffs(self): + return self.lamb_coeffs diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py new file mode 100644 index 0000000000000000000000000000000000000000..70282ec41714b202a7105925d4b051cd7571ecea --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py @@ -0,0 +1,365 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +import torch +import numpy as np +from deepspeed.accelerator import get_accelerator +from deepspeed.utils.torch import required_torch_version +from deepspeed import comm as dist + + +class ZeroOneAdam(torch.optim.Optimizer): + """ + Implements the 0/1 Adam algorithm. Currently GPU-only. + For usage example please see https://www.deepspeed.ai/tutorials/zero-one-adam/ + For technical details please read https://arxiv.org/abs/2202.06009 + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + var_freeze_step (int, optional): The latest step to update the variance, + using the notation from https://arxiv.org/abs/2202.06009, it denotes the + max{i|i in T_v}. Note that this is different from the freeze step from the + 1-bit Adam. The var_freeze_step is usually the end of the learning rate warmup + and thus does not require tuning. (default: 100000) + var_update_scaler (int, optional): The interval to update the variance. Note that + the update policy for variance follows an exponential rule, where var_update_scaler + denotes the kappa in the 0/1 Adam paper. (default: 16) + local_step_scaler (int, optional): The interval to scale the local steps interval + according to the learning rate policy. (default: 32678) + local_step_clipper (int, optional): The largest interval for local steps with + learning rate policy. This corresponds to the variable H in the 0/1 Adam paper. + (default: 16) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in 0/1 Adam! + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + cuda_aware (boolean, required): Set True if the underlying MPI implementation + supports CUDA-Aware communication. (default: False) + comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') + .. _Adam\\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + deepspeed=None, + lr=1e-3, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + var_freeze_step=100000, + var_update_scaler=16, + local_step_scaler=32678, + local_step_clipper=16, + amsgrad=False, + cuda_aware=False, + comm_backend_name='nccl'): + + if amsgrad: + raise RuntimeError('0/1 Adam does not support the AMSGrad variant.') + + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm) + + super(ZeroOneAdam, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.deepspeed = deepspeed + self.initialize = False + self.cuda_aware = cuda_aware + self.using_pipeline = False + + self.var_freeze_step = var_freeze_step + self.var_update_scaler = var_update_scaler + self.local_step_scaler = local_step_scaler + self.local_step_clipper = local_step_clipper + self.freeze_key = False + self.reinitial_error_buffer = False + + self.comm_backend_name = comm_backend_name + + assert dist.is_initialized(), "Please initialize the torch distributed backend." + # Empty initializer. Set handle based on the comm backend as follows. + self.comm_backend_handle = None + if self.comm_backend_name == 'nccl': + assert ( + required_torch_version(min_version=1.8) + ), "Please use torch 1.8 or greater to enable NCCL backend in 0/1 Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" + from deepspeed.runtime.comm.nccl import NcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'mpi': + from deepspeed.runtime.comm.mpi import MpiBackend + self.comm_backend_handle = MpiBackend(cuda_aware) + elif self.comm_backend_name == 'hccl': + from deepspeed.runtime.comm.hccl import HcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'compressed': + from deepspeed.runtime.comm.compressed import CompressedBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) + self.size = self.comm_backend_handle.size + + self.divider = int(self.size * 8 / np.gcd(self.size, 8)) + + def step(self, closure=None, grads=None): + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + output params (list of tensors, optional): A reduced precision copy + of the updated weights written out in addition to the regular + updated weights. Have to be of same type as gradients. (default: None) + scale (float, optional): factor to divide gradient tensor values + by before applying to weights. (default: 1) + """ + loss = None + if closure is not None: + loss = closure() + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + for group, grads_this_group in zip(self.param_groups, grads_group): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad in zip(group['params'], grads_this_group): + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('0/1 Adam does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + + if not self.initialize or 'worker_error' not in state.keys(): + # Some scalars to help scale the variance update/local step policies + state['var_interval'] = 1 + state['var_counter'] = 0 + state['local_step_interval'] = 1 + state['local_step_counter'] = 0 + state['lrs'] = 0 + state['tensor_size'] = torch.numel(p.data) + state['corrected_tensor_size'] = state['tensor_size'] + + if state['tensor_size'] % (self.size * self.divider) != 0: + state['corrected_tensor_size'] += ((self.size * self.divider) - (state['tensor_size'] % + (self.size * self.divider))) + state['server_chunk_size'] = state['corrected_tensor_size'] // self.size + get_accelerator().empty_cache() + state['worker_error'] = torch.zeros(state['corrected_tensor_size'], device=p.device) + state['server_error'] = torch.zeros(state['server_chunk_size'], device=p.device) + # Accumulation of momentum, i.e., the u variable in the 0/1 Adam paper + state['momentum_accumulator'] = torch.zeros_like(p.data) + get_accelerator().empty_cache() + # self.freeze_key = True + if not self.initialize and dist.get_rank() == 0: + print("Cupy Buffers Initialized Successfully.") + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + comm_buffer = state['momentum_accumulator'] + beta1, beta2 = group['betas'] + + state['step'] += 1 + + if self.initialize: + if self.freeze_key is False: + if state['step'] % state['var_interval'] == 0: + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + else: + if self.size > 1: + with torch.no_grad(): + grad_onebit = self.comm_backend_handle.compressed_allreduce( + grad, state['worker_error'], state['server_error'], self.deepspeed.local_rank) + if 'exp_avg_mask' in group: + if grad_onebit.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=grad_onebit.device) + grad_onebit.mul_(group['exp_avg_mask']) + exp_avg.mul_(beta1).add_(1 - beta1, grad_onebit) + else: + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + state['lrs'] += group['lr'] + grad = None + + if not self.initialize: + if self.size > 1: + comm_buffer.set_( + self.comm_backend_handle.compressed_allreduce(comm_buffer, state['worker_error'], + state['server_error'], + self.deepspeed.local_rank)) + if 'exp_avg_mask' in group: + if comm_buffer.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=comm_buffer.device) + comm_buffer.mul_(group['exp_avg_mask']) + + if self.initialize: + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + if group['weight_decay'] > 0.0: + update += group['weight_decay'] * p.data + with torch.no_grad(): + p.data.add_(-group['lr'] * update) + if self.freeze_key is True: + comm_buffer.add_(-group['lr'] * update) + if state['step'] % state['local_step_interval'] == 0 and self.freeze_key: + with torch.no_grad(): + p.data.add_(-1 * comm_buffer) + comm_buffer.mul_(exp_avg_sq.sqrt() + group['eps']) + if self.size > 1: + comm_buffer.copy_( + self.comm_backend_handle.compressed_allreduce(comm_buffer, state['worker_error'], + state['server_error'], + self.deepspeed.local_rank)) + if 'exp_avg_mask' in group: + if comm_buffer.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=comm_buffer.device) + comm_buffer.mul_(group['exp_avg_mask']) + exp_avg.zero_().add_(comm_buffer / state['lrs'], alpha=-1) + p.data.add_(comm_buffer / (exp_avg_sq.sqrt() + group['eps'])) + comm_buffer.zero_() + + state['lrs'] = 0 + + # According to 0/1 Adam theory, a fixed variance would allow more accurate estimation of momentum + # However, in practice, we can also disable the manual freezing of variance, since the interval of + # updating variance will increase exponentially, so that it has negligible effect on the estimation. + if self.freeze_key is False: + if state['step'] % state['var_interval'] == 0: + state['var_counter'] += 1 + if state['var_counter'] == self.var_update_scaler: + state['var_counter'] = 0 + state['var_interval'] *= 2 + if (state['step'] + 1) % state['var_interval'] == 0: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + else: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + else: + state['local_step_counter'] += 1 + if state['local_step_counter'] == self.local_step_scaler: + state['local_step_counter'] = 0 + state['local_step_interval'] = min(self.local_step_clipper, + state['local_step_interval'] * 2) + + if not self.initialize: + print('Pop out errors', flush=True) + self.freeze_key = False + state.pop('worker_error') + state.pop('server_error') + + if not self.initialize: + self.initialize = True + print(f"Finished the initialization step at rank {dist.get_rank()}") + return loss + + if self.state[self.param_groups[0]['params'][0]]['step'] > self.var_freeze_step: + self.freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + + if self.freeze_key is True and self.reinitial_error_buffer is False: + # We need to reinitialize the error buffers when local step > 1 since + # the errors will be logged for different metrics (gradient vs. accumulated momentum). + for group in self.param_groups: + for p in group['params']: + self.state[p]['worker_error'].zero_() + self.state[p]['server_error'].zero_() + self.reinitial_error_buffer = True + + return loss + + def load_state_dict(self, state_dict): + """ + Overrides load_state_dict() to add special handling when loading checkpoints + """ + # Because at different stage exp_avg_mask may change (e.g., + # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask + # in checkpoints but always use the one user provided in training script. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + # Thus here we keep the exp_avg_mask unchanged when loading checkpoint + for i, group in enumerate(self.param_groups): + if 'exp_avg_mask' in group: + state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] + elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: + state_dict['param_groups'][i].pop('exp_avg_mask') + super().load_state_dict(state_dict) + if self.state[self.param_groups[0]['params'][0]]['step'] < self.var_freeze_step: + self.var_freeze_key = False + if (self.state[self.param_groups[0]['params'][0]]['step'] + + 1) % self.state[self.param_groups[0]['params'][0]]['var_interval'] == 0: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + else: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + else: + self.var_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + self.reinitial_error_buffer = False + for group in self.param_groups: + for p in group['params']: + if 'worker_error' in self.state[p]: + self.state[p].pop('worker_error') + if 'server_error' in self.state[p]: + self.state[p].pop('server_error') + if 'momentum_accumulator' in self.state[p]: + self.state[p].pop('momentum_accumulator') diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..530355f846e2c04dfeaaf1886d79f67aa764908f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py @@ -0,0 +1,428 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from FP16_Optimizer in NVIDIA/apex +""" + +from deepspeed.moe.utils import split_params_grads_into_shared_and_expert_params +import torch +from torch._utils import _flatten_dense_tensors + +from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer +from deepspeed.runtime.utils import get_global_norm, CheckOverflow, get_weight_norm +from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE +from deepspeed.utils import logger +from deepspeed.utils.torch import required_torch_version +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT +from deepspeed.accelerator import get_accelerator +from deepspeed import comm as dist + + +class FP16_UnfusedOptimizer(DeepSpeedOptimizer): + """ + FP16 Optimizer without weight fusion to support LAMB optimizer + + For usage example please see, TODO: DeepSpeed V2 Tutorial + """ + + def __init__(self, + init_optimizer, + deepspeed=None, + static_loss_scale=1.0, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + mpu=None, + clip_grad=0.0, + fused_lamb_legacy=False): + + self.fused_lamb_legacy = fused_lamb_legacy + self._global_grad_norm = 0. + + if dist.get_rank() == 0: + logger.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ') + + if not get_accelerator().is_available(): + raise SystemError("Cannot use fp16 without accelerator.") + self.optimizer = init_optimizer + + # param groups + self.fp16_groups = [] + self.fp32_groups = [] + + # loop to deal with groups + for i, param_group in enumerate(self.optimizer.param_groups): + #fp16 weights that represents the actual model weights + self.fp16_groups.append(param_group['params']) + + #creating a fp32 copy of the weights that will be updated first then + #copied to fp16 weights + fp32_group = [p.clone().float().detach() for p in param_group['params']] + + #in case the internal optimizer needs it + for p in fp32_group: + p.requires_grad = True + + #setting the param groups in the optimizer to point to fp32 + #note these are not the weights used by the model + #the model uses the fp16 version that we added to fp16_group + self.fp32_groups.append(fp32_group) + param_group['params'] = self.fp32_groups[i] + + # we may have a way of fusing dynamic scale. Do not support for now + if dynamic_loss_scale: + self.dynamic_loss_scale = True + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = 2.0 + if dynamic_loss_args is None: + self.cur_scale = 1.0 * 2**16 + self.scale_window = 1000 + self.min_loss_scale = 0.25 + else: + self.cur_scale = dynamic_loss_args[INITIAL_LOSS_SCALE] + self.scale_window = dynamic_loss_args[SCALE_WINDOW] + self.min_loss_scale = dynamic_loss_args[MIN_LOSS_SCALE] + else: + self.dynamic_loss_scale = False + self.cur_iter = 0 + self.cur_scale = static_loss_scale + + self.custom_loss_scaler = False + self.external_loss_scale = None + + self.verbose = verbose + + self.clip_grad = clip_grad + self.norm_type = 2 + + if required_torch_version(max_version=0.4): + self.clip_grad_norm = torch.nn.utils.clip_grad_norm + else: + self.clip_grad_norm = torch.nn.utils.clip_grad_norm_ + + self.mpu = mpu + + self.overflow = False + self.overflow_checker = CheckOverflow(self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed) + + self.initialize_optimizer_states() + + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + # FP32 grad should never exist outside of the step function + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if set_to_none: + p.grad = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def step_fused_lamb(self, closure=None): + """ + Not supporting closure. + """ + # First compute norm for all group so we know if there is overflow + grads_groups_flat = [] + grads_groups = [] + norm_groups = [] + expert_norm_groups = [] + for i, group in enumerate(self.fp16_groups): + grads = [ + torch.zeros(p.size(), dtype=p.dtype, device=p.device) if p.grad is None else p.grad for p in group + ] + grads_groups.append(grads) + grads_groups_flat.append(_flatten_dense_tensors(grads)) + grads_for_norm, expert_grads_for_norm = split_params_grads_into_shared_and_expert_params(group) + norm_group_value = 0.0 + if len(grads_for_norm) > 0: + norm_group_value = get_weight_norm(_flatten_dense_tensors(grads_for_norm), mpu=self.mpu) + norm_groups.append(norm_group_value) + expert_norm_group_value = 0.0 + if len(expert_grads_for_norm) > 0: + expert_norm_group_value = get_weight_norm(_flatten_dense_tensors(expert_grads_for_norm), mpu=self.mpu) + expert_norm_groups.append(expert_norm_group_value) + + self.overflow = self.overflow_checker.check_using_norm(norm_groups + expert_norm_groups) + prev_scale = self.cur_scale + + self._update_scale(self.overflow) + if self.overflow: + if self.verbose: + logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " + "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) + return self.overflow + + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + combined_scale = self.unscale_and_clip_grads(self._global_grad_norm, apply_scale=False) + self.optimizer.step(grads=grads_groups, output_params=self.fp16_groups, scale=combined_scale) + + for fp32_group, fp16_group in zip(self.fp32_groups, self.fp16_groups): + for idx, (fp32_param, fp16_param) in enumerate(zip(fp32_group, fp16_group)): + + #remove the fp32 grad + fp32_param.grad = None + + #copy data from fp32 to fp16 + fp16_param.data.copy_(fp32_param.data) + + return self.overflow + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + def step(self, closure=None): + """ + Not supporting closure. + """ + + if self.fused_lamb_legacy: + return self.step_fused_lamb() + + self.overflow = self.overflow_checker.check() + prev_scale = self.cur_scale + + self._update_scale(self.overflow) + if self.overflow: + if self.verbose: + logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " + "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) + return self.overflow + + norm_groups = [] + for i, group in enumerate(self.fp16_groups): + grads_for_norm, _ = split_params_grads_into_shared_and_expert_params(group) + norm_group_value = 0.0 + if len(grads_for_norm) > 0: + norm_group_value = get_weight_norm(grads_for_norm, mpu=self.mpu) + norm_groups.append(norm_group_value) + + # copying gradients to fp32 to work with fp32 parameters + for fp32_param, fp16_param in zip(self.fp32_groups[i], self.fp16_groups[i]): + if fp16_param.grad is None: + fp32_param.grad = torch.zeros(fp16_param.size(), dtype=fp32_param.dtype, device=fp32_param.device) + else: + fp32_param.grad = fp16_param.grad.to(fp32_param.dtype) + + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + self.unscale_and_clip_grads(self._global_grad_norm) + + self.optimizer.step() + + for fp32_group, fp16_group in zip(self.fp32_groups, self.fp16_groups): + for idx, (fp32_param, fp16_param) in enumerate(zip(fp32_group, fp16_group)): + + #remove the fp32 grad + fp32_param.grad = None + + #copy data from fp32 to fp16 + fp16_param.data.copy_(fp32_param.data) + + return self.overflow + + def unscale_and_clip_grads(self, total_norm, apply_scale=True): + # compute combined scale factor for this group + combined_scale = self.cur_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad + if clip > 1: + combined_scale = clip * self.cur_scale + + if apply_scale: + for group in self.fp32_groups: + for param in group: + if param.grad is not None: + param.grad.data.mul_(1. / combined_scale) + + return combined_scale + + def backward(self, loss, create_graph=False, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + else: + scaled_loss = (loss.float()) * self.cur_scale + scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph) + + def _update_scale(self, skip): + if self.dynamic_loss_scale: + prev_scale = self.cur_scale + if skip: + self.cur_scale = max(self.cur_scale / self.scale_factor, self.min_loss_scale) + self.last_overflow_iter = self.cur_iter + if self.verbose: + logger.info("Grad overflow on iteration: %s", self.cur_iter) + logger.info(f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + # Ensure self.scale_window updates since last overflow + stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 + if (stable_interval > 0) and (stable_interval % self.scale_window == 0): + self.cur_scale *= self.scale_factor + if self.verbose: + logger.info(f"No Grad overflow for {self.scale_window} iterations") + logger.info(f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + if skip: + logger.info("Grad overflow on iteration %s", self.cur_iter) + logger.info("Using static loss scale of %s", self.cur_scale) + self.cur_iter += 1 + return + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + + param_groups = property(_get_param_groups, _set_param_groups) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + state_dict = {} + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['cur_scale'] = self.cur_scale + state_dict['cur_iter'] = self.cur_iter + if state_dict['dynamic_loss_scale']: + state_dict['last_overflow_iter'] = self.last_overflow_iter + state_dict['scale_factor'] = self.scale_factor + state_dict['scale_window'] = self.scale_window + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + state_dict['fp32_groups'] = self.fp32_groups + return state_dict + + # Refresh fp32 master params from fp16 copies + def refresh_fp32_params(self): + for current_group, saved_group in zip(self.fp32_groups, self.fp16_groups): + for current, saved in zip(current_group, saved_group): + current.data.copy_(saved.data) + + def load_state_dict(self, state_dict, load_optimizer_states=True): + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + # I think it should actually be ok to reload the optimizer before the model. + self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] + self.cur_scale = state_dict['cur_scale'] + self.cur_iter = state_dict['cur_iter'] + if state_dict['dynamic_loss_scale']: + self.last_overflow_iter = state_dict['last_overflow_iter'] + self.scale_factor = state_dict['scale_factor'] + self.scale_window = state_dict['scale_window'] + + if load_optimizer_states: + self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) + # At this point, the optimizer's references to the model's fp32 parameters are up to date. + # The optimizer's hyperparameters and internal buffers are also up to date. + # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still + # out of date. There are two options. + # 1: Refresh the master params from the model's fp16 params. + # This requires less storage but incurs precision loss. + # 2: Save and restore the fp32 master copies separately. + # We choose option 2. + # + # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device + # of their associated parameters, because it's possible those buffers might not exist yet in + # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been + # constructed in the same way as the one whose state_dict we are loading, the same master params + # are guaranteed to exist, so we can just copy_() from the saved master params. + for current_group, saved_group in zip(self.fp32_groups, state_dict['fp32_groups']): + for current, saved in zip(current_group, saved_group): + current.data.copy_(saved.data) + + def __repr__(self): + return repr(self.optimizer) + + def initialize_optimizer_states(self): + for i, group in enumerate(self.fp16_groups): + for param in group: + param.grad = torch.zeros(param.size(), + dtype=param.dtype, + device=get_accelerator().current_device_name()) + + for i, group in enumerate(self.fp32_groups): + for param in group: + param.grad = torch.zeros(param.size(), + dtype=param.dtype, + device=get_accelerator().current_device_name()) + + self.optimizer.step() + + for i, group in enumerate(self.fp16_groups): + for param in group: + param.grad = None + + for i, group in enumerate(self.fp32_groups): + for param in group: + param.grad = None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..b96228c90b02315fd63e547dbc28962009affd7c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py @@ -0,0 +1,445 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.inference.config import DeepSpeedInferenceConfig +from deepspeed.module_inject.replace_policy import replace_policies +from deepspeed.module_inject.utils import policy_to_ds_container +from .engine import DeepSpeedEngine +from .utils import TLinear, get_inactive_params +from deepspeed.runtime.zero import GatheredParameters +import time +import gc +import math +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from torch import nn +from deepspeed.utils import logger +from deepspeed.module_inject.layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding +from ..ops.transformer.inference.op_binding.workspace import WorkspaceOp + +try: + import transformers + OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding +except: + OPTLearnedPositionalEmbedding = None + + +class DeepSpeedHybridEngine(DeepSpeedEngine): + r"""DeepSpeed engine for training and inference.""" + inference_mp_group = None + + def __init__(self, args, model, **kwargs): + + super().__init__(args, model, **kwargs) + + # synch seed between all GPUs + _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) + dist.broadcast(_rng_state, 0) + get_accelerator().set_rng_state(_rng_state.cpu()) + + self.Z3_enabled = (self._config.zero_config.stage == 3) + self.gather_all_layers = self._config.hybrid_engine.pin_parameters + + # inference containers / fwds + self._inference_containers = [] + self._orig_modules = [] + self._orig_fwds = [] + self.create_inference_module() + + # Performance stats + self._t_start = None + self._total_latency = 0 + self._iters = 0 + self._training_start_time = None + self._generate_latency = 0 + self._training_latency = 0 + self._total_batch_size = None + self._gather_latency = 0 + + self.is_lora_fused = False + self.workspace = WorkspaceOp() + + def convert_to_linear_transposed(self, model): + + def _replace_linear_layer(r_module, parent_type=None, prev_type=None): + for name, child in r_module.named_children(): + if child.__class__ in [torch.nn.Linear] and \ + (parent_type is torch.nn.ModuleList or prev_type is torch.nn.ModuleList): + setattr(r_module, name, TLinear(child, name)) + else: + _replace_linear_layer(child, type(r_module), prev_type=parent_type) + return r_module + + _replace_linear_layer(model) + + def new_inference_container(self, orig_layer, policy_cls, layer_id): + policy = policy_cls(orig_layer, inference=True) + + if self._config.float16_config.enabled: + inference_dtype = torch.float16 + elif self._config.bfloat16_config.enabled: + inference_dtype = torch.bfloat16 + else: + inference_dtype = torch.float32 + + _container = policy_to_ds_container( + policy=policy, + config=DeepSpeedInferenceConfig( + set_empty_params=True, + dtype=inference_dtype, + max_out_tokens=self._config.hybrid_engine.max_out_tokens, + min_out_tokens=self._config.hybrid_engine.max_out_tokens, + transposed_mode=True, + ), + model_config=self.module.config if hasattr(self.module, 'config') else None, + layer_id=layer_id, + child=orig_layer) + + if self.mpu is not None: + if hasattr(self.mpu, 'get_model_parallel_world_size'): + _container.set_tensor_parallel_config(self.mpu.get_model_parallel_world_size(), + self.mpu.get_model_parallel_group()) + else: + _container.set_tensor_parallel_config(self.mpu.get_tensor_model_parallel_world_size(), + self.mpu.get_tensor_model_parallel_group()) + else: + _container.set_tensor_parallel_config(self._config.hybrid_engine.inference_tp_size, self.mp_group) + _container.initialize_tensors(enable_training=True) + _container.create_ds_model_config() + _container.create_module() + _container.set_params_wo_copy(Z3_enabled=self.Z3_enabled) + return _container + + def populate_all_inference_policies(self): + self.inference_policies = {} + for plcy in replace_policies: + _ = plcy(None) + if isinstance(plcy._orig_layer_class, list): + for orig_layer_class in plcy._orig_layer_class: + self.inference_policies.update({orig_layer_class: (self.new_inference_container, plcy)}) + elif plcy._orig_layer_class is not None: + self.inference_policies.update({plcy._orig_layer_class: (self.new_inference_container, plcy)}) + self.inference_policies.update({ + nn.Linear: (LinearLayer, ), + nn.Embedding: (EmbeddingLayer, ), + nn.LayerNorm: (Normalize, ), + OPTLearnedPositionalEmbedding: (OPTEmbedding, ) + }) + + def _fuse_lora_layer(self, layer_id): + self._inference_containers[layer_id].fuse_lora() + + def fuse_lora_weight(self): + for layer_id in range(len(self.layer_params)): + self._fuse_lora_layer(layer_id) + + def _unfuse_lora_layer(self, layer_id): + self._inference_containers[layer_id].unfuse_lora() + + def unfuse_lora_weight(self): + for layer_id in range(len(self.layer_params)): + self._unfuse_lora_layer(layer_id) + + def unfuse_lora_weight_non_pinned(self): + for layer_id in range(len(self.layer_params)): + non_active_params = get_inactive_params(self.layer_params[layer_id]) + non_active_lora_params = get_inactive_params(self.layer_lora_params[layer_id]) + non_active_params.extend(non_active_lora_params) + + with GatheredParameters(non_active_params): + self._unfuse_lora_layer(layer_id) + + def retake_inference_cache(self): + if self._config.hybrid_engine.release_inference_cache: + retake_success = self.workspace.retake_workspace() + + if not retake_success: + logger.warning("Unable to acquire workspace on first attempt, emptying cache and retrying.") + gc.collect() + get_accelerator().empty_cache() + retake_success = self.workspace.retake_workspace() + + if not retake_success: + raise RuntimeError("Unable to retake inference workspace.") + + def generate(self, *inputs, **kwargs): + if self._total_batch_size is None: + bsz = inputs[0].shape[0] if len(inputs) > 0 else \ + kwargs['input_ids'].shape[0] + self._total_batch_size = bsz * dist.get_world_size() + + self._t0 = time.time() + + if self.Z3_enabled and self.gather_all_layers: + if self._config.hybrid_engine.inference_tp_size > 1: + non_tp_params = [] + for other_layer in self._other_layers: + non_tp_params.extend(list(other_layer.parameters())) + + partition_size = self._config.hybrid_engine.tp_gather_partition_size + + layer_groups = math.ceil(len(self.layer_params) / partition_size) + for lg in range(layer_groups): + non_active_params = [] + non_active_lora_params = [] + for layer_id in range(lg * partition_size, min(len(self.layer_params), (lg + 1) * partition_size), + 1): + non_tp_params.extend(self.layer_params[layer_id][:4]) + non_active_params.extend(get_inactive_params(self.layer_params[layer_id])) + non_active_params.extend(get_inactive_params(self.layer_lora_params[layer_id])) + with GatheredParameters(non_active_params): + for layer_id in range(lg * partition_size, + min(len(self.layer_params), (lg + 1) * partition_size), 1): + if len(self.all_lora_params) > 0: + self._fuse_lora_layer(layer_id) + + if self.mpu is not None: + self._inference_containers[layer_id].apply_tensor_parallelism(self.mp_replace, + reversed_dim=True) + + # TODO(cmikeh2) Evaluate if this can be deferred when release_inference_cache + # is enabled. + gc.collect() + get_accelerator().empty_cache() + + self._gather_latency = time.time() - self._t0 + + input_shape = inputs[0].shape if len(inputs) > 0 else \ + kwargs['input_ids'].shape + output = torch.zeros( + (input_shape[0] * self._config.hybrid_engine.inference_tp_size, ) + input_shape[1:], + dtype=inputs[0].dtype if len(inputs) > 0 else kwargs['input_ids'].dtype, + device=inputs[0].device if len(inputs) > 0 else kwargs['input_ids'].device) + input_cont = inputs[0].contiguous() if len(inputs) > 0 else kwargs['input_ids'].contiguous() + dist.all_gather_into_tensor(output, input_cont, group=self.mp_group) + + if len(inputs) > 0: + inputs = (output, *inputs[1:]) + else: + kwargs['input_ids'] = output + + self.retake_inference_cache() + + non_active_params = get_inactive_params(non_tp_params) + with GatheredParameters(non_active_params): + generate_ret_vals = self._generate(*inputs, **kwargs) + + for layer_id in range(len(self.layer_params)): + self._inference_containers[layer_id].release_memory() + + rank = dist.get_rank(group=self.mp_group) + generate_ret_vals = generate_ret_vals[input_shape[0] * rank:input_shape[0] * (rank + 1)] + + else: + non_active_layers = get_inactive_params(self.all_layers_params) + non_active_lora_params = get_inactive_params(self.all_lora_params) + non_active_layers.extend(non_active_lora_params) + with GatheredParameters(non_active_layers): + self._gather_latency = time.time() - self._t0 + + if len(self.all_lora_params) > 0: + self.fuse_lora_weight() + + self.retake_inference_cache() + generate_ret_vals = self._generate(*inputs, **kwargs) + + if len(self.all_lora_params) > 0: + self.unfuse_lora_weight() + else: + if len(self.all_lora_params) > 0 and (not self.Z3_enabled): + self.fuse_lora_weight() + + self.retake_inference_cache() + generate_ret_vals = self._generate(*inputs, **kwargs) + + if len(self.all_lora_params) > 0: + if (not self.Z3_enabled): + self.unfuse_lora_weight() + else: + self.unfuse_lora_weight_non_pinned() + self.is_lora_fused = False + + if self._config.hybrid_engine.release_inference_cache: + self.workspace.release_workspace() + gc.collect() + get_accelerator().empty_cache() + + self._generate_latency = time.time() - self._t0 - self._gather_latency + + return generate_ret_vals + + def create_inference_containers(self, module, layer_id=0): + for name, child in module.named_children(): + if child.__class__ in self.inference_policies: + if self.inference_policies[child.__class__][0] == self.new_inference_container: + self._inference_containers.append(self.inference_policies[child.__class__][0]( + child, self.inference_policies[child.__class__][-1], layer_id)) + self._orig_modules.append(child) + self._orig_fwds.append(child.forward) + + self.layer_params.append(self._inference_containers[layer_id].get_all_params()) + + self.lora_params.append(self._inference_containers[layer_id].get_lora_params()) + self.layer_lora_params.append([]) + for lora_param in self.lora_params[layer_id]: + self.layer_lora_params[layer_id].extend(lora_param[:-1]) + self.all_lora_params.extend(lora_param[:-1]) + + layer_id += 1 + else: + if self.inference_policies[child.__class__][0] == LinearLayer: + self._other_layers.append(self.inference_policies[child.__class__][0](module=child, + mp_group=None, + skip_partition=True)) + else: + self._other_layers.append(self.inference_policies[child.__class__][0]( + weight=child.weight, bias=child.bias if hasattr(child, 'bias') else None)) + self._orig_modules_others.append(child) + self._orig_fwds_others.append(child.forward) + else: + self.create_inference_containers(child, layer_id=layer_id) + + def create_inference_module(self): + self.layer_params = [] + self.layer_lora_params = [] + self.lora_params = [] + self.all_lora_params = [] + + self._other_layers = [] + self._orig_modules_others = [] + self._orig_fwds_others = [] + + if self._config.hybrid_engine.inference_tp_size > 1: + if self.mpu is None: + global_rank = dist.get_rank() + world_size = dist.get_world_size() + mp_group_id = global_rank // self._config.hybrid_engine.inference_tp_size + num_mp_groups = world_size // self._config.hybrid_engine.inference_tp_size + for mp_group_id in range(num_mp_groups): + ranks = list( + range(mp_group_id * self._config.hybrid_engine.inference_tp_size, \ + (mp_group_id + 1) * self._config.hybrid_engine.inference_tp_size, \ + 1) + ) + mp_group = dist.new_group(ranks) + if global_rank in ranks: + # mp_group is used for broader collective + self.mp_group = mp_group + + # mp_replace is used for container tensor slicing + from deepspeed.module_inject import ReplaceWithTensorSlicing + self.mp_replace = ReplaceWithTensorSlicing( + mp_group=self.mp_group, + mp_size=self._config.hybrid_engine.inference_tp_size, + out_dim=0, + in_dim=1) + + else: + self.mp_group = self.mpu.get_model_parallel_group() if hasattr(self.mpu, 'get_model_parallel_group') else \ + self.mpu.get_tensor_model_parallel_group() + + from deepspeed.module_inject import ReplaceWithTensorSlicing + self.mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group, + mp_size=self._config.hybrid_engine.inference_tp_size, + out_dim=0, + in_dim=1) + else: + self.mp_group = None + self.mp_replace = None + self.populate_all_inference_policies() + self.all_layers_params = list(self.module.parameters()) + self.create_inference_containers(self.module) + + if len(self._inference_containers) > 0: + self._generate = self.module.generate + self.module.generate = self.generate + + self._t0 = time.time() + + def _zero3_forward(self, layer_id): + + def run_forward(*inputs, **kwargs): + non_active_params = get_inactive_params(self.layer_params[layer_id]) + non_active_lora_params = get_inactive_params(self.layer_lora_params[layer_id]) + non_active_params.extend(non_active_lora_params) + + with GatheredParameters(non_active_params): + if len(self.all_lora_params) > 0: + # Use the is_lora_fused flag to prevent multiple fusion in Z3 with non-pinned memory + if not self.is_lora_fused: + self._fuse_lora_layer(layer_id) + # Set the is_lora_fused to true when reaching the last layer + if layer_id == len(self.layer_params) - 1: + self.is_lora_fused = True + return self._inference_containers[layer_id].module.forward(*inputs, **kwargs) + + return run_forward + + def eval(self): + if self._t_start is not None: + latency = time.time() - self._t_start + self._total_latency = self._total_latency + latency + self._iters = self._iters + 1 + if not dist.is_initialized() or dist.get_rank() == 0: + if self._total_batch_size is not None: + cur_samples_p_sec = f'|CurSamplesPerSec={(1 / latency * self._total_batch_size):.2f} ' + avg_samples_p_sec = f'|AvgSamplesPerSec={(1 / (self._total_latency / self._iters) * self._total_batch_size):.2f}' + else: + cur_samples_p_sec = '' + avg_samples_p_sec = '' + others = latency - (self._generate_latency + self._training_latency) + print(f'|E2E latency={(latency):.2f}s ' + \ + f'|Gather latency={self._gather_latency:.2f}s ({(self._gather_latency / latency * 100):.2f}%) ' + f'|Generate time={(self._generate_latency):.2f}s ({(self._generate_latency / latency * 100):.2f}%) ' + \ + f'|Training time={(self._training_latency):.2f}s ({(self._training_latency / latency * 100):.2f}%) ' + \ + f'|Others={others:.2f} ({(others / latency * 100):.2f}%)' + \ + cur_samples_p_sec + \ + avg_samples_p_sec) + self._t_start = time.time() + self._training_latency = 0 + super().eval() + if len(self._inference_containers) > 0: + for i, (orig_module, inference_container) in enumerate(zip(self._orig_modules, + self._inference_containers)): + if self.Z3_enabled and not self.gather_all_layers: + orig_module.forward = self._zero3_forward(i) + else: + orig_module.forward = inference_container.module.forward + + inference_container.transform_for_inference() + + if not self.Z3_enabled or self.gather_all_layers: + for orig_module, inference_layer in zip(self._orig_modules_others, self._other_layers): + orig_module.forward = inference_layer.forward + if self.Z3_enabled: + gc.collect() + get_accelerator().empty_cache() + if self._t_start is None: + self._t_start = time.time() + + def train(self, mode=True): + if mode and len(self._orig_modules) > 0: + for inference_container, orig_module, orig_fwd in zip(self._inference_containers, self._orig_modules, + self._orig_fwds): + inference_container.transform_for_training() + orig_module.forward = orig_fwd + for orig_module, orig_fwd in zip(self._orig_modules_others, self._orig_fwds_others): + orig_module.forward = orig_fwd + super().train(mode) + if mode: + self._training_start_time = time.time() + + def step(self, lr_kwargs=None): + super().step(lr_kwargs=lr_kwargs) + + if len(self._inference_containers) > 0: + if not self.Z3_enabled: + for inference_container in self._inference_containers: + inference_container.reset_params() + + if self._training_start_time is not None: + self._training_latency += (time.time() - self._training_start_time) + self._training_start_time = time.time() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.py new file mode 100644 index 0000000000000000000000000000000000000000..bb0866e0d52d19bc6c5271007946ee6b1a6125f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.py @@ -0,0 +1,881 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Implementation of learning rate schedules. + +Taken and modified from PyTorch v1.0.1 source +https://github.com/pytorch/pytorch/blob/v1.1.0/torch/optim/lr_scheduler.py +""" + +import argparse +from torch.optim import Optimizer +import math +from deepspeed.utils import logger + +LR_SCHEDULE = 'lr_schedule' +LR_RANGE_TEST = 'LRRangeTest' +ONE_CYCLE = 'OneCycle' +WARMUP_LR = 'WarmupLR' +WARMUP_DECAY_LR = 'WarmupDecayLR' +WARMUP_COSINE_LR = 'WarmupCosineLR' +VALID_LR_SCHEDULES = [LR_RANGE_TEST, ONE_CYCLE, WARMUP_LR, WARMUP_DECAY_LR, WARMUP_COSINE_LR] + +LR_RANGE_TEST_MIN_LR = 'lr_range_test_min_lr' +LR_RANGE_TEST_STEP_RATE = 'lr_range_test_step_rate' +LR_RANGE_TEST_STEP_SIZE = 'lr_range_test_step_size' +LR_RANGE_TEST_STAIRCASE = 'lr_range_test_staircase' + +EDGE_VALUE = 'edge_value' +MID_VALUE = 'mid_value' + +CYCLE_FIRST_STEP_SIZE = 'cycle_first_step_size' +CYCLE_FIRST_STAIR_COUNT = 'cycle_first_stair_count' +CYCLE_SECOND_STEP_SIZE = 'cycle_second_step_size' +CYCLE_SECOND_STAIR_COUNT = 'cycle_second_stair_count' +DECAY_STEP_SIZE = 'decay_step_size' + +CYCLE_MIN_LR = 'cycle_min_lr' +CYCLE_MAX_LR = 'cycle_max_lr' +DECAY_LR_RATE = 'decay_lr_rate' + +CYCLE_MIN_MOM = 'cycle_min_mom' +CYCLE_MAX_MOM = 'cycle_max_mom' +DECAY_MOM_RATE = 'decay_mom_rate' + +WARMUP_MIN_LR = 'warmup_min_lr' +WARMUP_MAX_LR = 'warmup_max_lr' +WARMUP_NUM_STEPS = 'warmup_num_steps' +WARMUP_TYPE = 'warmup_type' +WARMUP_LOG_RATE = 'log' +WARMUP_LINEAR_RATE = 'linear' + +WARMUP_MIN_RATIO = 'warmup_min_ratio' +COS_MIN_RATIO = 'cos_min_ratio' + +TOTAL_NUM_STEPS = 'total_num_steps' + + +def add_tuning_arguments(parser): + group = parser.add_argument_group('Convergence Tuning', 'Convergence tuning configurations') + + # LR scheduler + group.add_argument('--lr_schedule', type=str, default=None, help='LR schedule for training.') + + # Learning rate range test + group.add_argument("--lr_range_test_min_lr", type=float, default=0.001, help='Starting lr value.') + group.add_argument("--lr_range_test_step_rate", type=float, default=1.0, help='scaling rate for LR range test.') + group.add_argument("--lr_range_test_step_size", type=int, default=1000, help='training steps per LR change.') + group.add_argument("--lr_range_test_staircase", + type=bool, + default=False, + help='use staircase scaling for LR range test.') + + # OneCycle schedule + group.add_argument("--cycle_first_step_size", + type=int, + default=1000, + help='size of first step of 1Cycle schedule (training steps).') + group.add_argument("--cycle_first_stair_count", + type=int, + default=-1, + help='first stair count for 1Cycle schedule.') + group.add_argument("--cycle_second_step_size", + type=int, + default=-1, + help='size of second step of 1Cycle schedule (default first_step_size).') + group.add_argument("--cycle_second_stair_count", + type=int, + default=-1, + help='second stair count for 1Cycle schedule.') + group.add_argument("--decay_step_size", + type=int, + default=1000, + help='size of intervals for applying post cycle decay (training steps).') + + # 1Cycle LR + group.add_argument("--cycle_min_lr", type=float, default=0.01, help='1Cycle LR lower bound.') + group.add_argument("--cycle_max_lr", type=float, default=0.1, help='1Cycle LR upper bound.') + group.add_argument("--decay_lr_rate", type=float, default=0.0, help='post cycle LR decay rate.') + + # 1Cycle Momentum + group.add_argument('--cycle_momentum', default=False, action='store_true', help='Enable 1Cycle momentum schedule.') + group.add_argument("--cycle_min_mom", type=float, default=0.8, help='1Cycle momentum lower bound.') + group.add_argument("--cycle_max_mom", type=float, default=0.9, help='1Cycle momentum upper bound.') + group.add_argument("--decay_mom_rate", type=float, default=0.0, help='post cycle momentum decay rate.') + + # Warmup LR + group.add_argument('--warmup_min_lr', type=float, default=0, help='WarmupLR minimum/initial LR value') + group.add_argument('--warmup_max_lr', type=float, default=0.001, help='WarmupLR maximum LR value.') + group.add_argument('--warmup_num_steps', type=int, default=1000, help='WarmupLR step count for LR warmup.') + group.add_argument('--warmup_type', + type=str, + default=WARMUP_LOG_RATE, + help='WarmupLR increasing function during warmup') + + # WarmUP cos LR + group.add_argument("--warmup_min_ratio", type=float, default=0.01, help='Cosine LR lower bound.') + group.add_argument("--cos_min_ratio", type=float, default=0.01, help='Cosine LR lower bound.') + + return parser + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser = add_tuning_arguments(parser) + + lr_sched_args, unknown_args = parser.parse_known_args() + return lr_sched_args, unknown_args + + +def override_lr_range_test_params(args, params): + if hasattr(args, LR_RANGE_TEST_MIN_LR) and args.lr_range_test_min_lr is not None: + params[LR_RANGE_TEST_MIN_LR] = args.lr_range_test_min_lr + + if hasattr(args, LR_RANGE_TEST_STEP_RATE) and args.lr_range_test_step_rate is not None: + params[LR_RANGE_TEST_STEP_RATE] = args.lr_range_test_step_rate + + if hasattr(args, LR_RANGE_TEST_STEP_SIZE) and args.lr_range_test_step_size is not None: + params[LR_RANGE_TEST_STEP_SIZE] = args.lr_range_test_step_size + + if hasattr(args, LR_RANGE_TEST_STAIRCASE) and args.lr_range_test_staircase is not None: + params[LR_RANGE_TEST_STAIRCASE] = args.lr_range_test_staircase + + +def override_1cycle_params(args, params): + if hasattr(args, CYCLE_FIRST_STEP_SIZE) and args.cycle_first_step_size is not None: + params[CYCLE_FIRST_STEP_SIZE] = args.cycle_first_step_size + + if hasattr(args, CYCLE_FIRST_STAIR_COUNT) and args.cycle_first_stair_count is not None: + params[CYCLE_FIRST_STAIR_COUNT] = args.cycle_first_stair_count + + if hasattr(args, CYCLE_SECOND_STEP_SIZE) and args.cycle_second_step_size is not None: + params[CYCLE_SECOND_STEP_SIZE] = args.cycle_second_step_size + + if hasattr(args, CYCLE_SECOND_STAIR_COUNT) and args.cycle_second_stair_count is not None: + params[CYCLE_SECOND_STAIR_COUNT] = args.cycle_second_stair_count + + if hasattr(args, DECAY_STEP_SIZE) and args.decay_step_size is not None: + params[DECAY_STEP_SIZE] = args.decay_step_size + + # 1Cycle LR params + if hasattr(args, CYCLE_MIN_LR) and args.cycle_min_lr is not None: + params[CYCLE_MIN_LR] = args.cycle_min_lr + + if hasattr(args, CYCLE_MAX_LR) and args.cycle_max_lr is not None: + params[CYCLE_MAX_LR] = args.cycle_max_lr + + if hasattr(args, DECAY_LR_RATE) and args.decay_lr_rate is not None: + params[DECAY_LR_RATE] = args.decay_lr_rate + + # 1Cycle MOM params + if hasattr(args, CYCLE_MIN_MOM) and args.cycle_min_mom is not None: + params[CYCLE_MIN_MOM] = args.cycle_min_mom + + if hasattr(args, CYCLE_MAX_MOM) and args.cycle_max_mom is not None: + params[CYCLE_MAX_MOM] = args.cycle_max_mom + + if hasattr(args, DECAY_MOM_RATE) and args.decay_mom_rate is not None: + params[DECAY_MOM_RATE] = args.decay_mom_rate + + +def override_warmupLR_params(args, params): + if hasattr(args, WARMUP_MIN_LR) and args.warmup_min_lr is not None: + params[WARMUP_MIN_LR] = args.warmup_min_lr + + if hasattr(args, WARMUP_MAX_LR) and args.warmup_max_lr is not None: + params[WARMUP_MAX_LR] = args.warmup_max_lr + + if hasattr(args, WARMUP_NUM_STEPS) and args.warmup_num_steps is not None: + params[WARMUP_NUM_STEPS] = args.warmup_num_steps + + if hasattr(args, WARMUP_TYPE) and args.warmup_type is not None: + params[WARMUP_TYPE] = args.warmup_type + + +def override_params(args, params): + # LR range test params + override_lr_range_test_params(args, params) + + # 1Cycle params + override_1cycle_params(args, params) + + # WarmupLR params + override_warmupLR_params(args, params) + + +def get_config_from_args(args): + if not hasattr(args, LR_SCHEDULE) or args.lr_schedule is None: + return None, '--{} not specified on command line'.format(LR_SCHEDULE) + + if not args.lr_schedule in VALID_LR_SCHEDULES: + return None, '{} is not supported LR schedule'.format(args.lr_schedule) + + config = {} + config['type'] = args.lr_schedule + config['params'] = {} + + if args.lr_schedule == LR_RANGE_TEST: + override_lr_range_test_params(args, config['params']) + elif args.lr_schedule == ONE_CYCLE: + override_1cycle_params(args, config['params']) + else: + override_warmupLR_params(args, config['params']) + + return config, None + + +def get_lr_from_config(config): + if not 'type' in config: + return None, 'LR schedule type not defined in config' + + if not 'params' in config: + return None, 'LR schedule params not defined in config' + + lr_schedule = config['type'] + lr_params = config['params'] + + if not lr_schedule in VALID_LR_SCHEDULES: + return None, '{} is not a valid LR schedule'.format(lr_schedule) + + if lr_schedule == LR_RANGE_TEST: + return lr_params[LR_RANGE_TEST_MIN_LR], '' + if lr_schedule == ONE_CYCLE: + return lr_params[CYCLE_MAX_LR], '' + # Warmup LR + return lr_params[WARMUP_MAX_LR], '' + + +def update_lr(param_groups, lrs): + for param_group, lr in zip(param_groups, lrs): + param_group['lr'] = lr + return [group['lr'] for group in param_groups] + + +""" +Only optimizers that are subclass of torch.optim.Optimizer are supported. So check the passed optimizer and wrapped +optimizer to see if requirement is satisfied. +TODO: Looking under the hood to examine the wrapped optimizer is a hack that requires a better long-term fix. +""" + + +def get_torch_optimizer(optimizer): + if isinstance(optimizer, Optimizer): + return optimizer + + if hasattr(optimizer, 'optimizer') and isinstance(optimizer.optimizer, Optimizer): + return optimizer.optimizer + + raise TypeError('{} is not a subclass of torch.optim.Optimizer'.format(type(optimizer).__name__)) + + +class LRRangeTest(object): + """Sets the learning rate of each parameter group according to + learning rate range test (LRRT) policy. The policy increases learning + rate starting from a base value with a constant frequency, as detailed in + the paper `A disciplined approach to neural network hyper-parameters: Part 1 `_ + + LRRT policy is used for finding maximum LR that trains a model without divergence, and can be used to + configure the LR boundaries for Cyclic LR schedules. + + LRRT changes the learning rate after every batch. + `step` should be called after a batch has been used for training. + + Args: + optimizer (Optimizer): Wrapped optimizer. + lr_range_test_min_lr (float or list): Initial learning rate which is the + lower boundary in the range test for each parameter group. + lr_range_test_step_size (int): Interval of training steps to increase learning rate. Default: 2000 + lr_range_test_step_rate (float): Scaling rate for range test. Default: 1.0 + lr_range_test_staircase (bool): Scale in staircase fashion, rather than continuous. Default: False. + last_batch_iteration (int): The index of the last batch. This parameter is used when + resuming a training job. Since `step()` should be invoked after each + batch instead of after each epoch, this number represents the total + number of *batches* computed, not the total number of epochs computed. + When last_batch_iteration=-1, the schedule is started from the beginning. + Default: -1 + + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = LRRangeTest(optimizer) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: + https://arxiv.org/abs/1803.09820 +""" + + def __init__(self, + optimizer: Optimizer, + lr_range_test_min_lr: float = 1e-3, + lr_range_test_step_size: int = 2000, + lr_range_test_step_rate: float = 1.0, + lr_range_test_staircase: bool = False, + last_batch_iteration: int = -1): + + self.optimizer = get_torch_optimizer(optimizer) + + if isinstance(lr_range_test_min_lr, list) or isinstance(lr_range_test_min_lr, tuple): + if len(lr_range_test_min_lr) != len(self.optimizer.param_groups): + raise ValueError("expected {} lr_range_test_min_lr, got {}".format(len(self.optimizer.param_groups), + len(lr_range_test_min_lr))) + self.min_lr = list(lr_range_test_min_lr) + else: + self.min_lr = [lr_range_test_min_lr] * len(self.optimizer.param_groups) + + self.step_size = lr_range_test_step_size + self.step_rate = lr_range_test_step_rate + self.last_batch_iteration = last_batch_iteration + self.staircase = lr_range_test_staircase + self.interval_fn = self._staircase_interval if lr_range_test_staircase else self._continuous_interval + + if last_batch_iteration == -1: + self._last_lr = update_lr(self.optimizer.param_groups, self.min_lr) + + def _staircase_interval(self): + return math.floor(float(self.last_batch_iteration + 1) / self.step_size) + + def _continuous_interval(self): + return float(self.last_batch_iteration + 1) / self.step_size + + def _get_increase(self): + return (1 + self.step_rate * self.interval_fn()) + + def get_lr(self): + lr_increase = self._get_increase() + return [lr_range_test_min_lr * lr_increase for lr_range_test_min_lr in self.min_lr] + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def step(self, batch_iteration=None): + if batch_iteration is None: + batch_iteration = self.last_batch_iteration + 1 + self.last_batch_iteration = batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + +class OneCycle(object): + """Sets the learning rate of each parameter group according to + 1Cycle learning rate policy (1CLR). 1CLR is a variation of the + Cyclical Learning Rate (CLR) policy that involves one cycle followed by + decay. The policy simultaneously cycles the learning rate (and momentum) + between two boundaries with a constant frequency, as detailed in + the paper `A disciplined approach to neural network hyper-parameters`_. + + 1CLR policy changes the learning rate after every batch. + `step` should be called after a batch has been used for training. + + This implementation was adapted from the github repo: `PyTorch `_. + + Args: + optimizer (Optimizer): Wrapped optimizer. + cycle_min_lr (float or list): Initial learning rate which is the + lower boundary in the cycle for each parameter group. + cycle_max_lr (float or list): Upper learning rate boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (cycle_max_lr - cycle_min_lr). + The lr at any cycle is the sum of cycle_min_lr + and some scaling of the amplitude; therefore + cycle_max_lr may not actually be reached depending on + scaling function. + decay_lr_rate(float): Decay rate for learning rate. Default: 0. + cycle_first_step_size (int): Number of training iterations in the + increasing half of a cycle. Default: 2000 + cycle_second_step_size (int): Number of training iterations in the + decreasing half of a cycle. If cycle_second_step_size is None, + it is set to cycle_first_step_size. Default: None + cycle_first_stair_count(int): Number of stairs in first half of cycle phase. This means + lr/mom are changed in staircase fashion. Default 0, means staircase disabled. + cycle_second_stair_count(int): Number of stairs in second half of cycle phase. This means + lr/mom are changed in staircase fashion. Default 0, means staircase disabled. + decay_step_size (int): Intervals for applying decay in decay phase. Default: 0, means no decay. + cycle_momentum (bool): If ``True``, momentum is cycled inversely + to learning rate between 'cycle_min_mom' and 'cycle_max_mom'. + Default: True + cycle_min_mom (float or list): Initial momentum which is the + lower boundary in the cycle for each parameter group. + Default: 0.8 + cycle_max_mom (float or list): Upper momentum boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (cycle_max_mom - cycle_min_mom). + The momentum at any cycle is the difference of cycle_max_mom + and some scaling of the amplitude; therefore + cycle_min_mom may not actually be reached depending on + scaling function. Default: 0.9 + decay_mom_rate (float): Decay rate for momentum. Default: 0. + last_batch_iteration (int): The index of the last batch. This parameter is used when + resuming a training job. Since `step()` should be invoked after each + batch instead of after each epoch, this number represents the total + number of *batches* computed, not the total number of epochs computed. + When last_batch_iteration=-1, the schedule is started from the beginning. + Default: -1 + + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = OneCycle(optimizer, 0.0001, 0.0010) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + + .. _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: https://arxiv.org/abs/1803.09820 + """ + + def __init__(self, + optimizer, + cycle_min_lr, + cycle_max_lr, + decay_lr_rate=0., + cycle_first_step_size=2000, + cycle_second_step_size=None, + cycle_first_stair_count=0, + cycle_second_stair_count=None, + decay_step_size=0, + cycle_momentum=True, + cycle_min_mom=0.8, + cycle_max_mom=0.9, + decay_mom_rate=0., + last_batch_iteration=-1): + + self.optimizer = get_torch_optimizer(optimizer) + + # Initialize cycle shape + self._initialize_cycle(cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count, + cycle_second_stair_count, decay_step_size) + + # Initialize cycle lr + self._initialize_lr(self.optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration) + + # Initialize cyclic momentum + self.cycle_momentum = cycle_momentum + if cycle_momentum: + self._initialize_momentum(self.optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, + last_batch_iteration) + # Initialize batch iteration tracker + self.last_batch_iteration = last_batch_iteration + + # Configure cycle shape + + def _initialize_cycle(self, cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count, + cycle_second_stair_count, decay_step_size): + cycle_first_step_size = float(cycle_first_step_size) + cycle_second_step_size = float( + cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size + + self.total_size = cycle_first_step_size + cycle_second_step_size + self.step_ratio = cycle_first_step_size / self.total_size + self.first_stair_count = cycle_first_stair_count + self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count + self.decay_step_size = decay_step_size + + if math.isclose(self.decay_step_size, 0): + self.skip_lr_decay = True + self.skip_mom_decay = True + else: + self.skip_lr_decay = False + self.skip_mom_decay = False + + # Configure lr schedule + def _initialize_lr(self, optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration): + self.min_lrs = [cycle_min_lr] * len(optimizer.param_groups) + if last_batch_iteration == -1: + for lr, group in zip(self.min_lrs, optimizer.param_groups): + group['lr'] = lr + + self.max_lrs = [cycle_max_lr] * len(optimizer.param_groups) + self.decay_lr_rate = decay_lr_rate + + if math.isclose(self.decay_lr_rate, 0): + self.skip_lr_decay = True + + # Configure momentum schedule + def _initialize_momentum(self, optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, last_batch_iteration): + if 'betas' not in optimizer.defaults: + optimizer_name = type(optimizer).__name__ + logger.warning( + f"cycle_momentum is disabled because optimizer {optimizer_name} does not support momentum, no betas attribute in defaults" + ) + self.cycle_momentum = False + return + + self.decay_mom_rate = decay_mom_rate + self.min_moms = [(cycle_min_mom, 0.99)] * len(optimizer.param_groups) + self.max_moms = [(cycle_max_mom, 0.99)] * len(optimizer.param_groups) + + if last_batch_iteration == -1: + for momentum, group in zip(self.min_moms, optimizer.param_groups): + group['betas'] = momentum + + if math.isclose(self.decay_mom_rate, 0): + self.skip_mom_decay = True + + def _get_scale_factor(self): + batch_iteration = (self.last_batch_iteration + 1) + cycle = math.floor(1 + batch_iteration / self.total_size) + x = 1. + batch_iteration / self.total_size - cycle + if x <= self.step_ratio: + scale_factor = x / self.step_ratio + else: + scale_factor = (x - 1) / (self.step_ratio - 1) + + return scale_factor + + def _get_cycle_mom(self): + scale_factor = self._get_scale_factor() + momentums = [] + for base_betas, max_betas in zip(self.min_moms, self.max_moms): + cycle_min_mom = base_betas[0] + cycle_max_mom = max_betas[0] + base_height = (cycle_max_mom - cycle_min_mom) * scale_factor + momentum = cycle_max_mom - base_height + momentums.append((momentum, base_betas[1])) + return momentums + + def _get_cycle_lr(self): + scale_factor = self._get_scale_factor() + lrs = [] + for cycle_min_lr, cycle_max_lr in zip(self.min_lrs, self.max_lrs): + base_height = (cycle_max_lr - cycle_min_lr) * scale_factor + lr = cycle_min_lr + base_height + lrs.append(lr) + + return lrs + + def _get_decay_mom(self, decay_batch_iteration): + if self.skip_mom_decay: + return self.max_moms + + decay_interval = decay_batch_iteration / self.decay_step_size + mom_decay_factor = (1 + self.decay_mom_rate * decay_interval) + momentums = [(beta0 * mom_decay_factor, beta1) for beta0, beta1 in self.max_moms] + + return momentums + + def _get_decay_lr(self, decay_batch_iteration): + """Calculates the learning rate at batch index. This function is used + after the cycle completes and post cycle decaying of lr/mom is enabled. + This function treats `self.last_batch_iteration` as the last batch index. + """ + if self.skip_lr_decay: + return self.min_lrs + + decay_interval = decay_batch_iteration / self.decay_step_size + lr_decay_factor = (1 + self.decay_lr_rate * decay_interval) + lrs = [cycle_min_lr / lr_decay_factor for cycle_min_lr in self.min_lrs] + + return lrs + + def get_lr(self): + """Calculates the learning rate at batch index. This function treats + `self.last_batch_iteration` as the last batch index. + """ + if self.last_batch_iteration < self.total_size: + return self._get_cycle_lr() + return self._get_decay_lr(self.last_batch_iteration - self.total_size + 1) + + def get_mom(self): + """Calculates the momentum at batch index. This function treats + `self.last_batch_iteration` as the last batch index. + """ + if not self.cycle_momentum: + return None + + if self.last_batch_iteration < self.total_size: + return self._get_cycle_mom() + return self._get_decay_mom(self.last_batch_iteration - self.total_size + 1) + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def step(self, batch_iteration=None): + """ Updates the optimizer with the learning rate for the last batch index. + `self.last_batch_iteration` is treated as the last batch index. + + If self.cycle_momentum is true, also updates optimizer momentum. + """ + if batch_iteration is None: + batch_iteration = self.last_batch_iteration + 1 + + self.last_batch_iteration = batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + if self.cycle_momentum: + momentums = self.get_mom() + for param_group, momentum in zip(self.optimizer.param_groups, momentums): + param_group['betas'] = momentum + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + +class WarmupLR(object): + """Increase the learning rate of each parameter group from min lr to max lr + over warmup_num_steps steps, and then fix at max lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + warmup_min_lr (float or list): minimum learning rate. Default: 0 + warmup_max_lr (float or list): maximum learning rate. Default: 0.001 + warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 + warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log + last_batch_iteration (int): The index of the last batch. Default: -1. + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = WarmupLR(optimizer) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + """ + + def __init__(self, + optimizer: Optimizer, + warmup_min_lr: float = 0.0, + warmup_max_lr: float = None, + warmup_num_steps: int = 1000, + warmup_type: str = WARMUP_LOG_RATE, + last_batch_iteration: int = -1): + + self.optimizer = get_torch_optimizer(optimizer) + + if warmup_max_lr is None: + warmup_max_lr = [group['lr'] for group in self.optimizer.param_groups][0] + + self.min_lrs = self._format_param(self.optimizer, warmup_min_lr, "min_lr") + self.max_lrs = self._format_param(self.optimizer, warmup_max_lr, "max_lr") + self.delta_lrs = [big - small for big, small in zip(self.max_lrs, self.min_lrs)] + self.warmup_num_steps = max(2, warmup_num_steps) + # Currently only support linear and log function + if warmup_type not in {WARMUP_LOG_RATE, WARMUP_LINEAR_RATE}: + logger.warning(f"Using unknown warmup_type: {warmup_type}. The increasing function " + f"is set to default (log)") + warmup_type = WARMUP_LOG_RATE + self.warmup_type = warmup_type + self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps) + self.last_batch_iteration = last_batch_iteration + # Initialize lr in optimizer + if last_batch_iteration == -1: + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def get_lr(self): + if self.last_batch_iteration < 0: + logger.warning("Attempting to get learning rate from scheduler before it has started") + return self.min_lrs + gamma = self._get_gamma() + return [min_lr + (delta_lr * gamma) for min_lr, delta_lr in zip(self.min_lrs, self.delta_lrs)] + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def step(self, last_batch_iteration=None): + if last_batch_iteration is None: + last_batch_iteration = self.last_batch_iteration + 1 + self.last_batch_iteration = last_batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + def _get_gamma(self): + if self.last_batch_iteration < self.warmup_num_steps: + if self.warmup_type == WARMUP_LOG_RATE: + return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) + elif self.warmup_type == WARMUP_LINEAR_RATE: + return self.last_batch_iteration / self.warmup_num_steps + return 1.0 + + def _format_param(self, optimizer, param_value, param_name): + if isinstance(param_value, list) or isinstance(param_value, tuple): + if len(param_value) != len(optimizer.param_groups): + raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name, + FileNotFoundError(param_value))) + return list(param_value) + return [param_value] * len(optimizer.param_groups) + + +class WarmupDecayLR(WarmupLR): + """Increase the learning rate of each parameter group from min lr to max lr + over warmup_num_steps steps, and then decay at linear rate over the remaining training steps. + + Args: + optimizer (Optimizer): Wrapped optimizer. + total_num_steps (int): total number of training steps + warmup_min_lr (float or list): minimum learning rate. Default: 0 + warmup_max_lr (float or list): maximum learning rate. Default: 0.001 + warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 + warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log + last_batch_iteration (int): The index of the last batch. Default: -1. + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = WarmupDecayLR(optimizer, 1000000) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + """ + + def __init__(self, + optimizer: Optimizer, + total_num_steps: int, + warmup_min_lr: float = 0.0, + warmup_max_lr: float = 0.001, + warmup_num_steps: int = 1000, + warmup_type: str = WARMUP_LOG_RATE, + last_batch_iteration: int = -1): + + self.total_num_steps = total_num_steps + super(WarmupDecayLR, self).__init__(optimizer, warmup_min_lr, warmup_max_lr, warmup_num_steps, warmup_type, + last_batch_iteration) + if self.total_num_steps < self.warmup_num_steps: + logger.warning('total_num_steps {} is less than warmup_num_steps {}'.format( + total_num_steps, warmup_num_steps)) + + def _get_gamma(self): + if self.last_batch_iteration < self.warmup_num_steps: + if self.warmup_type == WARMUP_LOG_RATE: + return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) + elif self.warmup_type == WARMUP_LINEAR_RATE: + return self.last_batch_iteration / self.warmup_num_steps + return max( + 0.0, + float(self.total_num_steps - self.last_batch_iteration) / + float(max(1.0, self.total_num_steps - self.warmup_num_steps))) + + +class WarmupCosineLR(object): + """Increase the learning rate of each parameter group from min lr ratio to max lr ratio + over warmup_num_steps steps, and then decay at cosine rate over the remaining training steps to min cosine ratio. + + Args: + optimizer (Optimizer): Wrapped optimizer. + total_num_steps (int): total number of training steps + warmup_min_ratio (float or list): warmup start learning rate ratio. Default: 0 + warmup_num_steps (int): number of steps to warm up from warmup_min_ratio to 1.0. Default: 1000 + warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log + cos_min_ratio (float): cosine end learning rate ratio. Default: 0.0001 + last_batch_iteration (int): The index of the last batch. Default: -1. + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = WarmupCosineLR(optimizer, 1000000) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + """ + + def __init__(self, + optimizer: Optimizer, + total_num_steps: int, + warmup_min_ratio: float = 0.0, + warmup_num_steps: int = 1000, + cos_min_ratio: float = 0.0001, + warmup_type: str = WARMUP_LOG_RATE, + last_batch_iteration: int = -1): + + self.optimizer = get_torch_optimizer(optimizer) + + self.total_num_steps = total_num_steps + self.last_batch_iteration = last_batch_iteration + self.cos_min_ratio = cos_min_ratio + + self.warmup_type = warmup_type + self.warmup_min_ratio = warmup_min_ratio + self.warmup_num_steps = max(2, warmup_num_steps) + self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps) + + if self.total_num_steps < self.warmup_num_steps: + logger.warning('total_num_steps {} is less than warmup_num_steps {}'.format( + total_num_steps, warmup_num_steps)) + self.org_lrs = [group['lr'] for group in self.optimizer.param_groups] + + # Initialize lrs in optimizer groups + if last_batch_iteration == -1: + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def get_lr_ratio(self): + if self.last_batch_iteration < 0: + logger.warning("Attempting to get learning rate from scheduler before it has started") + return [0.0] + + if self.last_batch_iteration < self.warmup_num_steps: + if self.warmup_type == WARMUP_LOG_RATE: + ratio = self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) + elif self.warmup_type == WARMUP_LINEAR_RATE: + ratio = self.last_batch_iteration / self.warmup_num_steps + ratio_delta = 1. - self.warmup_min_ratio + ratio = self.warmup_min_ratio + ratio * ratio_delta + return ratio + + real_last_step = self.last_batch_iteration - self.warmup_num_steps + 1 + real_total_steps = self.total_num_steps - self.warmup_num_steps + ratio_delta = 1. - self.cos_min_ratio + ratio = (1 + math.cos(math.pi * real_last_step / real_total_steps)) / 2 + ratio = max(0.0, self.cos_min_ratio + ratio_delta * ratio) + return ratio + + def step(self, last_batch_iteration=None): + if last_batch_iteration is None: + last_batch_iteration = self.last_batch_iteration + 1 + self.last_batch_iteration = last_batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def get_lr(self): + if self.last_batch_iteration < 0: + logger.warning("Attempting to get learning rate from scheduler before it has started") + return [0.0] + lr_ratio = self.get_lr_ratio() + return [org_lr * lr_ratio for org_lr in self.org_lrs] + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + def _format_param(self, optimizer, param_value, param_name): + if isinstance(param_value, list) or isinstance(param_value, tuple): + if len(param_value) != len(optimizer.param_groups): + raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name, + FileNotFoundError(param_value))) + return list(param_value) + return [param_value] * len(optimizer.param_groups) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5e60b03ac67112149c351045c3d9e117633d2b3f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .constants import * +from .writer_factory import CheckpointWriterFactory diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b07dfb2fcc29406002c779623bb3db131e1b46f2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a86f5f95921e1b13802640a5e515e5cdd1a6a83 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbe7b3206b90c8081a08d76dacce8a8c7cc73209 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4141c05accd43bc22d1f3e5e4dc5514c4edf2fd5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b62b0a0ccb79a671444198bf15f38aa788d24290 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62aebef8d3d924280c4d1b7347e1f983d1c32d3a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/config.py new file mode 100644 index 0000000000000000000000000000000000000000..d5a579fe31f09668e59647b433e427447df01367 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/config.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param +from .constants import * + +VALID_VALUES = { + CHECKPOINT_TAG_VALIDATION: CHECKPOINT_TAG_VALIDATION_MODES, + CHECKPOINT_WRITER_TYPE: CHECKPOINT_WRITER_TYPES, + CHECKPOINT_DATA_PARALLEL: CHECKPOINT_DATA_PARALLEL_UNITS +} + +CHECKPOINT_DEFAULT_DICT = { + CHECKPOINT_TAG_VALIDATION: CHECKPOINT_TAG_VALIDATION_DEFAULT, + CHECKPOINT_SERIALIZATION: CHECKPOINT_SERIALIZATION_DEFAULT, + CHECKPOINT_WRITER: CHECKPOINT_WRITER_DEFAULT +} + + +def _validate_config_values(config_name, config_dict, valid_values): + for key, value in config_dict.items(): + if value is None: + continue + if key in valid_values.keys(): + assert value in valid_values[key], \ + f"{config_name} contains invalid value {value} for {key}, expecting one of {valid_values[key]}" + + +def _make_upper_case(value): + return value if value is None else value.upper() + + +def get_checkpoint_writer_config(param_dict): + writer_dict = param_dict.get(CHECKPOINT_WRITER, None) + if writer_dict is None: + return CHECKPOINT_WRITER_DEFAULT + + writer_config = { + CHECKPOINT_WRITER_TYPE: + _make_upper_case(get_scalar_param(writer_dict, CHECKPOINT_WRITER_TYPE, CHECKPOINT_WRITER_TYPE_DEFAULT)), + CHECKPOINT_IO_BUFFER_SIZE: + get_scalar_param(writer_dict, CHECKPOINT_IO_BUFFER_SIZE, CHECKPOINT_IO_BUFFER_SIZE_DEFAULT), + CHECKPOINT_IO_BUFFER_DOUBLE: + get_scalar_param(writer_dict, CHECKPOINT_IO_BUFFER_DOUBLE, CHECKPOINT_IO_BUFFER_DOUBLE_DEFAULT), + CHECKPOINT_IO_STATISTICS: + get_scalar_param(writer_dict, CHECKPOINT_IO_STATISTICS, CHECKPOINT_IO_STATISTICS_DEFAULT), + CHECKPOINT_DATA_PARALLEL: + _make_upper_case(get_scalar_param(writer_dict, CHECKPOINT_DATA_PARALLEL, CHECKPOINT_DATA_PARALLEL_DEFAULT)), + CHECKPOINT_WRITER_DECOUPLED: + get_scalar_param(writer_dict, CHECKPOINT_WRITER_DECOUPLED, CHECKPOINT_WRITER_DECOUPLED_DEFAULT), + CHECKPOINT_IO_MULTIPLIER: + get_scalar_param(writer_dict, CHECKPOINT_IO_MULTIPLIER, CHECKPOINT_IO_MULTIPLIER_DEFAULT), + } + _validate_config_values(CHECKPOINT_WRITER, writer_config, VALID_VALUES) + + return writer_config + + +def get_checkpoint_config(param_dict): + checkpoint_dict = param_dict.get(CHECKPOINT, None) + if checkpoint_dict is None: + return CHECKPOINT_DEFAULT_DICT + + checkpoint_config = { + CHECKPOINT_TAG_VALIDATION: + get_scalar_param(checkpoint_dict, CHECKPOINT_TAG_VALIDATION, CHECKPOINT_TAG_VALIDATION_DEFAULT).upper(), + CHECKPOINT_SERIALIZATION: + get_scalar_param(checkpoint_dict, CHECKPOINT_SERIALIZATION, CHECKPOINT_SERIALIZATION_DEFAULT), + CHECKPOINT_WRITER: + get_checkpoint_writer_config(checkpoint_dict) + } + + _validate_config_values(CHECKPOINT, checkpoint_config, VALID_VALUES) + + return checkpoint_config diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..3b9bd549af92b1458e65d0b0702da659e6d2fd04 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/constants.py @@ -0,0 +1,85 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + + +######################################### +# Validation modes +######################################### +class ValidationMode: + WARN = "WARN" + IGNORE = "IGNORE" + FAIL = "FAIL" + + +######################################### +# Checkpoint config params +######################################### +# "checkpoint": {tag_validation=["Ignore"|"Warn"|"Fail"]} +CHECKPOINT_FORMAT = ''' +"checkpoint": { + "tag_validation": [Ignore|Warn|Fail], + "checkpoint_serialization": False, + "writer": { + "type": [mock|python|fast], + "decoupled": [True|False] + "io_buffer_size": 64e6, + "io_buffer_double": True, + "show_statistics": False, + "data_parallel": [replica|socket|machine], + "io_multiplier": 1, + } +} +''' +CHECKPOINT = "checkpoint" +CHECKPOINT_TAG_VALIDATION = "tag_validation" +CHECKPOINT_TAG_VALIDATION_DEFAULT = ValidationMode.WARN +CHECKPOINT_TAG_VALIDATION_MODES = [ValidationMode.WARN, ValidationMode.IGNORE, ValidationMode.FAIL] + +CHECKPOINT_SERIALIZATION = "checkpoint_serialization" +CHECKPOINT_SERIALIZATION_DEFAULT = True + +CHECKPOINT_WRITER = "writer" +CHECKPOINT_WRITER_DEFAULT = None + +CHECKPOINT_WRITER_TYPE = "type" + + +class CheckpointWriterType: + MOCK = "MOCK" + PYTHON = "PYTHON" + FAST = "FAST" + + +CHECKPOINT_WRITER_TYPE_DEFAULT = CheckpointWriterType.FAST +CHECKPOINT_WRITER_TYPES = [CheckpointWriterType.MOCK, CheckpointWriterType.PYTHON, CheckpointWriterType.FAST] + +CHECKPOINT_IO_BUFFER_SIZE = "io_buffer_size" +CHECKPOINT_IO_BUFFER_SIZE_DEFAULT = 64 * (1024**2) + +CHECKPOINT_IO_BUFFER_DOUBLE = "io_buffer_double" +CHECKPOINT_IO_BUFFER_DOUBLE_DEFAULT = True + +CHECKPOINT_IO_MULTIPLIER = "io_multiplier" +CHECKPOINT_IO_MULTIPLIER_DEFAULT = 1 + +CHECKPOINT_IO_STATISTICS = "show_statistics" +CHECKPOINT_IO_STATISTICS_DEFAULT = False + +CHECKPOINT_DATA_PARALLEL = "data_parallel" +CHECKPOINT_DATA_PARALLEL_DEFAULT = None + + +class CheckpointDataParallel: + REPLICA = "REPLICA" + SOCKET = "SOCKET" + MACHINE = "MACHINE" + + +CHECKPOINT_DATA_PARALLEL_UNITS = [ + CheckpointDataParallel.REPLICA, CheckpointDataParallel.SOCKET, CheckpointDataParallel.MACHINE +] + +CHECKPOINT_WRITER_DECOUPLED = "decoupled" +CHECKPOINT_WRITER_DECOUPLED_DEFAULT = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..15ed235e47066c139e0db662c2c718b436103434 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py @@ -0,0 +1,216 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass +from deepspeed.checkpoint.reshape_utils import partition_data +from deepspeed.runtime.zero.config import ZeroStageEnum +from .constants import * + + +@dataclass +class DataParallelWriterConfig(object): + world_size: int + rank: int + global_rank: int + local_rank: int + pure_dp: bool + + +class DataParallelWriterFactory(object): + + def __init__(self, uni_parallel_info, parallel_unit): + self._uni_parallel_info = uni_parallel_info + self._parallel_unit = parallel_unit + if parallel_unit == CheckpointDataParallel.SOCKET: + self._num_resources = uni_parallel_info.num_sockets + else: + self._num_resources = uni_parallel_info.num_machines + self._ranks_per_resource = max(1, self._uni_parallel_info.global_world_size // self._num_resources) + + def create_config(self, zero_stage, has_moe_layers): + if zero_stage == ZeroStageEnum.weights: + return self._create_config(1, 0) + + if has_moe_layers: + writer_config = self._get_expert_data_parallel_config() + else: + writer_config = self._get_data_parallel_config() + + if writer_config is None and zero_stage >= ZeroStageEnum.optimizer_states: + return self._create_config(1, 0) + + return writer_config + + def _create_config(self, world_size, rank): + return DataParallelWriterConfig(world_size=world_size, + rank=rank, + global_rank=self._uni_parallel_info.global_rank, + local_rank=self._uni_parallel_info.local_rank, + pure_dp=self._uni_parallel_info.pure_dp) + + def _get_expert_data_parallel_config(self): + ep_info = self._uni_parallel_info.ep_info + if self._parallel_unit is None: + dp_rank = ep_info.dp_rank + return self._create_config(1, 0) if dp_rank == 0 else None + + assert self._uni_parallel_info.pure_dp, \ + f'3D parallelism is not yet supported for data parallel checkpointing.' + + if self._parallel_unit == CheckpointDataParallel.REPLICA or ep_info.ep_world_size == 1: + return self._get_parallel_write_for_ddp(ep_info.dp_world_size, ep_info.dp_rank) + + return self._get_expert_parallel_write_for_2d() + + def _get_expert_parallel_write_for_2d(self): + ep_info = self._uni_parallel_info.ep_info + + def _get_expert_slice_resources(expert_resources, resource_name): + ep_world_size = ep_info.ep_world_size + slices_per_resource = min(self._ranks_per_resource, ep_world_size) + assert slices_per_resource <= len(expert_resources) + + ep_num_resources = len(expert_resources) + assert ep_num_resources % slices_per_resource == 0, f'{resource_name}: Expected ep_num_resources={ep_num_resources} to multiple of slices_per_resource={slices_per_resource} for ep_world_size={ep_world_size}' + + slice_partitions = partition_data(expert_resources, slices_per_resource) + # print( + # f'edp_resource_partition: self._uni_parallel_info.global_rank={self._uni_parallel_info.global_rank} expert_resources={expert_resources} slices_per_resource={slices_per_resource} ep_world_size={ep_world_size} slice_partitions={slice_partitions}' + # ) + resource_index = ep_info.ep_rank % slice_resources + return slice_partitions[resource_index] + + dp_ranks = ep_info.dp_peer_ranks + expert_resources = [r // self._ranks_per_resource for r in dp_ranks] + slice_resources = _get_expert_slice_resources(expert_resources, self._parallel_unit) + assert all([idx < self._num_resources for idx in expert_resources]), \ + f'Detected invalid resource index in expert_resources={expert_resources}, self._num_resources={self._num_resources}' + return self._assign_resources_to_tensor_slice(slice_resources, ep_info.ep_rank, dp_ranks) + + def _get_data_parallel_config(self): + mpu_info = self._uni_parallel_info.mpu_info + if self._parallel_unit is None: + dp_rank = self._uni_parallel_info.dp_rank if mpu_info is None else mpu_info.dp_rank + return self._create_config(1, 0) if dp_rank == 0 else None + + if self._uni_parallel_info.pure_dp: + return self._get_parallel_write_for_ddp(self._uni_parallel_info.global_world_size, + self._uni_parallel_info.global_rank) + + if self._parallel_unit == CheckpointDataParallel.REPLICA: + return self._create_config(mpu_info.dp_world_size, mpu_info.dp_rank) + + return self._get_parallel_write_for_3d() + + def _get_parallel_write_for_3d(self): + mpu_info = self._uni_parallel_info.mpu_info + my_global_rank = self._uni_parallel_info.global_rank + + def _expand_resources(resource_list, new_size): + old_size = len(resource_list) + if old_size >= new_size: + return resource_list + + assert new_size % old_size == 0, f'Expect new_size={new_size} to be multiple of old_size={old_size}' + multiplier = new_size // old_size + new_resource_list = [] + for r in resource_list: + new_resource_list += [r] * multiplier + # print(f'expand_resources: {my_global_rank=} {old_size=} {new_size=} {resource_list=} {new_resource_list=}') + return new_resource_list + + # Getting resource partition for a tensor slice is a 2-step process + # 1. Get resource partitions for all pipeline stages. A pipeline stage is a 2D grid of size TP x DP + def _get_pipeline_stage_resources(resource_indices): + num_resources = len(resource_indices) + pp_world_size = mpu_info.pp_world_size + if num_resources < pp_world_size: + resource_indices = _expand_resources(resource_indices, pp_world_size) + num_resources = pp_world_size + global_resource_partitions = partition_data(resource_indices, pp_world_size) + pp_rank = mpu_info.pp_rank + return global_resource_partitions[pp_rank] + + # 2. Get resource partition for tensor slice. A tensor slice is a 1D vector of size DP + def _get_tensor_slice_resources(resource_indices, resource_name): + pipe_stage_resources = _get_pipeline_stage_resources(resource_indices) + tp_world_size = mpu_info.tp_world_size + if len(pipe_stage_resources) < tp_world_size: + pipe_stage_resources = _expand_resources(pipe_stage_resources, tp_world_size) + tp_num_resources = len(pipe_stage_resources) + assert tp_num_resources % tp_world_size == 0, \ + f'{resource_name}: Expected tp_num_resources={tp_num_resources} to multiple of tp_world_size={tp_world_size}' + + pipe_stage_resource_partitions = partition_data(pipe_stage_resources, tp_world_size) + tp_rank = mpu_info.tp_rank + return pipe_stage_resource_partitions[tp_rank] + + def _get_model_parallel_slice_resources(): + # Get resources of my dp peer ranks + resources = [(r // self._ranks_per_resource) for r in mpu_info.dp_peer_ranks] + if len(resources) < self._ranks_per_resource: + resources = _expand_resources(resources, self._ranks_per_resource) + + resource_partitions = partition_data(resources, self._ranks_per_resource) + mp_rank = (mpu_info.pp_rank * mpu_info.tp_world_size) + mpu_info.tp_rank + slice_rank = mp_rank % self._ranks_per_resource + return resource_partitions[slice_rank] + + num_slices = mpu_info.tp_world_size * mpu_info.pp_world_size + if num_slices > self._ranks_per_resource: + slice_resources = _get_model_parallel_slice_resources() + else: + all_resources = list(range(self._num_resources)) + slice_resources = _get_tensor_slice_resources(all_resources, self._parallel_unit) + + return self._assign_resources_to_tensor_slice(slice_resources, mpu_info.tp_rank, mpu_info.dp_peer_ranks) + + def _get_slice_writers(self, slice_resources, my_dp_ranks): + resource_map = {} + for res in slice_resources: + resource_map[res] = [r for r in my_dp_ranks if (r // self._ranks_per_resource) == res] + + # Only one writer per resource, and we conventionally pick the first rank as writer. + return [ranks[0] for ranks in resource_map.values()] + + def _assign_resources_to_tensor_slice(self, slice_resources, my_slice_index, my_dp_ranks): + my_global_rank = self._uni_parallel_info.global_rank + slice_writer_ranks = self._get_slice_writers(slice_resources, my_dp_ranks) + my_resource_index = my_global_rank // self._ranks_per_resource + print( + f'resource_assign: my_global_rank={my_global_rank} my_slice_index={my_slice_index} my_dp_ranks={my_dp_ranks} slice_resources={slice_resources} slice_writer_ranks={slice_writer_ranks}' + ) + if my_resource_index in slice_resources and my_global_rank in slice_writer_ranks: + my_writer_index = (my_global_rank - slice_writer_ranks[0]) // self._ranks_per_resource + num_slice_writers = len(slice_writer_ranks) + print( + f'slice_writer: my_global_rank={my_global_rank} my_writer_index={my_writer_index} num_slice_writers={num_slice_writers}' + ) + return self._create_config(num_slice_writers, my_writer_index) + + return None + + def _get_parallel_write_for_ddp(self, dp_world_size, dp_rank): + if self._parallel_unit == CheckpointDataParallel.REPLICA: + return self._create_config(dp_world_size, dp_rank) + + num_machines = self._uni_parallel_info.num_machines + if self._parallel_unit == CheckpointDataParallel.SOCKET: + if dp_world_size == num_machines: + # There is one rank per machine + return self._create_config(num_machines, dp_rank) + + num_sockets = self._uni_parallel_info.num_sockets + ranks_per_socket = dp_world_size // num_sockets + if dp_rank % ranks_per_socket == 0: + return self._create_config(num_sockets, dp_rank // ranks_per_socket) + else: + return None + + ranks_per_machine = dp_world_size // num_machines + if dp_rank % ranks_per_machine == 0: + return self._create_config(num_machines, self._uni_parallel_info.machine_rank) + + return None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e212008a9277f1804217a8aed7936004f25ee5fb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/utils.py @@ -0,0 +1,84 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from dataclasses import dataclass +from deepspeed import comm as dist +from deepspeed.constants import CROSS_RANK, CROSS_SIZE, LOCAL_RANK +from .data_parallel_writer_factory import DataParallelWriterFactory + +# TODO: parse socket number from env. +SOCKETS_PER_MACHINE = 2 + + +@dataclass +class MPUInfo(object): + pp_world_size: int + pp_rank: int + tp_world_size: int + tp_rank: int + dp_world_size: int + dp_peer_ranks: list + dp_rank: int + + +def _create_model_parallel_info(mpu): + return MPUInfo(pp_world_size=mpu.get_pipeline_model_parallel_world_size(), + pp_rank=mpu.get_pipeline_model_parallel_rank(), + tp_world_size=mpu.get_tensor_model_parallel_world_size(), + tp_rank=mpu.get_tensor_model_parallel_rank(), + dp_world_size=mpu.get_data_parallel_world_size(), + dp_peer_ranks=mpu.get_data_parallel_group_ranks(), + dp_rank=mpu.get_data_parallel_rank()) + + +@dataclass +class ExpertParallelInfo(object): + ep_world_size: int + ep_rank: int + dp_world_size: int + dp_peer_ranks: list + dp_rank: int + + +def _create_expert_parallel_info(groups): + group_name = groups._get_max_expert_size_name() + return ExpertParallelInfo(ep_world_size=groups._get_expert_parallel_world_size(group_name), + ep_rank=groups._get_expert_parallel_rank(group_name), + dp_world_size=groups._get_expert_data_parallel_world_size(group_name), + dp_peer_ranks=groups._get_expert_data_parallel_group_ranks(group_name), + dp_rank=groups._get_expert_data_parallel_rank(group_name)) + + +@dataclass +class UniversalParallelInfo(object): + global_world_size: int + global_rank: int + local_rank: int + mpu_info: MPUInfo + ep_info: ExpertParallelInfo + pure_dp: bool + num_machines: int + machine_rank: int + num_sockets: int + + +def create_universal_parallel_info(groups, has_moe_layers): + return UniversalParallelInfo(global_world_size=dist.get_world_size(), + global_rank=dist.get_rank(), + local_rank=int(os.environ[LOCAL_RANK]), + mpu_info=None if groups.mpu is None else _create_model_parallel_info(groups.mpu), + ep_info=_create_expert_parallel_info(groups) if has_moe_layers else None, + pure_dp=groups.mpu is None + or groups.mpu.get_data_parallel_world_size() == dist.get_world_size(), + num_machines=int(os.environ[CROSS_SIZE]), + machine_rank=int(os.environ[CROSS_RANK]), + num_sockets=int(os.environ[CROSS_SIZE]) * SOCKETS_PER_MACHINE) + + +def create_data_parallel_writer_config(groups, parallel_unit, zero_stage, has_moe_layers): + uni_parallel_info = create_universal_parallel_info(groups, has_moe_layers) + writer_factory = DataParallelWriterFactory(uni_parallel_info, parallel_unit) + return writer_factory.create_config(zero_stage, has_moe_layers) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/writer_factory.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/writer_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..a8c324530ae52468219489ff2c9480288f2d547f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/writer_factory.py @@ -0,0 +1,95 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.ops.op_builder import AsyncIOBuilder, GDSBuilder +from deepspeed.io import MockFileWriter, PyFileWriter, FastFileWriter, FastFileWriterConfig +from deepspeed.runtime.swap_tensor.constants import * +from .constants import * +from deepspeed.accelerator import get_accelerator + + +class CheckpointWriterFactory(object): + + def __init__(self, writer_config, aio_config, dp_writer_config): + self._type = writer_config[CHECKPOINT_WRITER_TYPE] + self._io_buffer_size = writer_config[CHECKPOINT_IO_BUFFER_SIZE] + self._io_buffer_double = writer_config[CHECKPOINT_IO_BUFFER_DOUBLE] + self._data_parallel_writer = dp_writer_config + self._io_multiplier = writer_config[CHECKPOINT_IO_MULTIPLIER] + if self._data_parallel_writer.pure_dp: + self._show_statistics = writer_config[CHECKPOINT_IO_STATISTICS] and self._data_parallel_writer is not None + else: + self._show_statistics = writer_config[CHECKPOINT_IO_STATISTICS] and self._data_parallel_writer is not None + self._io_buffer = None + self._dnvme_handle = None + self._writer = None + self._use_gds = False + + if self._type == CheckpointWriterType.FAST: + self._use_gds = aio_config[AIO_USE_GDS] + if self._use_gds: + self._setup_for_gds(aio_config) + else: + self._setup_for_aio(aio_config) + print( + f'WriterFactory: self._data_parallel_writer={self._data_parallel_writer} self._show_statistics={self._show_statistics}' + ) + + def create_writer(self, file_path, optimize_dp_state): + assert self._writer is None, \ + f'Cannot create checkpoint writer for {file_path} because writer is currently used for {self._writer.file_path()}.\ + Must call writer.release() before reusing to avoid this error.' + + if self._type == CheckpointWriterType.MOCK: + self._writer = MockFileWriter(file_path) + elif self._type == CheckpointWriterType.PYTHON: + self._writer = PyFileWriter(file_path) + else: + if optimize_dp_state: + num_parallel_writers = self._data_parallel_writer.world_size * self._io_multiplier + writer_rank = self._data_parallel_writer.rank + file_path = f'{file_path}-{writer_rank}.{num_parallel_writers}' + # print(f'create_dp_writer: {self._data_parallel_writer.global_rank=} {writer_rank=} {num_parallel_writers=} {file_path=}') + else: + num_parallel_writers = 1 + writer_rank = 0 + # print(f'create_rank0_writer: {self._data_parallel_writer.global_rank=} {writer_rank=} {num_parallel_writers=} {file_path=}') + + config = FastFileWriterConfig(dnvme_handle=self._dnvme_handle, + pinned_tensor=self._io_buffer, + double_buffer=self._io_buffer_double, + num_parallel_writers=num_parallel_writers, + writer_rank=writer_rank, + global_rank=self._data_parallel_writer.global_rank) + self._writer = FastFileWriter(file_path=file_path, config=config) + + return self._writer + + def release_writer(self): + self._writer.close() + if self._show_statistics: + self._writer._dump_state() + self._writer = None + + def _setup_for_aio(self, aio_config): + self._io_buffer = torch.zeros(self._io_buffer_size, dtype=torch.uint8, device='cpu').pin_memory() + self._dnvme_handle = AsyncIOBuilder().load().aio_handle( + block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + def _setup_for_gds(self, aio_config): + self._io_buffer = torch.zeros(self._io_buffer_size, + dtype=torch.uint8, + device=get_accelerator().current_device_name()) + self._dnvme_handle = GDSBuilder().load().gds_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + self._dnvme_handle.pin_device_tensor(self._io_buffer) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cf58eb1801af7cc02f3cc2057a5c843cd078b575 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .module import PipelineModule, LayerSpec, TiedLayerSpec +from .topology import ProcessTopology diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de46cb67cfba49ca168ccf1763a60ce6f93ada8b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86ab8e38986732b58303dead4124039016424930 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..387ffcfdc213ae71153764368f28466d17df861e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff65b9693087591111f1e1c9008522e0388e1247 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d950a390fc006d7f30a41945296cb438ec43b0f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60203c98b07edda9f58f655dc5544aa36bf064f3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..6ef00748bf1f3b74991173c390410761c460c73f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/engine.py @@ -0,0 +1,1426 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from types import MethodType +from collections import OrderedDict +from functools import reduce +from operator import mul + +import torch +from deepspeed import comm as dist + +from deepspeed.utils import logger +from deepspeed.utils.timer import ThroughputTimer +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.bf16_optimizer import BF16_Optimizer + +from ..engine import DeepSpeedEngine, MEMORY_OPT_ALLREDUCE_SIZE +from deepspeed.utils.timer import FORWARD_MICRO_TIMER, FORWARD_GLOBAL_TIMER, BACKWARD_MICRO_TIMER, \ + BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_INNER_GLOBAL_TIMER, \ + BACKWARD_REDUCE_MICRO_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ + STEP_MICRO_TIMER, STEP_GLOBAL_TIMER + +from ..utils import PartitionedTensor +from ..dataloader import RepeatingLoader +from ..zero.config import ZeroStageEnum +from ..activation_checkpointing import checkpointing as ds_checkpointing + +from .module import PipelineModule, PipelineError +from . import p2p +from . import schedule + +TARGET_ID = -2 +LOG_STAGE = -2 +DATA_PARALLEL_ID = -2 + +BATCH_INPUT_TIMER = 'batch_input' +TRAIN_BATCH_TIMER = 'train_batch' +PIPE_SEND_OUTPUT_TIMER = 'pipe_send_output' +PIPE_SEND_GRAD_TIMER = 'pipe_send_grad' +PIPE_RECV_INPUT_TIMER = 'pipe_recv_input' +PIPE_RECV_GRAD_TIMER = 'pipe_recv_grad' + +# The buffer size to store the meta data for each tensor. +TENSOR_META_SIZE = 256 + + +def is_even(number): + return number % 2 == 0 + + +mem_alloced = 0 +mem_cached = 0 + + +def _tensor_bytes(tensor): + return tensor.numel() * tensor.element_size() + + +class PipelineEngine(DeepSpeedEngine): + """ A training engine hybrid pipeline, data, and model parallel training. + + This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule` + is provided. + """ + ID_TO_DTYPE = [ + torch.float32, torch.float64, torch.complex64, torch.complex128, torch.float16, torch.bfloat16, torch.uint8, + torch.int8, torch.int16, torch.int32, torch.int64, torch.bool + ] + DTYPE_TO_ID = {dtype: id_ for id_, dtype in enumerate(ID_TO_DTYPE)} + + def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs): + super().__init__(*super_args, **super_kwargs) + assert isinstance(self.module, PipelineModule), "model must base PipelineModule" + + assert self.zero_optimization_stage( + ) < ZeroStageEnum.gradients, "ZeRO-2 and ZeRO-3 are incompatible with pipeline parallelism" + + # We schedule the all-reduces, so disable it in super().backward() + self.enable_backward_allreduce = False + self.has_bool_tensors = has_bool_tensors + self.eval_return_logits = False + self.outputs = None + # BF16 Optimizer is hardcoded for fp32 gradient accumulation + self.using_bf16_optimizer = type(self.optimizer) == BF16_Optimizer + + # used to disable the pipeline all-reduce when used with 1-bit Adam/1-bit LAMB + self.pipeline_enable_backward_allreduce = True + + if self.elasticity_enabled(): + if not self.is_elastic_model_parallel_supported(): + assert not self.elasticity_enabled(), "Elasticity is not currently supported" \ + " with pipeline parallelism." + + # pipeline step for logging + self.log_batch_step_id = -1 + + self.micro_batch_size = self.train_micro_batch_size_per_gpu() + self.micro_batches = self.gradient_accumulation_steps() + + # Set Grid and Communication Groups + self.grid = self.module._grid + if self.grid.get_global_rank() == 0: + logger.info(f'CONFIG: micro_batches={self.micro_batches} ' + f'micro_batch_size={self.micro_batch_size}') + + self.global_rank = self.grid.get_global_rank() + + assert self.dp_world_size == self.grid.data_parallel_size + assert self.train_batch_size() == \ + self.micro_batch_size * self.micro_batches * self.grid.data_parallel_size + + # Set Stage Inf + self.num_stages = self.grid.pipe_parallel_size + self.stage_id = self.grid.get_stage_id() + self.prev_stage = self.stage_id - 1 + self.next_stage = self.stage_id + 1 + + self.data_iterator = None + self.batch_fn = None + + self._force_grad_boundary = False + + self.batch_timer = ThroughputTimer(self._config.timers_config, + batch_size=self.train_batch_size(), + logging_fn=self.tput_log, + monitor_memory=False, + steps_per_output=self.steps_per_print()) + + # PipelineEngine needs to handle data loading specially due to only the first + # and last stages loading inputs/labels. We construct a sampler that uses + if self.training_data: + self._build_data_iter(self.training_data) + + self.is_pipe_parallel = self.grid.pipe_parallel_size > 1 + self.is_data_parallel = self.grid.data_parallel_size > 1 + self.is_model_parallel = self.grid.model_parallel_size > 1 + + # Partition input/output buffers + # XXX temporarily disable while I revert some partition hacks. + assert isinstance(self._config.pipeline['pipe_partitioned'], bool) + assert isinstance(self._config.pipeline['grad_partitioned'], bool) + self.is_pipe_partitioned = self.is_model_parallel and self._config.pipeline['pipe_partitioned'] + self.is_grad_partitioned = self.is_model_parallel and self._config.pipeline['grad_partitioned'] + logger.info(f'is_pipe_partitioned= {self.is_pipe_partitioned} ' + f'is_grad_partitioned= {self.is_grad_partitioned}') + + model_parameters = filter(lambda p: p.requires_grad, self.module.parameters()) + num_params = sum([p.numel() for p in model_parameters]) + unique_params = num_params + # Subtract tied parameters if we don't own them + if self.module.tied_comms: + tied_params = 0 + for key, d in self.module.tied_comms.items(): + if self.global_rank != min(d['ranks']): + tied_params += sum(p.numel() for p in d['module'].parameters()) + unique_params -= tied_params + params_tensor = torch.LongTensor(data=[num_params, unique_params]).to(self.device) + dist.all_reduce(params_tensor, group=self.grid.get_model_parallel_group()) + params_tensor = params_tensor.tolist() + total_params = params_tensor[0] + unique_params = params_tensor[1] + if self.grid.data_parallel_id == 0: + logger.info(f'RANK={self.global_rank} ' + f'STAGE={self.stage_id} ' + f'LAYERS={self.module._local_stop - self.module._local_start} ' + f'[{self.module._local_start}, {self.module._local_stop}) ' + f'STAGE_PARAMS={num_params} ({num_params/1e6:0.3f}M) ' + f'TOTAL_PARAMS={total_params} ({total_params/1e6:0.3f}M) ' + f'UNIQUE_PARAMS={unique_params} ({unique_params/1e6:0.3f}M)') + + #initialize peer-2-peer communication and allreduce groups + if self.is_pipe_parallel: + p2p.init_process_groups(self.grid) + + # Pipeline buffers + self.num_pipe_buffers = 0 + self.pipe_buffers = { + 'inputs': [], # batch input and received activations + 'labels': [], # labels from batch input + 'outputs': [], # activations + 'output_tensors': [], # tensor object to preserve backward graph + } + self.pipe_recv_buf = None + self.grad_layer = None + self._grad_layer_buf = [] + + self.meta_buffer = None + + self.first_output_send = True + self.first_gradient_send = True + self.pipe_partition_input_meta_cache = None + self.pipe_partition_output_meta_cache = None + self.pipe_partition_grad_meta_cache = None + self.grad_partition_grad_layer_meta_cache = None + + #stores the loss for the current micro batch being processed + self.loss = torch.tensor(0.0).to(self.device) + + #stores the loss for the entire batch + self.total_loss = None + self.total_additional_losses = None + self.agg_loss = torch.tensor(0.0, requires_grad=False).to(self.device) + self.dp_group_loss = torch.tensor(0.0, requires_grad=False).to(self.device) + + # stores aggregated-DP train final loss and aggregated-DP additional losses, if any + # additional losses are stored as dict: {loss-name: agg-loss} + self.agg_train_loss = None + self.agg_additional_losses = None + + if self._config.pipeline['activation_checkpoint_interval'] > 0: + self.module.activation_checkpoint_interval = self._config.pipeline['activation_checkpoint_interval'] + # set use_reentrant default to True. + if self._config.pipeline.get('use_reentrant') is None: + self._config.pipeline['use_reentrant'] = True + if self._config.pipeline['use_reentrant'] is False: + # set activation_checkpoint_func to non_reentrant_checkpoint func. + self.module.activation_checkpoint_func = ds_checkpointing.non_reentrant_checkpoint + if self.grid.get_global_rank() == 0: + logger.info(f'CONFIG: activation_checkpoint_func=non_reentrant_checkpoint') + if self.module.activation_checkpoint_interval > 0: + self.module._precompute_checkpointable_values() + + self.module.checkpoint_parallel_write_pipeline = self._config.checkpoint_parallel_write_pipeline + + if self.is_last_stage(): + self.loss_model = self.module.loss_fn + + self.has_attention_mask = self.module.__class__.__name__ == 'GPT2ModelPipe' + # Initialize pipeline communicators. Just send a 0. + if is_even(self.stage_id): + if not self.is_last_stage(): + p2p.send(self.loss, self.next_stage) + if not self.is_first_stage(): + p2p.recv(self.loss, self.prev_stage) + else: + if not self.is_first_stage(): + p2p.recv(self.loss, self.prev_stage) + if not self.is_last_stage(): + p2p.send(self.loss, self.next_stage) + + # XXX look into timer reporting timing + # Initialize some timers because of early weirdness. + if self.wall_clock_breakdown(): + self.timers(FORWARD_MICRO_TIMER).start() + self.timers(FORWARD_MICRO_TIMER).stop() + self.timers(BACKWARD_MICRO_TIMER).start() + self.timers(BACKWARD_MICRO_TIMER).stop() + self.timers(BACKWARD_INNER_MICRO_TIMER).start() + self.timers(BACKWARD_INNER_MICRO_TIMER).stop() + self.timers(BACKWARD_REDUCE_MICRO_TIMER).start() + self.timers(BACKWARD_REDUCE_MICRO_TIMER).stop() + self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).start() + self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).stop() + self.timers(STEP_MICRO_TIMER).start() + self.timers(STEP_MICRO_TIMER).stop() + + self.dynamic_shape = self.module.dynamic_shape + + def set_has_attention_mask(self, value): + assert isinstance(value, bool) + self.has_attention_mask = value + + def _build_data_iter(self, dataset): + sampler = torch.utils.data.distributed.DistributedSampler(dataset, + num_replicas=self.dp_world_size, + rank=self.mpu.get_data_parallel_rank(), + shuffle=False) + # Build a loader and make it repeating. + pipe_dataloader = self.deepspeed_io(dataset, data_sampler=sampler) + pipe_dataloader = RepeatingLoader(pipe_dataloader) + self.set_dataloader(pipe_dataloader) + + def _exec_reduce_tied_grads(self): + # We need to run this first to write to self.averaged_gradients; + # since this class turns `enable_backward_allreduce` off, + # `self.overlapping_partition_gradients_reduce_epilogue()` defined in the DeepSpeedEngine + # never actually runs. I suspect this is because of efficiency problems; get_flat_partition in + # stage2.py might do something expensive; someone will have to look into that later. But + # in the meantime, this fixes ZeRO2 + Pipelining enough to run a demo. Further profiling + # needed to decide if it actually breaks everything. + # (see https://github.com/EleutherAI/gpt-neox/issues/62#issuecomment-761471944) + if self.zero_optimization_partition_gradients(): + self.optimizer.overlapping_partition_gradients_reduce_epilogue() + + weight_group_list = self.module.get_tied_weights_and_groups() + for weight, group in weight_group_list: + grad = weight._hp_grad if self.using_bf16_optimizer else weight.grad + if grad is not None: + dist.all_reduce(grad, group=group) + + def _exec_reduce_grads(self): + self._force_grad_boundary = True + if self.pipeline_enable_backward_allreduce: + if self.using_bf16_optimizer: + # PP+BF16 work for ZeRO Stage 1 + self._bf16_reduce_grads() + else: + self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE) + self._force_grad_boundary = False + + def _bf16_reduce_grads(self): + self.buffered_allreduce_fallback(grads=None, elements_per_buffer=MEMORY_OPT_ALLREDUCE_SIZE) + + def _reserve_pipe_buffers(self, num_buffers): + """Ensure that each pipeline buffer has at least ``num_buffers`` slots. + + This method only reserves slots and does not allocate tensors. + + Args: + num_buffers (int): The number of buffers to reserve. + """ + if self.num_pipe_buffers >= num_buffers: + return + + num_added = num_buffers - self.num_pipe_buffers + for key in self.pipe_buffers: + self.pipe_buffers[key].extend([None] * num_added) + self.num_pipe_buffers = num_buffers + + def reset_activation_shape(self): + """Reset the buffers when the shape of activation and gradient change. + For example, for curriculum learning that changes the seqlen of each + sample, we need to call this whenever the seqlen is going to change. + """ + self.first_output_send = True + self.pipe_recv_buf = None + self.grad_layer = None + self._grad_layer_buf = [] + self.meta_buffer = None + + self.pipe_partition_input_meta_cache = None + self.pipe_partition_output_meta_cache = None + self.pipe_partition_grad_meta_cache = None + self.grad_partition_grad_layer_meta_cache = None + + def train_batch(self, data_iter=None): + """Progress the pipeline to train the next batch of data. The engine will ingest + ``self.train_batch_size()`` total samples collectively across all workers. + + + An iterator that over training data should be provided as an argument + unless ``deepspeed.initialize()`` was provided a training set. In that event, + the training data will automatically be read. + + + .. warning:: + A total of ``self.gradient_accumulation_steps()`` entries will be pulled + from ``data_iter`` by each pipeline. There must be sufficient + data left in ``data_iter`` or else a ``StopIteration`` will halt training. + + DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` + that wraps data loaders to automatically restart upon a ``StopIteration``. + + Args: + data_iter (Iterator, optional): Iterator of training data. + + Returns: + The arithmetic mean of the losses computed this batch. + """ + if not torch._C.is_grad_enabled(): + raise RuntimeError(f'train_batch() requires gradients enabled. Use eval_batch() instead.') + + # Curriculum learning could change activation shape + if self.curriculum_enabled_legacy(): + new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ + self.global_steps + 1) + if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: + self.reset_activation_shape() + self.curriculum_scheduler_legacy.first_step = False + elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ + self.global_steps): + self.reset_activation_shape() + + if data_iter is not None: + self.set_dataiterator(data_iter) + + self.module.train() + self.total_loss = None + self.total_additional_losses = None + self._compute_loss = True + + # Do the work + self.timers(TRAIN_BATCH_TIMER).start() + sched = schedule.TrainSchedule(micro_batches=self.micro_batches, + stages=self.num_stages, + stage_id=self.stage_id) + self._exec_schedule(sched) + + with torch.no_grad(): + self.agg_train_loss = self._aggregate_total_loss() + + self.timers(TRAIN_BATCH_TIMER).stop() + + if self.steps_per_print() is not None and self.global_steps % self.steps_per_print() == 0: + if self.global_rank == 0: + elapsed = self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True) / 1000.0 + iter_time = elapsed / self.steps_per_print() + tput = self.train_batch_size() / iter_time + log_str = f'steps: {self.global_steps} loss: {self.agg_train_loss:0.4f} ' + if self.agg_additional_losses is not None: + for loss_name, loss_value in self.agg_additional_losses.items(): + log_str += f'{loss_name}: {loss_value.item():0.4f} ' + log_str += f'iter time (s): {iter_time:0.3f} samples/sec: {tput:0.3f}' + print(log_str) + else: + self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True) + + # Monitoring + if self.global_rank == 0 and self.monitor.enabled: + self.summary_events = [(f'Train/Samples/train_loss', self.agg_train_loss.mean().item(), + self.global_samples)] + self.monitor.write_events(self.summary_events) + + if self.steps_per_print() is not None and self.wall_clock_breakdown( + ) and self.global_steps % self.steps_per_print() == 0: + self.timers.log([ + PIPE_SEND_OUTPUT_TIMER, + PIPE_SEND_GRAD_TIMER, + PIPE_RECV_INPUT_TIMER, + PIPE_RECV_GRAD_TIMER, + ]) + + # TODO: should return precisely what loss returned and allow others to be queried? + return self.agg_train_loss + + def eval_batch(self, + data_iter, + return_logits=False, + compute_loss=True, + reduce_output='avg', + bcast_loss=True, + num_micro_batches=None): + """Evaluate the pipeline on a batch of data from ``data_iter``. The + engine will evaluate ``self.train_batch_size()`` total samples + collectively across all workers. + + This method is equivalent to: + + .. code-block:: python + + module.eval() + with torch.no_grad(): + output = module(batch) + + .. warning:: + A total of ``self.gradient_accumulation_steps()`` entries will be pulled + from ``data_iter`` by each pipeline. There must be sufficient + data left in ``data_iter`` or else a ``StopIteration`` will halt training. + + DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` + that wraps data loaders to automatically restart upon a ``StopIteration``. + + Args: + data_iter (Iterator): Iterator of data to evaluate. + + Returns: + The arithmetic mean of the losses computed this batch. + """ + self.eval_return_logits = return_logits + self.module.eval() + + # Curriculum learning could change activation shape + if self.curriculum_enabled_legacy(): + new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ + self.global_steps + 1) + if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: + self.reset_activation_shape() + self.curriculum_scheduler_legacy.first_step = False + elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ + self.global_steps): + self.reset_activation_shape() + + eval_output = None + + self._compute_loss = compute_loss + + # Use the provided data iterator + train_iterator = self.data_iterator + self.set_dataiterator(data_iter) + + # set the number micro batches in case the user chose value than training + micro_batches = self.micro_batches if num_micro_batches is None else num_micro_batches + + # Do the work + sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=self.num_stages, stage_id=self.stage_id) + + # prevent dead-lock with multiple evals sequence + dist.barrier() + + with torch.no_grad(): + self._exec_schedule(sched) + + if self.is_last_stage(): + eval_output = self._reduce_outputs(self.fwd_outputs, reduce=reduce_output, micro_batches=micro_batches) + + if compute_loss and (bcast_loss or self.monitor.enabled): + eval_output = self._bcast_pipe_scalar(eval_output) + + if self.global_rank == 0 and self.monitor.enabled: + self.summary_events = [(f'Train/Samples/eval_loss', eval_output.mean().item(), self.global_samples)] + self.monitor.write_events(self.summary_events) + + # Restore the training iterator + self.set_dataiterator(train_iterator) + + # Reset any buffers that may have been populated during the forward passes. + #ds_checkpointing.reset() + self.eval_return_logits = False + if return_logits: + outputs = self.outputs + self.outputs = None + return eval_output, outputs + return eval_output + + def set_train_batch_size(self, train_batch_size): + """Adjust the global batch size by increasing or decreasing the number of + micro-batches (i.e., gradient accumulation steps). The size of each micro-batch + (i.e., ``train_micro_batch_size_per_gpu``) is not changed. + Args: + train_batch_size (int): The new global batch size for training. + Raises: + ValueError: if ``train_batch_size`` is not divisible by the + configured micro-batch size and data parallelism. + """ + super().set_train_batch_size(train_batch_size) + self.micro_batches = self.gradient_accumulation_steps() + + def is_first_stage(self): + """True if this process is in the first stage in the pipeline.""" + return self.stage_id == 0 + + def is_last_stage(self): + """True if this process is in the last stage in the pipeline.""" + return self.stage_id == self.num_stages - 1 + + def _reduce_outputs(self, outputs, reduce='avg', reduce_dp=True, micro_batches=None): + if reduce is None: + return outputs + + if reduce.lower() == 'avg': + # first sum over all microbatches + if torch.is_tensor(outputs[0]): + reduced = sum(outputs) + else: + assert isinstance(outputs, (list, tuple)) + reduced = [torch.zeros_like(o) for o in outputs[0]] + for idx, out in outputs: + reduced[idx] += out + + # Average over the microbatches + reduced = self._scale_loss_by_gas(reduced, eval_micro_batches=micro_batches) + + # Average over DP groups + if reduce_dp and self.is_data_parallel: + if torch.is_tensor(reduced): + dist.all_reduce(reduced, group=self.mpu.get_data_parallel_group()) + reduced /= self.dp_world_size + else: + for idx in range(len(reduced)): + dist.all_reduce(reduced[idx], group=self.mpu.get_data_parallel_group()) + reduced[idx] /= self.dp_world_size + + return reduced + else: + raise NotImplementedError(f'reduction type {reduce} not supported.') + + def _bcast_pipe_scalar(self, data, src_rank=None, dtype=torch.float32): + # Default to last stage (e.g., for broadcasting loss) + if src_rank is None: + src_rank = self.grid.stage_to_global(self.num_stages - 1) + assert src_rank in self.grid.pp_group + + if self.global_rank == src_rank: + result = data.clone().detach().type(dtype).to(self.device) + else: + result = torch.Tensor([0.]).type(dtype).to(self.device) + + dist.broadcast(tensor=result, src=src_rank, group=self.mpu.get_pipe_parallel_group()) + + return result + + def _aggregate_total_loss(self): + # Scale loss, average among DP ranks, and bcast loss to the rest of my DP group + if self.is_last_stage(): + # Scale loss and additional losses, if any + loss = self._scale_loss_by_gas(self.total_loss) + self.agg_additional_losses = self.total_additional_losses + if self.agg_additional_losses is not None: + self.agg_additional_losses = OrderedDict({ + loss_name: self._scale_loss_by_gas(_loss.clone().detach()) + for loss_name, _loss in self.agg_additional_losses.items() + }) + + self.dp_group_loss = loss.clone().detach() + agg_loss = self.dp_group_loss.clone().detach() + #print(f'RANK={self.global_rank} bcast SENDER src={self.global_rank} group={self.grid.pp_group}', flush=True) + + # Average loss across all data-parallel groups + if self.is_data_parallel: + if self.agg_additional_losses is None: + dist.all_reduce(agg_loss, group=self.mpu.get_data_parallel_group()) + agg_loss /= self.dp_world_size + else: + # use a single reduce op for agg_loss and additional losses, if any + assert '__train_loss__' not in self.agg_additional_losses.keys() + tensors = OrderedDict({'__train_loss__': agg_loss}) + tensors.update(self.agg_additional_losses.items()) + flat_tensor = torch.cat([t.clone().reshape(-1).detach() for t in tensors.values()]) + dist.all_reduce(flat_tensor, group=self.mpu.get_data_parallel_group()) + flat_tensor /= self.dp_world_size + offset = 0 + reduced_tensor = {} + for name, t in tensors.items(): + n_elem = t.numel() + reduced_tensor[name] = flat_tensor[offset:offset + n_elem].clone().detach().reshape(t.shape) + offset += n_elem + agg_loss = reduced_tensor['__train_loss__'] + self.agg_additional_losses = OrderedDict( + {name: reduced_tensor[name] + for name in self.agg_additional_losses.keys()}) + + assert self.global_rank in self.grid.pp_group + losses = [self.dp_group_loss, agg_loss] + if self.agg_additional_losses is not None: + losses += list(self.agg_additional_losses.values()) + losses = torch.stack(losses).float() + if self.is_pipe_parallel: + dist.broadcast(tensor=losses, src=self.global_rank, group=self.mpu.get_pipe_parallel_group()) + else: + # Get loss from last stage + src_rank = self.grid.stage_to_global(self.num_stages - 1) + assert src_rank in self.grid.pp_group + # losses to reduce are: dp_group_loss, agg_loss, model additional losses + # therefore: 2 + n_additional_losses + additional_losses = self.module.get_additional_losses() + n_additional_losses = 0 if additional_losses is None else len(additional_losses) + losses = torch.Tensor([0.] * (2 + n_additional_losses)).to(self.device) + dist.broadcast(tensor=losses, src=src_rank, group=self.grid.get_pipe_parallel_group()) + self.dp_group_loss = losses[0].clone().detach() + agg_loss = losses[1].clone().detach() + if additional_losses is not None: + self.agg_additional_losses = OrderedDict({ + name: losses[2 + i].clone().detach() + for i, name in enumerate(additional_losses.keys()) + }) + return agg_loss + + def set_dataloader(self, loader): + """""" + if self.is_first_stage() or self.is_last_stage(): + self.training_dataloader = loader + self.data_iterator = iter(self.training_dataloader) + + def set_dataiterator(self, iterator): + """ Store an iterator to sample for training data. """ + if self.is_first_stage() or self.is_last_stage(): + self.training_dataloader = None + self.data_iterator = iterator + + def set_batch_fn(self, fn): + """Execute a post-processing function on input data. + + Args: + fn (function): The function to run. + """ + self.batch_fn = fn + + def is_gradient_accumulation_boundary(self): + """True if the engine is executing a gradient reduction or optimizer step instruction. + + This is overridden from :class:`DeepSpeedEngine` to force reductions + and steps when the pipeline engine is instructed to do so. + + Returns: + bool: whether reductions and optimizer steps should occur. + """ + return self._force_grad_boundary + + def log_for_device(self, *msg): + if LOG_STAGE == self.stage_id or LOG_STAGE == -1: + if DATA_PARALLEL_ID == self.grid.data_parallel_id or DATA_PARALLEL_ID == -1: + print( + f'RANK={dist.get_rank()} ' + f'PIPE-ID={self.stage_id} ' + f'DATA-ID={self.grid.data_parallel_id} ' + f'MBATCH-ID={self.microbatch_id} ' + f'STEP-ID={self.log_batch_step_id} ' + '::', + *msg, + flush=True) + + def tput_log(self, *msg): + if self.global_rank == 0 and self.global_steps % self.steps_per_print() == 0: + print(*msg) + + def _next_batch(self): + # If using 3D parallelism, only some first-stage ranks may do IO + batch = None + if self.data_iterator is not None: + batch = next(self.data_iterator) + + # Any post-processing, like broadcasting across a slice-parallel group. + if self.batch_fn: + batch = self.batch_fn(batch) + + return batch + + def _exec_forward_pass(self, buffer_id): + self.tput_timer.start() + self.mem_status('BEFORE FWD', reset_max=True) + + if isinstance(self.pipe_buffers['inputs'][buffer_id], tuple): + inputs = tuple(t.clone() for t in self.pipe_buffers['inputs'][buffer_id]) + else: + inputs = self.pipe_buffers['inputs'][buffer_id].clone() + + # collect the partitioned input from the previous stage + if self.is_pipe_partitioned and not self.is_first_stage(): + if self.pipe_partition_input_meta_cache is None: + self.pipe_partition_input_meta_cache = inputs[0].to('cpu') + part_input = PartitionedTensor.from_meta(meta=self.pipe_partition_input_meta_cache, + local_part=inputs[1], + group=self.grid.get_slice_parallel_group()) + + inputs = (part_input.full(), *inputs[2:]) + inputs[0].requires_grad = True + # skip mask + #inputs[1].requires_grad = True + part_input = None + inputs = inputs[0] if len(inputs) == 1 else inputs + self.pipe_buffers['inputs'][buffer_id] = inputs + + # inputs has no gradient because it is from a cloned tensor + outputs = super().forward(inputs) + + # Reset activation checkpointing buffers. + # Need to call this between evaluation iterations + if not self.module.training: + ds_checkpointing.reset() + + # Partition the outputs if we are not the last stage + if self.is_pipe_partitioned and not self.is_last_stage(): + if isinstance(outputs, tuple): + first_output = outputs[0] + # TODO: Improve pipe partitioning to pass multiple tensors that require grads + assert all([torch.is_tensor(elt) and elt.requires_grad is False for elt in outputs[1:]]) + outputs_tail = outputs[1:] + elif torch.is_tensor(outputs): + first_output = outputs + outputs_tail = [] + else: + raise ValueError("expecting a tensor or a tuple of tensors") + part = PartitionedTensor(tensor=first_output, group=self.grid.get_slice_parallel_group()) + # Clear the large output data, but save the computation graph + first_output.data = torch.zeros(1, device=first_output.data.device) + self.pipe_buffers['output_tensors'][buffer_id] = first_output + # Inject the partitioned tensor into the output before sending + outputs = (part.to_meta(), part.data(), *outputs_tail) + part = None + + self.pipe_buffers['outputs'][buffer_id] = outputs + + # Optionally compute loss on the last device + if self.is_last_stage(): + if self._compute_loss and self.module.loss_fn is not None: + labels = self.pipe_buffers['labels'][buffer_id] + self.loss = self.module.loss_fn(outputs, labels) + else: + # Some models just return loss from forward() + self.loss = outputs + if self.eval_return_logits: + self.outputs = outputs + + if isinstance(self.loss, torch.Tensor): + self.fwd_outputs.append(self.loss.detach()) + else: + self.fwd_outputs.append([l.detach() for l in self.loss]) + + def add_to_total_loss(_total_loss, _loss): + if isinstance(_loss, torch.Tensor): + if _total_loss is None: + _total_loss = torch.zeros_like(_loss) + _total_loss += _loss.detach() + else: + if _total_loss is None: + _total_loss = [torch.zeros_like(_l) for _l in _loss] + for _idx, _l in enumerate(_loss): + _total_loss[_idx] += _l.detach() + return _total_loss + + self.total_loss = add_to_total_loss(self.total_loss, self.loss) + + # aggregate additional losses across gradient accumulation steps + additional_losses = self.module.get_additional_losses() + if additional_losses is not None: + if self.total_additional_losses is None: + self.total_additional_losses = OrderedDict() + for name, loss in additional_losses.items(): + total = self.total_additional_losses[name] if name in self.total_additional_losses else None + self.total_additional_losses[name] = add_to_total_loss(total, loss) + + def _exec_backward_pass(self, buffer_id): + assert self.optimizer is not None, "must provide optimizer during " \ + "init in order to use backward" + + self.mem_status('BEFORE BWD', reset_max=True) + + # The last stage just runs backward on the loss using DeepSpeed's typical + # mechanisms. + if self.is_last_stage(): + super().backward(self.loss) + self.mem_status('AFTER BWD') + return + + outputs = self.pipe_buffers['outputs'][buffer_id] + + if self.wall_clock_breakdown(): + self.timers(BACKWARD_MICRO_TIMER).start() + self.timers(BACKWARD_GLOBAL_TIMER).start() + self.timers(BACKWARD_INNER_MICRO_TIMER).start() + self.timers(BACKWARD_INNER_GLOBAL_TIMER).start() + + # Reconstruct if we previously partitioned the output. We must be + # careful to also restore the computational graph of the tensors we partitioned. + if self.is_pipe_partitioned: + if self.is_grad_partitioned: + if self.pipe_partition_output_meta_cache is None: + self.pipe_partition_output_meta_cache = outputs[0].to('cpu') + part_output = PartitionedTensor.from_meta(meta=self.pipe_partition_output_meta_cache, + local_part=outputs[1], + group=self.grid.get_slice_parallel_group()) + self.pipe_buffers['output_tensors'][buffer_id].data = part_output.full() + outputs = (self.pipe_buffers['output_tensors'][buffer_id], *outputs[2:]) + else: + # Already restored from partition + self.pipe_buffers['output_tensors'][buffer_id].data = outputs[0] + outputs = (self.pipe_buffers['output_tensors'][buffer_id], *outputs[1:]) + + grad_tensors = self.grad_layer + if self.is_grad_partitioned: + #print(f'RANK={self.global_rank} BEFORE-BWD restoring grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') + if self.grad_partition_grad_layer_meta_cache is None: + self.grad_partition_grad_layer_meta_cache = self.grad_layer[0].to('cpu') + part_grad = PartitionedTensor.from_meta(meta=self.grad_partition_grad_layer_meta_cache, + local_part=self.grad_layer[1], + group=self.grid.get_slice_parallel_group()) + grad_tensors = (part_grad.full(), *grad_tensors[2:]) + part_grad = None + #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') + + if self.using_bf16_optimizer and not self.is_last_stage(): + # manually call because we don't call optimizer.backward() + self.optimizer.clear_lp_grads() + + # This handles either a single tensor or tuple of tensors. + if isinstance(outputs, tuple): + out_tensors = [t for t in outputs if t.is_floating_point()] + assert len(out_tensors) == len(grad_tensors) + torch.autograd.backward(tensors=out_tensors, grad_tensors=grad_tensors) + else: + torch.autograd.backward(tensors=(outputs, ), grad_tensors=(grad_tensors, )) + + if self.using_bf16_optimizer and not self.is_last_stage(): + # manually call because we don't call optimizer.backward() + if not self._config.bfloat16_config.immediate_grad_update: + self.optimizer.update_hp_grads(clear_lp_grads=False) + + # Free up the memory from the output of forward() + self.pipe_buffers['output_tensors'][buffer_id] = None + self.pipe_buffers['outputs'][buffer_id] = None + grad_tensors = None + + if self.wall_clock_breakdown(): + self.timers(BACKWARD_INNER_MICRO_TIMER).stop() + self.timers(BACKWARD_INNER_GLOBAL_TIMER).stop() + self.timers(BACKWARD_MICRO_TIMER).stop() + self.timers(BACKWARD_GLOBAL_TIMER).stop() + + self.mem_status('AFTER BWD') + + def _exec_load_micro_batch(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(BATCH_INPUT_TIMER).start() + + batch = self._next_batch() + + if self.is_first_stage(): + loaded = None + if torch.is_tensor(batch[0]): + loaded = batch[0].clone().to(self.device).detach() + if self._config.pipeline['activation_checkpoint_interval'] > 0 and self._config.pipeline[ + 'use_reentrant']: + loaded.requires_grad = loaded.is_floating_point() + else: + assert isinstance(batch[0], (tuple, list)) + # Assume list or tuple + loaded = [] + for x in batch[0]: + assert torch.is_tensor(x) + mine = x.clone().detach().to(self.device) + if self._config.pipeline['activation_checkpoint_interval'] > 0 and self._config.pipeline[ + 'use_reentrant']: + mine.requires_grad = mine.is_floating_point() + loaded.append(mine) + loaded = tuple(loaded) + + self.pipe_buffers['inputs'][buffer_id] = loaded + + if self.is_last_stage(): + loaded = batch[1] + if torch.is_tensor(batch[1]): + loaded = batch[1].to(self.device) + # XXX: torch 1.6.0 DataLoader will auto convert tuple to list + elif isinstance(batch[1], (tuple, list)): + loaded = [] + for x in batch[1]: + assert torch.is_tensor(x) + x = x.to(self.device).detach() + loaded.append(x) + loaded = tuple(loaded) + + self.pipe_buffers['labels'][buffer_id] = loaded + + if self.wall_clock_breakdown(): + self.timers(BATCH_INPUT_TIMER).stop() + + def _send_tensor_meta(self, buffer, recv_stage): + """ Communicate metadata about upcoming p2p transfers. + + Metadata is communicated in this order: + * type (0: tensor, 1: list) + * num_tensors if type=list + foreach tensor in buffer: + * ndims + * shape + """ + meta_buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device) + if isinstance(buffer, torch.Tensor): + meta_buf_list = [ + 0, # type of data (0: tensor, 1: list (unused), 2: tuple) + self.DTYPE_TO_ID[buffer.dtype], # dtype + len(buffer.size()) # ndims + ] + meta_buf_list.extend(buffer.size()) + assert len( + meta_buf_list + ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}" + meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32)) + p2p.send(meta_buffer, recv_stage) + + elif isinstance(buffer, tuple): + meta_buf_list = [ + 2, # type of data (0: tensor, 1: list (unused), 2: tuple) + len(buffer) # num_tensors + ] + + for tensor in buffer: + assert isinstance(tensor, torch.Tensor) + meta_buf_list.append(self.DTYPE_TO_ID[tensor.dtype]) + meta_buf_list.append(len(tensor.size())) + meta_buf_list.extend(tensor.size()) + + assert len( + meta_buf_list + ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}" + meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32)) + p2p.send(meta_buffer, recv_stage) + + else: + raise NotImplementedError(f'Could not send meta type {type(buffer)}') + + # Useful for performance debugging. + ''' + if self.grid.data_parallel_id == 0: + print(f'STAGE={self.stage_id} pipe-send-volume: {send_bytes/1024**2:0.2f}MB') + ''' + + def _recv_tensor_meta(self, send_stage): + """Receive metadata about upcoming p2p transfers and return allocated buffers. + + Returns: + Allocated buffer for receiving from send_stage. + """ + buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device) + p2p.recv(buffer, send_stage) + + recv_type = buffer[0].item() + + # A single tensor will be sent. + if recv_type == 0: + recv_dtype = self.ID_TO_DTYPE[buffer[1].item()] + recv_ndims = buffer[2].item() + recv_shape = buffer[3:3 + recv_ndims].tolist() + return self._allocate_or_extend_buffers(0, recv_shape, recv_dtype) + + # List or tuple of tensors (recv_type == 1 (list) is currently unused) + elif recv_type == 1 or recv_type == 2: + num_tensors = buffer[1].item() + + buffers = [] + offset = 2 + for idx in range(num_tensors): + recv_dtype = self.ID_TO_DTYPE[buffer[offset].item()] + recv_ndims = buffer[offset + 1].item() + recv_shape = buffer[offset + 2:offset + 2 + recv_ndims].tolist() + offset += 2 + recv_ndims + + buffers.append(self._allocate_or_extend_buffers(idx, recv_shape, recv_dtype)) + + # Convert to tuples if requested. + if recv_type == 2: + buffers = tuple(buffers) + return buffers + + else: + raise NotImplementedError(f'Could not receive type {type(recv_type)}') + + def _exec_send_activations(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_OUTPUT_TIMER).start() + + outputs = self.pipe_buffers['outputs'][buffer_id] + + # NCCL does not like to send torch.BoolTensor types, so cast the mask to half(). + # We could do char, but with half() we can eventually flatten with other fp16 + # messages (TODO) + if self.has_attention_mask or self.has_bool_tensors: + outputs = list(outputs) + outputs[-1] = outputs[-1].half() + outputs = tuple(outputs) + + if self.dynamic_shape or self.first_output_send: + self.first_output_send = False + self._send_tensor_meta(outputs, self.next_stage) + + if isinstance(outputs, torch.Tensor): + p2p.send(outputs, self.next_stage) + elif isinstance(outputs, tuple): + for idx, buffer in enumerate(outputs): + p2p.send(buffer, self.next_stage) + else: + raise NotImplementedError('Could not send output of type ' + f'{type(outputs)}') + + # Restore the boolean tensor + if self.has_attention_mask or self.has_bool_tensors: + outputs = list(outputs) + outputs[-1] = outputs[-1].bool() + outputs = tuple(outputs) + + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_OUTPUT_TIMER).stop() + + def _exec_send_grads(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_GRAD_TIMER).start() + + inputs = self.pipe_buffers['inputs'][buffer_id] + + # Partition the gradient + if self.is_grad_partitioned: + if isinstance(inputs, tuple): + first_input = inputs[0] + assert all([torch.is_tensor(elt) for elt in inputs[1:]]) + inputs_grad_tail = [elt.grad for elt in inputs[1:]] + elif torch.is_tensor(inputs): + first_input = inputs + inputs_grad_tail = [] + else: + raise ValueError("expecting a tensor or a tuple of tensors") + assert torch.is_tensor(first_input) + part = PartitionedTensor(tensor=first_input.grad, group=self.grid.get_slice_parallel_group()) + + inputs = (part.to_meta(), part.data(), *inputs_grad_tail) + + # XXX Terrible hack + # Drop the attention mask from the input buffer here. It does not have + # a grad that needs to be communicated. We free the buffer immediately + # after, so no need to restore it. The receiver also has a hack that skips + # the recv. This is because NCCL does not let us send torch.BoolTensor :-(. + if self.has_attention_mask or self.has_bool_tensors: + inputs = list(inputs) + inputs.pop() + inputs = tuple(inputs) + + if isinstance(inputs, torch.Tensor): + assert inputs.grad is not None + p2p.send(inputs.grad, self.prev_stage) + else: + # XXX terrible hacky branch + if self.is_grad_partitioned: + # First two sends are partitioned gradient + p2p.send(inputs[0], self.prev_stage) + p2p.send(inputs[1], self.prev_stage) + else: + for idx, buffer in enumerate(inputs): + # Skip tensors that will not produce a grad + if not buffer.is_floating_point(): + assert buffer.grad is None + continue + assert buffer.grad is not None + p2p.send(buffer.grad, self.prev_stage) + + # We can free up the input buffer now + self.pipe_buffers['inputs'][buffer_id] = None + + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_GRAD_TIMER).stop() + + def _exec_recv_activations(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_INPUT_TIMER).start() + + recvd = None + + # Allocate the buffer if necessary + if self.dynamic_shape or self.pipe_recv_buf is None: + self.pipe_recv_buf = self._recv_tensor_meta(self.prev_stage) + + if isinstance(self.pipe_recv_buf, torch.Tensor): + p2p.recv(self.pipe_recv_buf, self.prev_stage) + recvd = self.pipe_recv_buf.clone().detach() + recvd.requires_grad = recvd.is_floating_point() + else: + assert isinstance(self.pipe_recv_buf, tuple) + recvd = [None] * len(self.pipe_recv_buf) + for idx, buffer in enumerate(self.pipe_recv_buf): + assert torch.is_tensor(buffer) + # XXX hardcode meta type + if self.is_pipe_partitioned and idx == 0 and buffer.dtype != torch.long: + if self.meta_buffer is None: + self.meta_buffer = torch.zeros(buffer.size(), dtype=torch.long, device=self.device) + buffer = self.meta_buffer + + p2p.recv(buffer, self.prev_stage) + recvd[idx] = buffer.clone().detach() + + # NCCL does not like to send torch.BoolTensor types, so un-cast the + # attention mask + if self.has_attention_mask or self.has_bool_tensors: + recvd[-1] = recvd[-1].bool() + + recvd = tuple(recvd) + + for buffer in recvd: + buffer.requires_grad = buffer.is_floating_point() + + self.pipe_buffers['inputs'][buffer_id] = recvd + + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_INPUT_TIMER).stop() + + def _exec_recv_grads(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_GRAD_TIMER).start() + + outputs = self.pipe_buffers['outputs'][buffer_id] + # XXX these shapes are hardcoded for Megatron + # Restore partitioned output if it was partitioned and we are sending full gradients + if self.is_pipe_partitioned and not self.is_grad_partitioned: + if self.pipe_partition_grad_meta_cache is None: + self.pipe_partition_grad_meta_cache = outputs[0].to('cpu') + part_output = PartitionedTensor.from_meta(meta=self.pipe_partition_grad_meta_cache, + local_part=outputs[1], + group=self.grid.get_slice_parallel_group()) + outputs[0].data = part_output.full() + outputs = (outputs[0], *outputs[2:]) + # save for backward + self.pipe_buffers['outputs'][buffer_id] = outputs + + # Allocate gradient if necessary + if self.dynamic_shape or self.grad_layer is None: + if isinstance(outputs, torch.Tensor): + self.grad_layer = self._allocate_or_extend_buffers(0, list(outputs.size()), outputs.dtype) + else: + # XXX This is a HACK + # When we exchange activations/gradients, the two pipe stages + # need to issue the send/recv with the same buffer sizes or + # else there is a deadlock. The is_floating_point() filter is + # used to avoid sending gradients for tensors that do not + # produce gradients. When TP>1, we partition the first + # activations/gradients across TP ranks to save communication + # volume and memory. That partitioned tensor is represented as + # two tensors: a 1/TPth chunk of the original data and also a + # small LongTensor storing the metadata used to reconstruct on + # the other side. When combined, the floating point filter also + # filtered out the metadata tensor. This quick (hacky) fix just + # branches on is_grad_partitioned so we don't filter out the + # metadata tensor. + if self.is_grad_partitioned: + sizes_and_dtypes = [(list(t.size()), t.dtype) + for t in outputs[:2]] + [(list(t.size()), t.dtype) + for t in outputs[2:] if t.is_floating_point()] + else: + sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs if t.is_floating_point()] + + self.grad_layer = [ + self._allocate_or_extend_buffers(i, size, dtype) + for i, (size, dtype) in enumerate(sizes_and_dtypes) + ] + + if isinstance(self.grad_layer, torch.Tensor): + p2p.recv(self.grad_layer, self.next_stage) + else: + assert isinstance(outputs, tuple) + for idx, buffer in enumerate(self.grad_layer): + # XXX GPT-2 hack + if self.is_grad_partitioned and idx == 0 and buffer.dtype != torch.long: + buffer.data = torch.zeros(buffer.size(), dtype=torch.long, device=self.device) + p2p.recv(buffer, self.next_stage) + + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_GRAD_TIMER).stop() + + def _exec_optimizer_step(self, lr_kwargs=None): + if self.wall_clock_breakdown(): + self.timers(STEP_MICRO_TIMER).start() + self.timers(STEP_GLOBAL_TIMER).start() + self.mem_status('BEFORE STEP', reset_max=True) + + self._force_grad_boundary = True + self._take_model_step(lr_kwargs) + self._force_grad_boundary = False + + self.mem_status('AFTER STEP') + + if self.global_rank == 0 and self.monitor.enabled: + self.summary_events = [(f'Train/Samples/lr', self.get_lr()[0], self.global_samples)] + if self.fp16_enabled() and hasattr(self.optimizer, 'cur_scale'): + self.summary_events.append( + (f'Train/Samples/loss_scale', self.optimizer.cur_scale, self.global_samples)) + self.monitor.write_events(self.summary_events) + + if self.wall_clock_breakdown(): + self.timers(STEP_MICRO_TIMER).stop() + self.timers(STEP_GLOBAL_TIMER).stop() + if self.global_steps % self.steps_per_print() == 0: + self.timers.log([ + BATCH_INPUT_TIMER, + FORWARD_MICRO_TIMER, + BACKWARD_MICRO_TIMER, + BACKWARD_INNER_MICRO_TIMER, + BACKWARD_REDUCE_MICRO_TIMER, + STEP_MICRO_TIMER, + ]) + if self.global_steps % self.steps_per_print() == 0: + self.timers.log([ + FORWARD_GLOBAL_TIMER, + BACKWARD_GLOBAL_TIMER, + BACKWARD_INNER_GLOBAL_TIMER, + BACKWARD_REDUCE_GLOBAL_TIMER, + STEP_GLOBAL_TIMER, + ]) + + def _allocate_zeros(self, shape, **kwargs): + """ Allocate a tensor of zeros on the engine's device. + + Arguments: + shape: the shape of the tensor to allocate + kwargs: passed to torch.zeros() + + Returns: + A tensor from torch.zeros() allocated on self.device. + """ + if "dtype" not in kwargs: + if self.fp16_enabled(): + kwargs["dtype"] = torch.half + if self.bfloat16_enabled(): + kwargs["dtype"] = torch.bfloat16 + + return torch.zeros(shape, device=self.device, **kwargs) + + def _allocate_buffer(self, shape, num_buffers=-1, **kwargs): + buffers = [] + if num_buffers == -1: + num_buffers = self.num_pipe_buffers + for count in range(num_buffers): + buffers.append(self._allocate_zeros(shape, **kwargs)) + return buffers + + def _allocate_or_extend_buffers(self, idx, shape, dtype): + numel = reduce(mul, shape) if len(shape) > 0 else 1 + if len(self._grad_layer_buf) <= idx or self._grad_layer_buf[idx].numel() < numel: + new_buf = self._allocate_buffer(shape, dtype=dtype, num_buffers=1)[0] + if len(self._grad_layer_buf) <= idx: + self._grad_layer_buf.append(new_buf) + else: + self._grad_layer_buf[idx] = new_buf + return self._grad_layer_buf[idx] + else: + return self._grad_layer_buf[idx].flatten()[:numel].view(shape) + + def forward(self, *args, **kwargs): + """Disabled for pipeline parallel training. See ``train_batch()``. """ + raise PipelineError("Only train_batch() is accessible in pipeline mode.") + + def backward(self, *args, **kwargs): + """Disabled for pipeline parallel training. See ``train_batch()``. """ + raise PipelineError("Only train_batch() is accessible in pipeline mode.") + + def step(self, *args, **kwargs): + """Disabled for pipeline parallel training. See ``train_batch()``. """ + raise PipelineError("Only train_batch() is accessible in pipeline mode.") + + def mem_status(self, msg, print_rank=-1, reset_max=False): + return + global mem_alloced, mem_cached + if not self.global_steps == 0 or not self.global_steps == 9: + #return + pass + if self.mpu.get_data_parallel_rank() != 0: + return + + if self.global_rank != 0: + return + + rank = self.global_rank + if print_rank != -1 and rank != print_rank: + return + + get_accelerator().synchronize() + + if reset_max: + get_accelerator().reset_max_memory_cached() + get_accelerator().reset_max_memory_allocated() + + new_alloced = get_accelerator().memory_allocated() + new_cached = get_accelerator().memory_cached() + + delta_alloced = new_alloced - mem_alloced + delta_cached = new_cached - mem_cached + + mem_cached = new_cached + mem_alloced = new_alloced + + max_alloced = get_accelerator().max_memory_allocated() + max_cached = get_accelerator().max_memory_cached() + + # convert to GB for printing + new_alloced /= 1024**3 + new_cached /= 1024**3 + delta_alloced /= 1024**3 + delta_cached /= 1024**3 + max_alloced /= 1024**3 + max_cached /= 1024**3 + + print( + f'RANK={rank} STAGE={self.stage_id} STEP={self.global_steps} MEMSTATS', msg, + f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' + f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') + + def module_state_dict(self, exclude_frozen_parameters=False): + """Override hack to save a pipe model and return the directory path of the save. + + This method should only be called by DeepSpeed's ``save_checkpoint()``. The + recommended way of saving a ``PipelineModule`` outside of ``save_checkpoint()`` + is ``save_state_dict()``. + + Returns: + None + """ + assert isinstance(self.module, PipelineModule) + assert self._curr_ckpt_path is not None, \ + "PipelineEngine expects module_state_dict() to be called from save_checkpoint()" + + self.module.save_state_dict(self._curr_ckpt_path, + checkpoint_engine=self.checkpoint_engine, + exclude_frozen_params=exclude_frozen_parameters) + return None + + def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): + """Override hack to instead use a directory path. + + This is important because pipeline models checkpoint by layer instead of rank. + + If ``state_dict`` is not ``None`` or a ``str``, we revert to ``super()`` expecting a ``dict``. + + Args: + state_dict (str, None): unused + strict (bool, optional): Strict state loading. Defaults to True. + """ + assert custom_load_fn is None, "custom_load_fn not supported w. pipeline parallelism" + state_dict = checkpoint if self.has_moe_layers else checkpoint['module'] + if (state_dict is not None) and (not isinstance(state_dict, str)): + super().load_module_state_dict(state_dict, strict) + return + + self.module.load_state_dir(load_dir=self._curr_ckpt_path, + strict=strict, + checkpoint_engine=self.checkpoint_engine) + + # A map of PipeInstruction types to methods. Each method will be executed with the + # kwargs provided to the PipeInstruction from the scheduler. + _INSTRUCTION_MAP = { + schedule.OptimizerStep: _exec_optimizer_step, + schedule.ReduceGrads: _exec_reduce_grads, + schedule.ReduceTiedGrads: _exec_reduce_tied_grads, + schedule.LoadMicroBatch: _exec_load_micro_batch, + schedule.ForwardPass: _exec_forward_pass, + schedule.BackwardPass: _exec_backward_pass, + schedule.SendActivation: _exec_send_activations, + schedule.RecvActivation: _exec_recv_activations, + schedule.SendGrad: _exec_send_grads, + schedule.RecvGrad: _exec_recv_grads, + } + + def _exec_schedule(self, pipe_schedule): + # Reserve and reset buffers. + self._reserve_pipe_buffers(pipe_schedule.num_pipe_buffers()) + self.fwd_outputs = [] + + # For each step in the schedule + for step_cmds in pipe_schedule: + # For each instruction in the step + for cmd in step_cmds: + if type(cmd) not in self._INSTRUCTION_MAP: + raise RuntimeError(f'{self.__class__.__name__} does not understand instruction {repr(cmd)}') + + # Equivalent to: self._exec_forward_pass(buffer_id=0) + self._exec_instr = MethodType(self._INSTRUCTION_MAP[type(cmd)], self) + self._exec_instr(**cmd.kwargs) + + def get_additional_losses(self): + return self.agg_additional_losses diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/module.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/module.py new file mode 100644 index 0000000000000000000000000000000000000000..ad2463f0769f66fb597d80e82b1a811f80b12439 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/module.py @@ -0,0 +1,695 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import glob + +import re as regex + +from functools import partial + +import torch +import torch.nn as nn +from deepspeed import comm as dist + +from deepspeed.utils import logger +from .. import utils as ds_utils +from ..activation_checkpointing import checkpointing +from .topology import PipeDataParallelTopology, PipelineParallelGrid +from deepspeed.runtime.state_dict_factory import SDLoaderFactory +from deepspeed.accelerator import get_accelerator +from deepspeed.checkpoint.utils import clone_tensors_for_torch_save + + +class PipelineError(Exception): + """Errors related to the use of deepspeed.PipelineModule """ + + +class LayerSpec: + """Building block for specifying pipeline-parallel modules. + + LayerSpec stores the type information and parameters for each stage in a + PipelineModule. For example: + + .. code-block:: python + + nn.Sequence( + torch.nn.Linear(self.in_dim, self.hidden_dim, bias=False), + torch.nn.Linear(self.hidden_hidden, self.out_dim) + ) + + becomes + + .. code-block:: python + + layer_specs = [ + LayerSpec(torch.nn.Linear, self.in_dim, self.hidden_dim, bias=False), + LayerSpec(torch.nn.Linear, self.hidden_hidden, self.out_dim)] + ] + """ + + def __init__(self, typename, *module_args, **module_kwargs): + self.typename = typename + self.module_args = module_args + self.module_kwargs = module_kwargs + + if not issubclass(typename, nn.Module): + raise RuntimeError('LayerSpec only supports torch.nn.Module types.') + + if dist.is_initialized(): + self.global_rank = dist.get_rank() + else: + self.global_rank = -1 + + def __repr__(self): + return ds_utils.call_to_str(self.typename.__name__, self.module_args, self.module_kwargs) + + def build(self, log=False): + """Build the stored specification.""" + if log: + logger.info(f'RANK={self.global_rank} building {repr(self)}') + + return self.typename(*self.module_args, **self.module_kwargs) + + +class TiedLayerSpec(LayerSpec): + + def __init__(self, key, typename, *module_args, forward_fn=None, tied_weight_attr=['weight'], **module_kwargs): + super().__init__(typename, *module_args, **module_kwargs) + self.key = key + self.forward_fn = forward_fn + self.tied_weight_attr = [tied_weight_attr] if type(tied_weight_attr) == str else tied_weight_attr + + +class PipelineModule(nn.Module): + """Modules to be parallelized with pipeline parallelism. + + The key constraint that enables pipeline parallelism is the + representation of the forward pass as a sequence of layers + and the enforcement of a simple interface between them. The + forward pass is implicitly defined by the module ``layers``. The key + assumption is that the output of each layer can be directly fed as + input to the next, like a ``torch.nn.Sequence``. The forward pass is + implicitly: + + .. code-block:: python + + def forward(self, inputs): + x = inputs + for layer in self.layers: + x = layer(x) + return x + + .. note:: + Pipeline parallelism is not compatible with ZeRO-2 and ZeRO-3. + + Args: + layers (Iterable): A sequence of layers defining pipeline structure. Can be a ``torch.nn.Sequential`` module. + num_stages (int, optional): The degree of pipeline parallelism. If not specified, ``topology`` must be provided. + topology (``deepspeed.runtime.pipe.ProcessTopology``, optional): Defines the axes of parallelism axes for training. Must be provided if ``num_stages`` is ``None``. + loss_fn (callable, optional): Loss is computed ``loss = loss_fn(outputs, label)`` + seed_layers(bool, optional): Use a different seed for each layer. Defaults to False. + seed_fn(type, optional): The custom seed generating function. Defaults to random seed generator. + base_seed (int, optional): The starting seed. Defaults to 1234. + partition_method (str, optional): The method upon which the layers are partitioned. Defaults to 'parameters'. + activation_checkpoint_interval (int, optional): The granularity activation checkpointing in terms of number of layers. 0 disables activation checkpointing. + activation_checkpoint_func (callable, optional): The function to use for activation checkpointing. Defaults to ``deepspeed.checkpointing.checkpoint``. + checkpointable_layers (list[str], optional): List of layer class names that are eligible for checkpointing. For GPT models, + ParallelTransformerLayerPipe is always checkpointed regardless of this list. If None, all layers with parameters are + considered checkpointable. Defaults to None. + dynamic_shape: Allows dynamic shapes of inputs. This might have a performance impact. + """ + + def __init__(self, + layers, + num_stages=None, + topology=None, + loss_fn=None, + seed_layers=False, + seed_fn=None, + base_seed=1234, + partition_method='parameters', + activation_checkpoint_interval=0, + activation_checkpoint_func=checkpointing.checkpoint, + checkpointable_layers=None, + dynamic_shape=False): + + super().__init__() + + if num_stages is None and topology is None: + raise RuntimeError('must provide num_stages or topology') + + self.micro_offset = 0 + + self.loss_fn = loss_fn + + self.checkpointable_layers = checkpointable_layers + if checkpointable_layers is not None: + assert isinstance(checkpointable_layers, list), "param `checkpointable_layers` must be type of list." + + self.seed_layers = seed_layers + self.seed_fn = seed_fn + self.base_seed = base_seed + if dist.get_rank() == 0: + try: + seed_str = self.seed_fn.__name__ + except AttributeError: + seed_str = None + print(f'SEED_LAYERS={self.seed_layers} BASE_SEED={self.base_seed} SEED_FN={seed_str}') + + # Setup world info + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + self.global_rank = dist.get_rank(group=self.world_group) + self.world_size = dist.get_world_size(group=self.world_group) + self.local_rank = int(os.environ.get("LOCAL_RANK", None)) + assert self.local_rank is not None + + if topology: + self._topo = topology + self.num_stages = self._topo.get_dim('pipe') + else: + self.num_stages = num_stages + if topology is None: + if self.world_size % self.num_stages != 0: + raise RuntimeError( + f'num_stages ({self.num_stages}) must divide distributed world size ({self.world_size})') + dp = self.world_size // num_stages + topology = PipeDataParallelTopology(num_pp=num_stages, num_dp=dp) + self._topo = topology + + # Construct communicators for pipeline topology + self._grid = PipelineParallelGrid(process_group=self.world_group, topology=self._topo) + + self.stage_id = self._topo.get_coord(self.global_rank).pipe + + # Initialize partition information + self._layer_specs = list(layers) + self._num_layers = len(self._layer_specs) + self._local_start = 0 + self._local_stop = None + self._partition_layers(method=partition_method) + + self.forward_funcs = [] + self.fwd_map = {} + self.tied_modules = nn.ModuleDict() + self.tied_weight_attrs = {} + + # Offset the random seed by the stage ID. + #newseed = get_accelerator().initial_seed() + self._grid.get_stage_id() + #ds_utils.set_random_seed(newseed) + + self.activation_checkpoint_interval = activation_checkpoint_interval + + self.activation_checkpoint_func = activation_checkpoint_func + + #storage for precomputed checkpointeble results + self.is_checkpointable_results = [] + self.is_checkpointable_results_interval = None + + # if configuration use_reentrant = False, self.activation_checkpoint_func will be set to ``checkpointing.non_reentrant_checkpoint`` + + #with torch.random.fork_rng(devices=[get_accelerator().current_device_name()]): + self._build() + self.to(get_accelerator().device_name(self.local_rank)) + + self.tied_comms = self._index_tied_modules() + self._synchronize_tied_weights() + + self.dynamic_shape = dynamic_shape + + def _precompute_checkpointable_values(self): + if self.activation_checkpoint_interval > 0 and self.is_checkpointable_results_interval != self.activation_checkpoint_interval: + num_layers = len(self.forward_funcs) + self.interval_was_zero = False + for start_idx in range(0, num_layers, self.activation_checkpoint_interval): + end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) + funcs = self.forward_funcs[start_idx:end_idx] + self.is_checkpointable_results.append(self._is_checkpointable(funcs)) + self.is_checkpointable_results_interval = self.activation_checkpoint_interval + + def _build(self): + specs = self._layer_specs + + for local_idx, layer in enumerate(specs[self._local_start:self._local_stop]): + layer_idx = local_idx + self._local_start + if self.seed_layers: + if self.seed_fn: + self.seed_fn(self.base_seed + layer_idx) + else: + ds_utils.set_random_seed(self.base_seed + layer_idx) + + # Recursively build PipelineModule objects + if isinstance(layer, PipelineModule): + raise NotImplementedError('RECURSIVE BUILD NOT YET IMPLEMENTED') + + # LayerSpec objects contain an nn.Module that should be allocated now. + elif isinstance(layer, nn.Module): + name = str(layer_idx) + self.forward_funcs.append(layer) + self.fwd_map.update({name: len(self.forward_funcs) - 1}) + self.add_module(name, layer) + + # TiedLayerSpec objects contain an nn.Module that should be allocated now. + elif isinstance(layer, TiedLayerSpec): + # Build and register the module if we haven't seen it before. + if layer.key not in self.tied_modules: + self.tied_modules[layer.key] = layer.build() + self.tied_weight_attrs[layer.key] = layer.tied_weight_attr + + if layer.forward_fn is None: + # Just use forward() + self.forward_funcs.append(self.tied_modules[layer.key]) + else: + # User specified fn with args (module, input) + self.forward_funcs.append(partial(layer.forward_fn, self.tied_modules[layer.key])) + + # LayerSpec objects contain an nn.Module that should be allocated now. + elif isinstance(layer, LayerSpec): + module = layer.build() + name = str(layer_idx) + self.forward_funcs.append(module) + self.fwd_map.update({name: len(self.forward_funcs) - 1}) + self.add_module(name, module) + + # Last option: layer may be a functional (e.g., lambda). We do nothing in + # that case and just use it in forward() + else: + self.forward_funcs.append(layer) + + # All pipeline parameters should be considered as model parallel in the context + # of our FP16 optimizer + for p in self.parameters(): + p.ds_pipe_replicated = False + + def _get_frozen_parameter_names(self, layer): + """ Get names of frozen parameters in the layer. + + Returns: + A list of frozen parameter names + """ + if isinstance(layer, LayerSpec): + l = layer.build() + return [n for n, p in l.named_parameters() if not p.requires_grad] + elif isinstance(layer, nn.Module): + return [n for n, p in layer.named_parameters() if not p.requires_grad] + + return [] + + def _count_layer_params(self): + """Count the trainable parameters in individual layers. + + This routine will only build one layer at a time. + + Returns: + A list of the number of parameters in each layer. + """ + param_counts = [0] * len(self._layer_specs) + for idx, layer in enumerate(self._layer_specs): + if isinstance(layer, LayerSpec): + l = layer.build() + params = filter(lambda p: p.requires_grad, l.parameters()) + param_counts[idx] = sum(p.numel() for p in params) + elif isinstance(layer, nn.Module): + params = filter(lambda p: p.requires_grad, layer.parameters()) + param_counts[idx] = sum(p.numel() for p in params) + return param_counts + + def _find_layer_type(self, layername): + idxs = [] + typeregex = regex.compile(layername, regex.IGNORECASE) + for idx, layer in enumerate(self._layer_specs): + name = None + if isinstance(layer, LayerSpec): + name = layer.typename.__name__ + elif isinstance(layer, nn.Module): + name = layer.__class__.__name__ + else: + try: + name = layer.__name__ + except AttributeError: + continue + if typeregex.search(name): + idxs.append(idx) + + if len(idxs) == 0: + raise RuntimeError(f"Partitioning '{layername}' found no valid layers to partition.") + return idxs + + def forward(self, forward_input): + # We need to offset the seed by the microbatch ID. Save it in a local var to + # ensure it is preserved in the closure. Otherwise checkpointed forward funcs + # will see a different offset. + self.micro_offset += 1 + + def exec_range_func(start, end): + ''' Helper function to be used with checkpoint() + Adapted from torch.utils.checkpoint:checkpoint_sequential() + ''' + local_micro_offset = self.micro_offset + 1 + + def exec_func(*inputs): + # Single tensor inputs need to be unwrapped + if len(inputs) == 1: + inputs = inputs[0] + for idx, layer in enumerate(self.forward_funcs[start:end]): + self.curr_layer = idx + self._local_start + if self.seed_layers: + new_seed = (self.base_seed * local_micro_offset) + self.curr_layer + if self.seed_fn: + self.seed_fn(new_seed) + else: + ds_utils.set_random_seed(new_seed) + + inputs = layer(inputs) + return inputs + + return exec_func + + if self.activation_checkpoint_interval == 0: + func = exec_range_func(0, len(self.forward_funcs)) + x = func(forward_input) + else: + num_layers = len(self.forward_funcs) + x = forward_input + for start_idx, is_checkpointable_result in \ + zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results): + + end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) + + funcs = self.forward_funcs[start_idx:end_idx] + # Since we either pass tensors or tuples of tensors without unpacking, we + # need to be careful not to double-wrap tensors with tuple. + if not isinstance(x, tuple): + x = (x, ) + + if is_checkpointable_result: + x = self.activation_checkpoint_func(exec_range_func(start_idx, end_idx), *x) + else: + x = exec_range_func(start_idx, end_idx)(*x) + return x + + def _partition_layers(self, method='uniform'): + num_stages = self._topo.get_dim('pipe') + stage_id = self._topo.get_coord(self.global_rank).pipe + + if self.global_rank == 0: + logger.info(f'Partitioning pipeline stages with method {method}') + + method = method.lower() + + # Each stage gets a simple uniform number of layers. + if method == 'uniform': + num_layers = len(self._layer_specs) + self.parts = ds_utils.partition_uniform(num_items=num_layers, num_parts=num_stages) + elif method == 'parameters': + param_counts = self._count_layer_params() + self.parts = ds_utils.partition_balanced(weights=param_counts, num_parts=num_stages) + elif method.startswith('type:'): + layertype = method.split(':')[1] + binary_weights = [0] * len(self._layer_specs) + for idx in self._find_layer_type(layertype): + binary_weights[idx] = 1 + self.parts = ds_utils.partition_balanced(weights=binary_weights, num_parts=num_stages) + elif method == 'profile': + raise NotImplementedError(f'Partitioning method {method} not implemented.') + else: + raise NotImplementedError(f'Partitioning method {method} not implemented.') + + # Print some information on the partitioning. + if self.global_rank == 0: + for stage in range(num_stages): + start = self.parts[stage] + stop = self.parts[stage + 1] + print(f'stage={stage} layers={stop - start}') + for idx, layer in enumerate(self._layer_specs[start:stop]): + name = str(layer) + if isinstance(layer, LayerSpec): + name = layer.typename.__name__ + if isinstance(layer, nn.Module): + name = layer.__class__.__name__ + else: + try: + name = layer.__name__ + except AttributeError: + pass + print(f' {idx+start:2d}: {name}') + if self.loss_fn: + try: + print(f' loss: {self.loss_fn.__name__}') + except AttributeError: + print(f' loss: {self.loss_fn.__class__.__name__}') + + self._set_bounds(start=self.parts[stage_id], stop=self.parts[stage_id + 1]) + + @staticmethod + def _recursive_getattr(module: torch.nn.Module, attr_name: str) -> torch.Tensor: + '''Allow getting an attribute like "linear.weight"''' + weight = module + for item in attr_name.split("."): + weight = getattr(weight, item) + return weight + + def allreduce_tied_weight_gradients(self): + '''All reduce the gradients of the tied weights between tied stages''' + for key, comm in self.tied_comms.items(): + for attr_name in comm['weight_attr']: + weight = self._recursive_getattr(self.tied_modules[key], attr_name) + dist.all_reduce(weight.grad, group=comm['group']) + + def get_tied_weights_and_groups(self): + weight_group_list = [] + for key, comm in self.tied_comms.items(): + for attr_name in comm['weight_attr']: + weight = self._recursive_getattr(self.tied_modules[key], attr_name) + weight_group_list.append((weight, comm['group'])) + return weight_group_list + + def _synchronize_tied_weights(self): + for key, comm in self.tied_comms.items(): + for attr_name in comm['weight_attr']: + dist.broadcast( + self._recursive_getattr(comm['module'], attr_name), + src=min(comm['ranks']), + group=comm['group'], + ) + + def _index_tied_modules(self): + ''' Build communication structures for tied modules. ''' + tied_comms = {} + if self._topo.get_dim('pipe') == 1: + return tied_comms + + specs = self._layer_specs + tie_keys = set(s.key for s in specs if isinstance(s, TiedLayerSpec)) + # Since Python 3.7, "Dictionary order is guaranteed to be insertion order." + # Sort tie_keys here so that orders of self.tied_comms.items() are consistent + # among ranks. + for key in sorted(tie_keys): + # Find the layers that the tied module appears in + tied_layers = [] + for idx, layer in enumerate(specs): + if isinstance(layer, TiedLayerSpec) and layer.key == key: + tied_layers.append(idx) + # Find all stages with this tied module + # TODO: Would be nice to remove the nested data/model parallelism loops and + # TODO: instead generalize in some way, since we really just care about the + # TODO: stage that owns the tied layer. Then loop over each (dp, mp, ...) + # TODO: fiber to generate process groups. + tied_stages = set(self.stage_owner(idx) for idx in tied_layers) + for dp in range(self._grid.data_parallel_size): + for mp in range(self._grid.get_slice_parallel_world_size()): + tied_ranks = [] + for s in sorted(tied_stages): + if self._grid.get_slice_parallel_world_size() > 1: + tied_ranks.append(self._grid.stage_to_global(stage_id=s, data=dp, model=mp)) + else: + tied_ranks.append(self._grid.stage_to_global(stage_id=s, data=dp)) + group = dist.new_group(ranks=tied_ranks) + + # Record this tied module if we own a local copy of it. + if self.global_rank in tied_ranks: + assert key in self.tied_modules + if key in self.tied_modules: + tied_comms[key] = { + 'ranks': tied_ranks, + 'group': group, + 'weight_attr': self.tied_weight_attrs[key], + 'module': self.tied_modules[key], + } + # Only count the tied module once in the eyes of the FP16 optimizer + if self.global_rank != tied_ranks[0]: + for p in self.tied_modules[key].parameters(): + p.ds_pipe_replicated = True + ''' + if len(tied_comms) > 0: + print(f'RANK={self.global_rank} tied_comms={tied_comms}') + ''' + + return tied_comms + + def partitions(self): + return self.parts + + def stage_owner(self, layer_idx): + assert 0 <= layer_idx < self._num_layers + for stage in range(self._topo.get_dim('pipe')): + if self.parts[stage] <= layer_idx < self.parts[stage + 1]: + return stage + raise RuntimeError(f'Layer {layer_idx} not owned? parts={self.parts}') + + def _set_bounds(self, start=None, stop=None): + """Manually define the range of layers that will be built on this process. + + These boundaries are treated as list slices and so start is inclusive and stop is + exclusive. The default of None for both results in all layers being built + locally. + """ + self._local_start = start + self._local_stop = stop + + def set_checkpoint_interval(self, interval): + assert interval >= 0 + self.checkpoint_interval = interval + + def topology(self): + """ ProcessTopology object to query process mappings. """ + return self._topo + + def mpu(self): + return self._grid + + def num_pipeline_stages(self): + return self._topo.get_dim('pipe') + + def ckpt_prefix(self, checkpoints_path, tag): + """Build a prefix for all checkpoint files written by this module. """ + # All checkpoint files start with this + rank_name = 'module' + + # Data parallelism is omitted from the naming convention because we are agnostic + # to this in the checkpoint. + omit_dims = frozenset(['data']) + axes = [a for a in self._grid._topo.get_axis_names() if a not in omit_dims] + for dim in axes: + rank = getattr(self._grid._topo.get_coord(rank=self.global_rank), dim) + rank_name += f'-{dim}_{rank:02d}' + + ckpt_name = os.path.join(checkpoints_path, str(tag), rank_name) + return ckpt_name + + def ckpt_layer_path(self, ckpt_dir, local_layer_idx): + """Customize a prefix for a specific pipeline module layer. """ + idx = local_layer_idx + self._local_start + layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}') + rank_repr = self._grid._topo.get_rank_repr(rank=self.global_rank) + if rank_repr != '': + layer_ckpt_path += f'-{rank_repr}' + layer_ckpt_path += '-model_states.pt' + return layer_ckpt_path + + def ckpt_layer_path_list(self, ckpt_dir, local_layer_idx): + """Get all ckpt file list for a specific pipeline module layer. """ + idx = local_layer_idx + self._local_start + layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}-') + layer_ckpt_path += "*model_states.pt" + ckpt_files = glob.glob(layer_ckpt_path) + ckpt_files.sort() + return ckpt_files + + def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=False): + # TODO: Need to validate interaction of checkpoint_parallel_write_pipeline and fastwriter + + # Processes having the same model parallel rank on different data parallel instances + # have identical layer weights. We can distribute the task of saving the layer weights + # among the data parallel ranks. For example, if a pipeline stage has 9 layers and + # if there are 2 data parallel instances, rank 0 will save the first 5 layers and + # rank 1 will save the last 4. + dp_rank = self._grid.data_parallel_id + dp_size = self._grid.data_parallel_size + num_layers = len(self.forward_funcs) + if self.checkpoint_parallel_write_pipeline: + # spread layers evenly across data parallel ranks + offsets = ds_utils.partition_uniform(num_layers, dp_size) + start, end = offsets[dp_rank], offsets[dp_rank + 1] + else: + # data parallel rank 0 writes all layers + if dp_rank != 0: + return + start, end = 0, num_layers + layer_list = self.forward_funcs[start:end] + + checkpoint_engine.makedirs(save_dir, exist_ok=True) + for idx, layer in enumerate(layer_list): + model_ckpt_path = self.ckpt_layer_path(save_dir, start + idx) + if not hasattr(layer, 'state_dict'): + continue + + orig_state_dict = layer.state_dict() + if exclude_frozen_params: + for n in self._get_frozen_parameter_names(layer): + del orig_state_dict[n] + final_state_dict = clone_tensors_for_torch_save(orig_state_dict) + checkpoint_engine.save(state_dict=final_state_dict, path=model_ckpt_path) + + def load_state_dir(self, load_dir, checkpoint_engine, strict=True): + for idx, layer in enumerate(self.forward_funcs): + # Functions, etc. will not have state_dicts + if not hasattr(layer, 'load_state_dict'): + continue + + # get all checkpoint files for the layer. + model_ckpt_list = self.ckpt_layer_path_list(load_dir, idx) + mp_rank = self._grid.get_slice_parallel_rank() + mp_world_size = self._grid.get_slice_parallel_world_size() + + sd_loader = SDLoaderFactory.get_sd_loader(model_ckpt_list, + version=2.0, + checkpoint_engine=checkpoint_engine) + load_path, checkpoint, _ = sd_loader.load(mp_world_size, mp_rank, module_key=None, is_pipe_parallel=True) + + layer.load_state_dict(checkpoint, strict=strict) + + # if self._grid.data_parallel_id == 0: + # logger.info( + # f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}' + # ) + + self._synchronize_tied_weights() + + def _is_checkpointable(self, funcs): + + if self.activation_checkpoint_func is not checkpointing.non_reentrant_checkpoint: + # This hook excludes the embedding layer + # because only non_reentrant_checkpoint can accept inputs with requires_grad=False + # otherwise, the backward of the embedding layer won't receive gradients. + if self.__class__.__name__ in ('GPTModelPipe', 'GPT2ModelPipe'): + # For GPT models, checkpoint both transformer layers and any additional + # layers specified in checkpointable_layers (if provided) + return all('ParallelTransformerLayerPipe' in f.__class__.__name__ or ( + self.checkpointable_layers is not None and f.__class__.__name__ in self.checkpointable_layers) + for f in funcs) + + if self.checkpointable_layers is not None: + # For non-GPT models, only checkpoint layers specified in checkpointable_layers + return all(f.__class__.__name__ in self.checkpointable_layers for f in funcs) + + # Default behavior: checkpoint any layer that has parameters + params = [f.parameters() for f in funcs if isinstance(f, torch.nn.Module)] + return any(len(list(p)) > 0 for p in params) + + def get_additional_losses(self): + """ Returns model specific additional losses for reporting + + Return a dictionary of {"loss name": loss_value} or None if no additional losses. + """ + return None + + def compile(self, *args, **kwargs): + for idx, layer in enumerate(self.forward_funcs): + if isinstance(layer, nn.Module): + layer.compile(*args, **kwargs) + else: + new_layer = torch.compile(layer, *args, **kwargs) + self.forward_funcs[idx] = new_layer diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/p2p.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/p2p.py new file mode 100644 index 0000000000000000000000000000000000000000..ed6d80b8d4fbf5131d773d72bec8574f3105bac9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/p2p.py @@ -0,0 +1,182 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import msgpack +import typing + +import torch +from deepspeed import comm as dist + +from deepspeed.utils.torch import required_torch_version +from deepspeed.accelerator import get_accelerator + +_groups = None +_grid = None + +_async = [] + + +def can_send_recv() -> bool: + return required_torch_version(min_version=1.8) + + +#initializes adjacent process groups +#run this only after deepspeed.init_distributed() has been called +def init_process_groups(grid): + global _groups, _grid + _grid = grid + + assert _grid.pipe_parallel_size > 1, "There is no pipeline parallelism" + + if not can_send_recv(): + _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups] + + +def _is_valid_send_recv(src_stage, dest_stage): + first_stage = 0 + last_stage = _grid.pipe_parallel_size - 1 + assert abs(src_stage-dest_stage) == 1 or \ + (src_stage == first_stage and dest_stage == last_stage) or \ + (src_stage == last_stage and dest_stage == first_stage), \ + "Functionality currently limited to send and receive between adjacent ranks only" + + +def send(tensor, dest_stage, async_op=False): + global _groups + assert async_op == False, "Doesn't support async_op true" + src_stage = _grid.get_stage_id() + _is_valid_send_recv(src_stage, dest_stage) + + dest_rank = _grid.stage_to_global(stage_id=dest_stage) + if async_op: + global _async + op = dist.isend(tensor, dest_rank) + _async.append(op) + else: + + if can_send_recv(): + return dist.send(tensor, dest_rank) + else: + group = _get_send_recv_group(src_stage, dest_stage) + src_rank = _grid.stage_to_global(stage_id=src_stage) + return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) + + +def recv(tensor, src_stage, async_op=False): + global _groups + assert async_op == False, "Doesn't support async_op true" + dest_stage = _grid.get_stage_id() + _is_valid_send_recv(src_stage, dest_stage) + + src_rank = _grid.stage_to_global(stage_id=src_stage) + + if async_op: + global _async + op = dist.irecv(tensor, src_rank) + _async.append(op) + else: + if can_send_recv(): + return dist.recv(tensor, src_rank) + else: + group = _get_send_recv_group(src_stage, dest_stage) + return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) + + +def wait(): + global _async + for op in _async: + op.wait() + _async = [] + + get_accelerator().synchronize() + + +def send_obj(msg: typing.Any, dest: int): + """Send an arbitrary python object to ``dest``. + + Note: ``msg`` must be serializable by msgpack. + + WARN: This incurs a CPU -> GPU transfer and should be used sparingly + for performance reasons. + + Args: + msg (typing.Any): The object to send. + dest (int): Destination rank. + """ + # serialize the message + msg = msgpack.packb(msg) + # construct a tensor to send + msg = torch.ByteTensor(torch.ByteStorage.from_buffer(msg)).to(get_accelerator().device_name()) + + # Send meta and message + length_tensor = torch.tensor([len(msg)], dtype=torch.long).to(get_accelerator().device_name()) + dist.send(length_tensor, dst=dest) + dist.send(msg, dst=dest) + + +def recv_obj(sender: int) -> typing.Any: + """Receive an arbitrary python object from ``sender``. + + WARN: This incur a CPU <-> GPU transfers and should be used sparingly + for performance reasons. + + Args: + sender (int): The rank sending the message. + """ + # Get message meta + length = torch.tensor([0], dtype=torch.long).to(get_accelerator().device_name()) + dist.recv(length, src=sender) + + # Receive and deserialize + msg = torch.empty(length.item(), dtype=torch.uint8).to(get_accelerator().device_name()) + dist.recv(msg, src=sender) + + msg = msgpack.unpackb(msg.cpu().numpy().tobytes()) + + def _to(x): + """Recursively move to the current device.""" + if torch.is_tensor(x): + return x.to(get_accelerator().device_name()) + if isinstance(x, (tuple, list)): + ret = [_to(x_) for x_ in x] + if isinstance(x, tuple): + ret = tuple(ret) + return ret + # handle kwargs + if isinstance(x, dict): + ret = dict() + for key, val in x.items(): + ret[_to(key)] = _to(val) + return ret + + # Anything else is a no-op + return x + + msg = _to(msg) + return msg + + +def _get_send_recv_group(src_stage, dest_stage): + '''the group id is always the smaller rank unless its a wrap around''' + + stage_id = None + + first_stage = 0 + last_stage = _grid.pipe_parallel_size - 1 + + if (src_stage == first_stage and dest_stage == last_stage + or dest_stage == first_stage and src_stage == last_stage): + stage_id = last_stage + elif src_stage > dest_stage: + stage_id = dest_stage + else: + stage_id = src_stage + '''group_id corresponds to group of [group_id, group_id+1] + unless group_id is the rank of the last stage + in which case group_id corresponds to group[group_id-num_stages+1, group_id] + ''' + group_id = _grid.stage_to_global(stage_id=stage_id) + + return _groups[group_id] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/schedule.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..21bf24dad8edcffd54d05d0c7092ab225e1e39dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/schedule.py @@ -0,0 +1,494 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..utils import call_to_str + +from abc import ABC, abstractmethod + + +class PipeSchedule(ABC): + """Directs the execution of a pipeline engine by generating sequences of + :class:`PipeInstruction`. + + Schedules are generators that yield sequences of + :class:`PipeInstruction` to process the micro-batches in one batch. + Each yielded step is atomic in the sense that a barrier + synchronization can be placed between successive steps without + deadlock. + + Below is an example schedule that implements data parallelism with gradient accumulation: + + .. code-block:: python + + class DataParallelSchedule(PipeSchedule): + def steps(self): + for step_id in range(self.micro_batches): + cmds = [ + LoadMicroBatch(buffer_id=0), + ForwardPass(buffer_id=0), + BackwardPass(buffer_id=0), + ] + if step_id == self.micro_batches - 1: + cmds.extend([ + ReduceGrads(), + OptimizerStep(), + ]) + yield cmds + + def num_pipe_buffers(self): + return 1 + + Args: + micro_batches (int): The number of micro-batches that comprise a batch. + stages (int): The number of pipeline stages. + stage_id (int): The pipe stage that will execute the generated schedule. + """ + + def __init__(self, micro_batches, stages, stage_id): + super().__init__() + self.micro_batches = micro_batches + self.stages = stages + self.stage_id = stage_id + self.prev_stage = self.stage_id - 1 + self.next_stage = self.stage_id + 1 + + @abstractmethod + def steps(self): + """Yield a list of :class:`PipeInstruction` for each step in the schedule. + + .. note:: + Schedules must implement ``steps()`` to define the schedule. + + Returns: + Instructions to be executed as one step of the pipeline + """ + pass + + def num_pipe_buffers(self): + """The number of pipeline buffers that will be used by this stage. + + .. note:: + Schedules should specialize ``num_pipe_buffers()`` for memory savings at scale. + + Returns: + The number of buffers for the engine to allocate. + """ + return self.micro_batches + + def _valid_micro_batch(self, micro_batch_id): + return 0 <= micro_batch_id < self.micro_batches + + def _valid_stage(self, stage_id): + return 0 <= stage_id < self.stages + + @property + def stage(self): + """Stage index used to configure this schedule.""" + return self.stage_id + + @property + def num_stages(self): + """The number of total pipeline stages used to configure this schedule.""" + return self.stages + + @property + def num_micro_batches(self): + """The number of total micro_batches used to configure this schedule.""" + return self.micro_batches + + @property + def is_first_stage(self): + """True if the configured ``stage_id`` is the first stage in the pipeline.""" + return self.stage_id == 0 + + @property + def is_last_stage(self): + """True if the configured ``stage_id`` is the last stage in the pipeline.""" + return self.stage_id == self.stages - 1 + + def _buffer_idx(self, micro_batch_id): + """Map a micro-batch index to a pipeline buffer index. + + This method uses a cyclic allocation strategy. + + Args: + micro_batch_id (int): The micro-batch index relative to the beginning of the schedule. + + Returns: + int: The index of the buffer that should store data. + """ + assert self._valid_micro_batch(micro_batch_id) + return micro_batch_id % self.num_pipe_buffers() + + def __iter__(self): + self.it = None + return self + + def __next__(self): + if self.it is None: + self.it = self.steps() + return next(self.it) + + +class InferenceSchedule(PipeSchedule): + """A schedule for inferencing batches using pipeline parallelism. + """ + + def steps(self): + """""" + prev_micro_batch_id = -1 + total_steps = self.micro_batches + self.stages - 1 + for step_id in range(total_steps): + cmds = [] + micro_batch_id = step_id - self.stage_id + + # Alternate send/recv buffers + if _is_even(self.stage_id): + recv_buf = step_id % 2 + send_buf = (step_id + 1) % 2 + else: + recv_buf = (step_id + 1) % 2 + send_buf = step_id % 2 + + if self.is_first_stage or self.is_last_stage: + if self._valid_micro_batch(micro_batch_id): + cmds.append(LoadMicroBatch(recv_buf)) + + if _is_even(self.stage_id): + if self._valid_stage(self.next_stage): + if self._valid_micro_batch(micro_batch_id - 1): + cmds.append(SendActivation(send_buf)) + if self._valid_stage(self.prev_stage): + if self._valid_micro_batch(micro_batch_id): + cmds.append(RecvActivation(recv_buf)) + else: + if self._valid_stage(self.prev_stage): + if self._valid_micro_batch(micro_batch_id): + cmds.append(RecvActivation(recv_buf)) + + if self._valid_stage(self.next_stage): + if self._valid_micro_batch(micro_batch_id - 1): + cmds.append(SendActivation(send_buf)) + + if self._valid_micro_batch(micro_batch_id): + cmds.append(ForwardPass(recv_buf)) + + yield cmds + + def num_pipe_buffers(self): + """Only two pipeline buffers are required for inferencing. + + Returns: + ``2`` + """ + return 2 + + +class TrainSchedule(PipeSchedule): + """A schedule for training a batch using hybrid parallelism. + + Pipeline parallelism is extracted through gradient accumulation and thus + convergence follows that of a data parallel approach with the same batch + size. + """ + + def steps(self): + """""" + prev_micro_batch_id = -1 + total_steps = 2 * (self.micro_batches + self.stages - 1) + for step_id in range(total_steps): + # Map the step of the pipeline to the micro-batch id and also whether it is a + # forward or backward pass step. + micro_batch_id, is_forward = self._step_to_micro_batch(step_id) + + if self._valid_micro_batch(prev_micro_batch_id): + prev_buffer = self._buffer_idx(prev_micro_batch_id) + if self._valid_micro_batch(micro_batch_id): + curr_buffer = self._buffer_idx(micro_batch_id) + + cmds = [] + + # Exchange activations + if is_forward: + if self._valid_micro_batch(prev_micro_batch_id) and self._valid_stage(self.prev_stage): + cmds.append(SendGrad(prev_buffer)) + if self._valid_micro_batch(micro_batch_id) and self._valid_stage(self.prev_stage): + cmds.append(RecvActivation(curr_buffer)) + else: + if self._valid_micro_batch(micro_batch_id) and self._valid_stage(self.next_stage): + cmds.append(RecvGrad(curr_buffer)) + if self._valid_micro_batch(prev_micro_batch_id) and self._valid_stage(self.next_stage): + cmds.append(SendActivation(prev_buffer)) + + # First/last stage loads + if self.stage_id == 0 or self.stage_id == self.stages - 1: + if is_forward and self._valid_micro_batch(micro_batch_id): + cmds.append(LoadMicroBatch(curr_buffer)) + + # Computation + if self._valid_micro_batch(micro_batch_id): + if is_forward: + cmds.append(ForwardPass(curr_buffer)) + else: + cmds.append(BackwardPass(curr_buffer)) + + # Model step at the end of the batch + if step_id == total_steps - 1: + cmds.append(ReduceTiedGrads()) + cmds.append(ReduceGrads()) + cmds.append(OptimizerStep()) + + # Prepare state for next time + prev_micro_batch_id = micro_batch_id + yield cmds + + def num_pipe_buffers(self): + """Return the number of pipeline buffers required for this stage. + + This is equivalent to the maximum number of in-flight forward passes, + since we need to remember the activations of forward passes in order + to run backpropagation. For synchronous 1F1B, this is equivalent to + the index difference between this stage and the last stage. + """ + buffers = min(self.stages - self.stage_id, self.micro_batches) + return max(2, buffers) + + def _step_to_micro_batch(self, step_id): + if _is_even(step_id) and _is_even(self.stage_id): + micro_batch_id = self._even_step_forward_id(step_id) + is_forward = True + + elif _is_odd(step_id) and _is_odd(self.stage_id): + micro_batch_id = self._odd_step_forward_id(step_id) + is_forward = True + + elif _is_even(step_id) and _is_odd(self.stage_id): + micro_batch_id = self._even_step_backward_id(step_id) + is_forward = False + + elif _is_odd(step_id) and _is_even(self.stage_id): + micro_batch_id = self._odd_step_backward_id(step_id) + is_forward = False + + else: + assert False + + return micro_batch_id, is_forward + + def _even_step_forward_id(self, step_id): + base = step_id // 2 + micro_batch_id = int(base - self.stage_id // 2) + return micro_batch_id + + def _odd_step_forward_id(self, step_id): + base = (step_id - 1) // 2 + micro_batch_id = int(base - self.stage_id // 2) + return micro_batch_id + + def _even_step_backward_id(self, step_id): + base = step_id // 2 + micro_batch_id = int(base - self.stages + (self.stage_id + 1) // 2) + return micro_batch_id + + def _odd_step_backward_id(self, step_id): + base = ((step_id - 1) // 2) - self.stages + 1 + micro_batch_id = int(base + self.stage_id // 2) + return micro_batch_id + + +class DataParallelSchedule(PipeSchedule): + """An example schedule that trains using traditional data parallelism with gradient + accumulation. + """ + + def steps(self): + """""" + for step_id in range(self.micro_batches): + cmds = [ + LoadMicroBatch(buffer_id=0), + ForwardPass(buffer_id=0), + BackwardPass(buffer_id=0), + ] + if step_id == self.micro_batches - 1: + cmds.extend([ + ReduceGrads(), + OptimizerStep(), + ]) + yield cmds + + def num_pipe_buffers(self): + """Only one pipeline buffer needed. + """ + return 1 + + +class PipeInstruction: + """Base class for all instructions to be executed by the pipeline engine. + + All keyword arguments are stored as members similar to a ``namedtuple``. These are + then accessible to the :class:`PipeEngine` during execution. + + Args: + kwargs (optional): keyword arguments to store as members + """ + + def __init__(self, **kwargs): + self.name = self.__class__.__name__ + self.kwargs = kwargs + for key, val in kwargs.items(): + setattr(self, key, val) + + def __repr__(self): + return call_to_str(self.name, **self.kwargs) + + +class OptimizerStep(PipeInstruction): + """Performs one step with the optimizer and zeros gradients. + + .. note:: Should be issued after :class:`ReduceGrads` and :class:`ReduceTiedGrads`. + + .. note:: Can be a synchronization point among data-parallel ranks. + """ + pass + + +class ReduceGrads(PipeInstruction): + """Reduce the computed gradients among data-parallel processes within the stage. + """ + pass + + +class ReduceTiedGrads(PipeInstruction): + """Reduce the computed gradients of tied modules within a pipeline-parallel group. + + .. warning:: + The stages included in this synchronization point are not known until + the model is partitioned among pipeline stages. In the worst case, it + includes all pipeline stages. This instruction should be scheduled + carefully to avoid deadlocks. + """ + pass + + +class BufferOpInstruction(PipeInstruction): + """A pipeline instruction that operates on pipeline buffer(s). + + Args: + buffer_id (int): the index of the pipeline buffer() to modify. + """ + + def __init__(self, buffer_id, **kwargs): + super().__init__(buffer_id=buffer_id, **kwargs) + + +# IO +class LoadMicroBatch(BufferOpInstruction): + """Load a micro-batch into a buffer. + + Roughly: + + .. code-block:: python + + buffers['inputs'][buffer_id] = next(data_iter) + """ + pass + + +# Compute +class ForwardPass(BufferOpInstruction): + """Compute a forward pass. + + Roughly: + + .. code-block:: python + + buffers['outputs'][buffer_id] = forward(buffers['inputs'][buffer_id]) + """ + pass + + +class BackwardPass(BufferOpInstruction): + """Compute a backward pass and accumulate gradients. + + Roughly: + + .. code-block:: python + + outputs = buffers['outputs'][buffer_id] + gradients = buffers['gradients'][buffer_id] + torch.autograd.backward(tensors=outputs, + grad_tensors=gradients) + """ + pass + + +# Communication +class SendActivation(BufferOpInstruction): + """Send activations to the next stage in the pipeline. + + Roughly: + + .. code-block:: python + + send(buffers['outputs'][buffer_id]) + + .. note:: + The communication is blocking and must be paired with a :class:`RecvActivation` + on the next pipeline stage to avoid deadlock. + """ + pass + + +class RecvActivation(BufferOpInstruction): + """Receive activations from the previous stage in the pipeline. + + Roughly: + + .. code-block:: python + + buffers['inputs'][buffer_id] = recv() + + .. note:: + The communication is blocking and must be paired with a :class:`SendActivation` + on the previous pipeline stage to avoid deadlock. + """ + pass + + +class SendGrad(BufferOpInstruction): + """Send computed gradients to the previous pipeline stage. + with respect to the received activations + + .. note:: + Only received tensors with ``requires_grad==True`` will produce gradients. + Missing gradients will be replaced with ``None`` on the receiving stage. + + .. note:: + The communication is blocking and must be paired with a :class:`RecvGrad` + on the previous pipeline stage to avoid deadlock. + """ + pass + + +class RecvGrad(BufferOpInstruction): + """Receive computed gradients the next pipeline stage. + + .. note:: + Only activations with ``requires_grad==True`` will produce gradients. + Missing gradients will be replaced with ``None``. + + .. note:: + The communication is blocking and must be paired with a :class:`SendGrad` + on the next pipeline stage to avoid deadlock. + """ + pass + + +def _is_even(x): + return x % 2 == 0 + + +def _is_odd(x): + return x % 2 != 0 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/topology.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..4b107762636412ce2d455fe004fac863f6ef8f91 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/topology.py @@ -0,0 +1,472 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed import comm as dist + +from collections import namedtuple +from itertools import product as cartesian_product + + +class ProcessTopology: + """ Manages the mapping of n-dimensional Cartesian coordinates to linear + indices. This mapping is used to map the rank of processes to the grid + for various forms of parallelism. + + Each axis of the tensor is accessed by its name. The provided ordering + of the axes defines the layout of the topology. ProcessTopology uses a "row-major" + layout of the tensor axes, and so axes=['x', 'y'] would map coordinates (x,y) and + (x,y+1) to adjacent linear indices. If instead axes=['y', 'x'] was used, coordinates + (x,y) and (x+1,y) would be adjacent. + + Some methods return ProcessCoord namedtuples. + """ + + def __init__(self, axes, dims): + """Create a mapping of n-dimensional tensor coordinates to linear indices. + + Arguments: + axes (list): the names of the tensor axes + dims (list): the dimension (length) of each axis of the topology tensor + """ + + self.axes = axes # names of each topology axis + self.dims = dims # length of each topology axis + + # This is actually a class that lets us hash {'row':3, 'col':2} mappings + self.ProcessCoord = namedtuple('ProcessCoord', axes) + + self.mapping = {} + ranges = [range(d) for d in dims] + # example: 1, (0,0,1) + for global_rank, coord in enumerate(cartesian_product(*ranges)): + key = {axis: coord[self.axes.index(axis)] for axis in self.axes} + key = self.ProcessCoord(**key) + # for example, {ProcessCoord(row=0, col=1) : 1} + self.mapping[key] = global_rank + + def get_rank(self, **coord_kwargs): + """Return the global rank of a process via its coordinates. + + Coordinates are specified as kwargs. For example: + + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> X.get_rank(x=0, y=1) + 1 + """ + if len(coord_kwargs) != len(self.axes): + raise ValueError('get_rank() does not support slices. Use filter_match())') + + key = self.ProcessCoord(**coord_kwargs) + assert key in self.mapping, f'key {coord_kwargs} invalid' + return self.mapping[key] + + def get_axis_names(self): + """Return a list of the axis names in the ordering of the topology. """ + return self.axes + + def get_rank_repr(self, rank, omit_axes=['data', 'pipe'], inner_sep='_', outer_sep='-'): + """Return a string representation of a rank. + + This method is primarily used for checkpointing model data. + + For example: + >>> topo = Topo(axes=['a', 'b'], dims=[2, 2]) + >>> topo.get_rank_repr(rank=3) + 'a_01-b_01' + >>> topo.get_rank_repr(rank=3, omit_axes=['a']) + 'b_01' + + Args: + rank (int): A rank in the topology. + omit_axes (list, optional): Axes that should not be in the representation. Defaults to ['data', 'pipe']. + inner_sep (str, optional): [description]. Defaults to '_'. + outer_sep (str, optional): [description]. Defaults to '-'. + + Returns: + str: A string representation of the coordinate owned by ``rank``. + """ + omit_axes = frozenset(omit_axes) + axes = [a for a in self.get_axis_names() if a not in omit_axes] + names = [] + for ax in axes: + ax_rank = getattr(self.get_coord(rank=rank), ax) + names.append(f'{ax}{inner_sep}{ax_rank:02d}') + return outer_sep.join(names) + + def get_dim(self, axis): + """Return the number of processes along the given axis. + + For example: + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> X.get_dim('y') + 3 + """ + if axis not in self.axes: + return 0 + return self.dims[self.axes.index(axis)] + + def get_coord(self, rank): + """Return the coordinate owned by a process rank. + + The axes of the returned namedtuple can be directly accessed as members. For + example: + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> coord = X.get_coord(rank=1) + >>> coord.x + 0 + >>> coord.y + 1 + """ + for coord, idx in self.mapping.items(): + if idx == rank: + return coord + raise ValueError(f'rank {rank} not found in topology.') + + def get_axis_comm_lists(self, axis): + """ Construct lists suitable for a communicator group along axis ``axis``. + + Example: + >>> topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2]) + >>> topo.get_axis_comm_lists('pipe') + [ + [0, 4], # data=0, model=0 + [1, 5], # data=0, model=1 + [2, 6], # data=1, model=0 + [3, 7], # data=1, model=1 + ] + + Returns: + A list of lists whose coordinates match in all axes *except* ``axis``. + """ + + # We don't want to RuntimeError because it allows us to write more generalized + # code for hybrid parallelisms. + if axis not in self.axes: + return [] + + # Grab all axes but `axis` + other_axes = [a for a in self.axes if a != axis] + + lists = [] + + # Construct all combinations of coords with other_axes + ranges = [range(self.get_dim(a)) for a in other_axes] + for coord in cartesian_product(*ranges): + other_keys = {a: coord[other_axes.index(a)] for a in other_axes} + # now go over all ranks in `axis`. + sub_list = [] + for axis_key in range(self.get_dim(axis)): + key = self.ProcessCoord(**other_keys, **{axis: axis_key}) + sub_list.append(self.mapping[key]) + lists.append(sub_list) + + return lists + + def filter_match(self, **filter_kwargs): + """Return the list of ranks whose coordinates match the provided criteria. + + Example: + >>> X = ProcessTopology(axes=['pipe', 'data', 'model'], dims=[2, 2, 2]) + >>> X.filter_match(pipe=0, data=1) + [2, 3] + >>> [X.get_coord(rank) for rank in X.filter_match(pipe=0, data=1)] + [ProcessCoord(pipe=0, data=1, model=0), ProcessCoord(pipe=0, data=1, model=1)] + + Arguments: + **filter_kwargs (dict): criteria used to select coordinates. + + Returns: + The list of ranks whose coordinates match filter_kwargs. + """ + + def _filter_helper(x): + for key, val in filter_kwargs.items(): + if getattr(x, key) != val: + return False + return True + + coords = filter(_filter_helper, self.mapping.keys()) + return [self.mapping[coord] for coord in coords] + + def get_axis_list(self, axis, idx): + """Returns the list of global ranks whose coordinate in an axis is idx. + + For example: + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> X.get_axis_list(axis='x', idx=0) + [0, 1, 2] + >>> X.get_axis_list(axis='y', idx=0) + [0, 3] + """ + + # This could be faster by generating the desired keys directly instead of + # filtering. + axis_num = self.axes.index(axis) + ranks = [self.mapping[k] for k in self.mapping.keys() if k[axis_num] == idx] + return ranks + + def world_size(self): + return len(self.mapping) + + def __str__(self): + return str(self.mapping) + + +def _prime_factors(N): + """ Returns the prime factorization of positive integer N. """ + if N <= 0: + raise ValueError("Values must be strictly positive.") + + primes = [] + while N != 1: + for candidate in range(2, N + 1): + if N % candidate == 0: + primes.append(candidate) + N //= candidate + break + return primes + + +class PipeDataParallelTopology(ProcessTopology): + """ A topology specialization for hybrid data and pipeline parallelism. + + Uses data parallelism on the last dimension to encourage gradient + reductions to use high-bandwidth intra-node links and lower-volume + pipeline communications to use low-bandwidth inter-node links. + """ + + def __init__(self, num_pp, num_dp): + super().__init__(axes=['pipe', 'data'], dims=[num_pp, num_dp]) + + +class PipeModelDataParallelTopology(ProcessTopology): + """ A topology for hybrid pipeline, model, and data parallelism. """ + + def __init__(self, num_pp, num_mp, num_dp): + super().__init__(axes=['pipe', 'data', 'model'], dims=[num_pp, num_dp, num_mp]) + + +class PipelineParallelGrid: + """Implements a grid object that stores the data parallel ranks + corresponding to each of the model parallel stages + + The grid object organizes the processes in a distributed pytorch job + into a 2D grid, of stage_id and data_parallel_id. + + self.stage_id and self.data_parallel_id stores the stage id + and the data parallel id of current process. + + self.dp_group groups the processes by stage_id. + self.dp_group[i], is a list containing all process ranks whose + stage_id is i. + + self.p2p_groups stores a list of tuple, where each tuple + stores process ranks of adjacent stages for a given data_parallel_id. + For example if num_stage is 5 then a tuple [7,8] represents stages [3, 4], + with data_parallel id = 1. A stage wrap around will appear as non-adjacent ranks, + for example tuple [4,0] with representing wrap-around stage 4 and 0, for + data_parallel_id = 0, or similarly [9,5] represents wrapped around stages [4,0] + for data_parallel_id = 1. + """ + + def __init__(self, topology=None, process_group=None): + # TODO use process_group if provided + self.global_rank = dist.get_rank() + self.world_size = dist.get_world_size() + if topology is not None: + if self.global_rank == 0: + print('Using topology:', topology) + self._topo = topology + else: + num_pp = 1 + num_dp = 1 + for idx, prime in enumerate(_prime_factors(self.world_size)): + if idx % 2 == 0: + num_pp *= prime + else: + num_dp *= prime + self._topo = PipeDataParallelTopology(num_dp=num_dp, num_pp=num_pp) + self.data_parallel_size = max(self._topo.get_dim('data'), 1) + self.pipe_parallel_size = max(self._topo.get_dim('pipe'), 1) + self.model_parallel_size = max(self._topo.get_dim('model'), 1) + self.slice_parallel_size = self.model_parallel_size + assert self._is_grid_valid(), "Invalid Grid" + + self.stage_id = self.get_stage_id() + self.data_parallel_id = self.get_data_parallel_id() + + # Create new ProcessGroups for all model parallelism. DeepSpeedLight uses these + # to detect overflow, etc. + self.ds_model_proc_group = None + self.ds_model_rank = -1 + for dp in range(self.data_parallel_size): + ranks = sorted(self._topo.get_axis_list(axis='data', idx=dp)) + if self.global_rank == 0: + #print(f'RANK={self.global_rank} building DeepSpeed model group: {ranks}') + pass + proc_group = dist.new_group(ranks=ranks) + if self.global_rank in ranks: + self.ds_model_proc_group = proc_group + self.ds_model_world_size = len(ranks) + self.ds_model_rank = ranks.index(self.global_rank) + assert self.ds_model_rank > -1 + assert self.ds_model_proc_group is not None + + # Create new ProcessGroup for gradient all-reduces - these are the data parallel groups + self.dp_group = [] + self.dp_groups = self._topo.get_axis_comm_lists('data') + for g in self.dp_groups: + proc_group = dist.new_group(ranks=g) + if self.global_rank in g: + self.dp_group = g + self.dp_proc_group = proc_group + + self.is_first_stage = (self.stage_id == 0) + self.is_last_stage = (self.stage_id == (self.pipe_parallel_size - 1)) + + self.p2p_groups = self._build_p2p_groups() + + # Create new ProcessGroup for pipeline collectives - these are pipe parallel groups + self.pp_group = [] + self.pp_proc_group = None + self.pipe_groups = self._topo.get_axis_comm_lists('pipe') + for ranks in self.pipe_groups: + if self.global_rank == 0: + #print(f'RANK={self.global_rank} building pipeline group: {ranks}') + pass + proc_group = dist.new_group(ranks=ranks) + if self.global_rank in ranks: + self.pp_group = ranks + self.pp_proc_group = proc_group + assert self.pp_proc_group is not None + + # Create new ProcessGroup for model (tensor-slicing) collectives + + # Short circuit case without model parallelism. + # TODO: it would be nice if topology had bcast semantics to avoid this branching + # case? + if self.model_parallel_size == 1: + for group_rank in range(self.world_size): + group_rank = [group_rank] + group = dist.new_group(ranks=group_rank) + if group_rank[0] == self.global_rank: + self.slice_group = group_rank + self.slice_proc_group = group + return + else: + self.mp_group = [] + self.model_groups = self._topo.get_axis_comm_lists('model') + for g in self.model_groups: + proc_group = dist.new_group(ranks=g) + if self.global_rank in g: + self.slice_group = g + self.slice_proc_group = proc_group + + def get_stage_id(self): + return self._topo.get_coord(rank=self.global_rank).pipe + + def get_data_parallel_id(self): + return self._topo.get_coord(rank=self.global_rank).data + + def _build_p2p_groups(self): + """Groups for sending and receiving activations and gradients across model + parallel stages. + """ + comm_lists = self._topo.get_axis_comm_lists('pipe') + p2p_lists = [] + for rank in range(self.world_size): + for l in comm_lists: + assert len(l) == self.pipe_parallel_size + if rank in l: + idx = l.index(rank) + buddy_rank = l[(idx + 1) % self.pipe_parallel_size] + p2p_lists.append([rank, buddy_rank]) + break # next global rank + assert len(p2p_lists) == self.world_size + return p2p_lists + + def _is_grid_valid(self): + ranks = 1 + for ax in self._topo.get_axis_names(): + ranks *= self._topo.get_dim(ax) + return ranks == dist.get_world_size() + + #returns the global rank of the process with the provided stage id + #which has the same data_parallel_id as caller process + def stage_to_global(self, stage_id, **kwargs): + me = self._topo.get_coord(self.global_rank) + transform = me._replace(pipe=stage_id, **kwargs)._asdict() + return self._topo.get_rank(**transform) + + def topology(self): + return self._topo + + # MPU functions for DeepSpeed integration + def get_global_rank(self): + return self.global_rank + + def get_pipe_parallel_rank(self): + """ The stage of the pipeline this rank resides in. """ + return self.get_stage_id() + + def get_pipeline_model_parallel_rank(self): + return self.get_pipe_parallel_rank() + + def get_pipe_parallel_world_size(self): + """ The number of stages in the pipeline. """ + return self.pipe_parallel_size + + def get_pipeline_model_parallel_world_size(self): + return self.get_pipe_parallel_world_size() + + def get_pipe_parallel_group(self): + """ The group of ranks within the same pipeline. """ + return self.pp_proc_group + + def get_data_parallel_rank(self): + """ Which pipeline this rank resides in. """ + return self.data_parallel_id + + def get_data_parallel_world_size(self): + """ The number of pipelines. """ + return self.data_parallel_size + + def get_data_parallel_group(self): + """ The group of ranks within the same stage of all pipelines. """ + return self.dp_proc_group + + def get_data_parallel_group_ranks(self): + """ List of ranks in the data parallel group. """ + return self.dp_group + + # These are model parallel groups across all types of model parallelism. + # Deepspeed uses them to detect overflow, etc. + def get_model_parallel_rank(self): + return self.ds_model_rank + + def get_model_parallel_world_size(self): + return self.ds_model_world_size + + def get_model_parallel_group(self): + return self.ds_model_proc_group + + # For Megatron-style tensor slicing + def get_slice_parallel_rank(self): + if 'model' in self._topo.get_axis_names(): + return self._topo.get_coord(rank=self.global_rank).model + else: + return 0 + + def get_tensor_model_parallel_rank(self): + return self.get_slice_parallel_rank() + + def get_slice_parallel_world_size(self): + return self.slice_parallel_size + + def get_tensor_model_parallel_world_size(self): + return self.get_slice_parallel_world_size() + + def get_slice_parallel_group(self): + return self.slice_proc_group diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/precision_config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/precision_config.py new file mode 100644 index 0000000000000000000000000000000000000000..1c7fab8bd234ef3ca87a8feccd2fd8aca463e915 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/precision_config.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import DeepSpeedConfigModel +from .fp16.loss_scaler import ( + INITIAL_LOSS_SCALE, + SCALE_WINDOW, + DELAYED_SHIFT, + CONSECUTIVE_HYSTERESIS, + MIN_LOSS_SCALE, +) + +######################################### +# BFLOAT16 support +######################################### +# BFLOAT16 feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +BFLOAT16_FORMAT = ''' +BFLOAT16 parameters should be of the format: +"bf16": { + "enabled": true, + "immediate_grad_update": false, + "check_grad_overflow": false +} +''' +BFLOAT16 = "bf16" +BFLOAT16_OLD = "bfloat16" # keeping for backwards compatibility + + +def get_bfloat16_config(param_dict): + bf16_config_dict = param_dict.get(BFLOAT16, None) + if bf16_config_dict is None: + bf16_config_dict = param_dict.get(BFLOAT16_OLD, {}) + return DeepSpeedBF16Config(**bf16_config_dict) + + +class DeepSpeedBF16Config(DeepSpeedConfigModel): + """ + For bfloat16 configuration + """ + + enabled: bool = False + """ + Enable bfloat16 mixed-precision training/inference + """ + + immediate_grad_update: bool = False + """ + Apply gradient updates immediately rather than delayed. + """ + + check_grad_overflow: bool = False + """ + Check for gradient overflows and underflows + """ + + +######################################### +# FP16 support +######################################### +# FP16 feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +FP16_FORMAT = ''' +FP16 parameters should be of the format: +"fp16": { + "enabled": true, + "auto_cast": false, + "loss_scale": 0, + "initial_scale_power": 16, + "loss_scale_window": 1000, + "hysteresis": 2, + "consecutive_hysteresis": false, + "min_loss_scale": 1 +} +''' +FP16 = "fp16" + + +def get_float16_config(param_dict): + fp16_config_dict = param_dict.get(FP16, {}) + return DeepSpeedFP16Config(**fp16_config_dict) + + +class DeepSpeedFP16Config(DeepSpeedConfigModel): + """ + For float16 configuration + """ + + enabled: bool = False + """ + Enable fp16 mixed-precision training/inference + """ + + auto_cast: bool = False + """ + Automatically cast inputs to fp16 + """ + + loss_scale: float = 0 + """ + Loss scaling value. Default value of 0 means dynamic loss scaling instead of static loss scale. + """ + + initial_scale_power: int = 16 + """ + For dynamic loss scaling, set initial loss scale to 2^{initial_scale_power}. + """ + + loss_scale_window: int = 1000 + """ + Iteration intervals for raising/lowering dynamic loss scale value. + """ + + hysteresis: int = 2 + """ + Delay shift in dynamic loss scaling. + """ + + consecutive_hysteresis: bool = False + """ + Refill hysteresis if iteration does not overflow/underflow. + """ + + min_loss_scale: int = 1 + """ + Minimum dynamic loss scale value. + """ + + fp16_master_weights_and_grads: bool = False + """ + Maintain master weights in optimizer state as fp16 instead of fp32 (valid with DeepSpeedCPUAdam only). + """ + + def initial_dynamic_scale(self): + return 2**self.initial_scale_power + + def dynamic_loss_scale_args(self): + return { + INITIAL_LOSS_SCALE: 2**self.initial_scale_power, + SCALE_WINDOW: self.loss_scale_window, + DELAYED_SHIFT: self.hysteresis, + CONSECUTIVE_HYSTERESIS: self.consecutive_hysteresis, + MIN_LOSS_SCALE: self.min_loss_scale, + } diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/progressive_layer_drop.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/progressive_layer_drop.py new file mode 100644 index 0000000000000000000000000000000000000000..bd3be8f13011a4c9322839ff8a631778e7fb6847 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/progressive_layer_drop.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +from deepspeed.utils import log_dist + + +class ProgressiveLayerDrop(object): + r""" Progressive Layer Dropping (PLD) for model training. + This implements the PLD technique for compressed model training + from this paper: https://arxiv.org/pdf/2010.13369.pdf + Args: + theta (float): a hyper-parameter that controls the trade-off between training time and robustness. + The lower the theta value, the faster the training speed. Default value: 0.5. + gamma (float): a hyper-parameter that controls how fast the drop ratio increases. Default value: 0.001. + """ + + def __init__(self, theta=0.5, gamma=0.001): + super().__init__() + + self.theta = theta + self.gamma = gamma + self.current_theta = 1.0 + log_dist(f'Enabled progressive layer dropping (theta = {self.theta})', ranks=[0]) + + def get_state(self): + kwargs = {'progressive_layer_drop': True, 'pld_theta': self.get_theta()} + return kwargs + + def get_theta(self): + return self.current_theta + + def update_state(self, global_step): + + def _prob(x, gamma, p): + return (1. - p) * np.exp(-gamma * x) + p + + self.current_theta = _prob(global_step, self.gamma, self.theta) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/quantize.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..ad553138b2d7c8095aa1568f56a1f1a0b25e16d3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/quantize.py @@ -0,0 +1,180 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import math +from deepspeed.utils import logger +from deepspeed.ops.quantizer import ds_quantizer + +TWO_D_PARAMS = 6 + + +class Quantizer(object): + + def __init__(self, + q_groups=1, + q_mixed_fp16=False, + q_change_ratio=0.01, + q_type=0, + q_rounding=0, + q_verbose=False, + q_eigenvalue=False, + use_quantizer_kernel=False, + layer_num=0): + + self.q_groups = q_groups + self.q_mixed_fp16 = q_mixed_fp16 + self.q_change_ratio = q_change_ratio + self.q_type = q_type + self.qsteps = 0 + self.quantize_real_ratio = 1.000 + self.q_verbose = q_verbose + self.q_eigenvalue = q_eigenvalue + self.use_quantizer_kernel = use_quantizer_kernel + self.q_rounding = q_rounding + self.layer_num = layer_num + + def any_precision_switch(self): + # Temporary disabled functionality + if self.layer_num == 0: + return True + result = False + for index in range(self.layer_num): + if self.q_start_bits[index] != self.q_target_bits: + next_step = self.qsteps + (TWO_D_PARAMS * (self.layer_num if self.layer_num != 0 else 1)) + if next_step >= self.q_period[index]: + result = True + return result + + def quantize(self, parameter_group, overflow, eigenvalue_enabled, block_eigenvalue={}): + + if overflow and not eigenvalue_enabled: + return + + self.step() + + self.update_fp16_ratio() + + for i in range(len(parameter_group)): + for p in parameter_group[i]: + if len(p.size()) > 1 and hasattr(p, "start_bits") and p.start_bits: + param_id = id(p) + if block_eigenvalue is None: + eigenvalue, layer_id = None, 0 + else: + eigenvalue, layer_id = block_eigenvalue[param_id] if param_id in block_eigenvalue else (None, + 0) + if eigenvalue is not None: + factor = 1 + math.floor(eigenvalue * 4) + p.data = self.compute_quantization(p.data, layer_id, factor) + else: + p.data = self.compute_quantization(p, layer_id) + + def step(self): + self.qsteps += 1 + + def quantize_highbit(self, inputs, num_bits): + + q_range = 2**num_bits + input_flat = inputs.reshape(self.q_groups, -1) + g_min = input_flat.amin(dim=-1, keepdim=True) + g_max = input_flat.amax(dim=-1, keepdim=True) + + # Random number generator (Uniform) + if self.q_rounding == 'nearest': + p = 0. + else: + p = input_flat.new(input_flat.shape).uniform_(-0.5, 0.5) + + if self.q_type == 'symmetric': + scale = 2 * torch.max(torch.abs(g_min), torch.abs(g_max)) / q_range + zero_point = 0. + input_flat = (input_flat / scale + p).round().clamp(-(q_range >> 1), (q_range >> 1) - 1) * scale + elif self.q_type == 'asymmetric': + scale = (g_max - g_min) / q_range + zero_point = (g_min / scale).round() * scale + input_flat = ((input_flat - zero_point) / scale + p).round().clamp(0, (q_range - 1)) * scale + zero_point + output = input_flat.reshape(inputs.shape).contiguous() + return output + + def quantize_tenary(self, inputs): + input_flat = inputs.reshape(self.q_groups, -1) + n = input_flat.shape[1] + m = input_flat.norm(p=1, dim=1).div(n) + thres = (0.7 * m).view(-1, 1) #.expand_as(input_flat) + pos = (input_flat > thres).type(inputs.type()) + neg = (input_flat < -thres).type(inputs.type()) + mask = (input_flat.abs() > thres).type(inputs.type()) + alpha = ((mask * input_flat).abs().sum(dim=1) / mask.sum(dim=1)).view(-1, 1) + output = alpha * pos - alpha * neg + output = output.reshape(inputs.shape).contiguous() + return output + + def quantize_binary(self, inputs): + input_flat = inputs.reshape(self.q_groups, -1) + n = input_flat.shape[1] + m = input_flat.norm(p=1, dim=1, keepdim=True).div(n) + output = input_flat.sign().mul(m) + output = output.reshape(inputs.shape).contiguous() + return output + + def mixed_fp16_quantize(self, input, input_q, index): + if self.q_mixed_fp16 and self.q_start_bits[index] >= (self.q_target_bits - 1): + input_q = input * self.quantize_real_ratio + (1 - self.quantize_real_ratio) * input_q + return input_q + return input_q + + def compute_quantization(self, input, index=0, factor=1): + # fixing the quantization bits based on the training steps + # when reducing 1 bit at each period, we increase the period + # to go slowly toward the target quantization bits + # the period and starting bit can be configured + + if input.start_bits != input.target_bits: + if self.qsteps >= input.q_period: + self.quantize_real_ratio = 1.0 + input.q_period <<= 1 + input.q_period *= factor + input.start_bits -= 1 + if self.q_verbose: + logger.info( + f'Quantization settings: current bit-precision = {input.start_bits}, step = {self.qsteps}, quantization period = {input.q_period}, index = {index}' + ) + assert (input.start_bits >= input.target_bits), \ + 'Quantization bit is lower than target precision bits!' + + if self.use_quantizer_kernel: + if input.start_bits <= 2: + raise ValueError('Quantization bit is too low, please do it without quantization kernel!') + input_q = ds_quantizer(input.data.clone(), + self.q_groups, + input.start_bits, + asym=False if self.q_type == 'symmetric' else True, + sr=False if self.q_rounding == 'nearest_neighbor' else True) + else: + if input.start_bits >= 3: + input_flat = self.quantize_highbit(input.data, input.start_bits) + elif input.start_bits == 2: + assert self.q_type == 'symmetric', 'Quantization type is not symmetric!' + assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!' + input_flat = self.quantize_tenary(input.data) + elif input.start_bits == 1: + assert self.q_type == 'symmetric', 'Quantization type is not symmetric!' + assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!' + input_flat = self.quantize_binary(input.data) + if self.use_quantizer_kernel: + return self.mixed_fp16_quantize(input.data, input_q, index) + else: + if self.q_mixed_fp16 and input.start_bits >= input.target_bits - 1: + input_flat = self.quantize_real_ratio * input.data + \ + (1 - self.quantize_real_ratio) * input_flat + return input_flat + + def update_fp16_ratio(self): + if self.q_mixed_fp16: + if self.quantize_real_ratio > 0: + self.quantize_real_ratio -= self.q_change_ratio + else: + self.quantize_real_ratio = 0.000 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8cb728da375e05de5023b4196d686c25e1c4e63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) The DeepSpeed Contributors +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fbc27b819b0be9c03bf9ed941bbf946490c291ee Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc0cbfff9bef56f88258c00593ded44f65a4ed84 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c70881c86b51a22080fdb21502cdc7a93560ed7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py new file mode 100644 index 0000000000000000000000000000000000000000..c16880e15d42467f29f8103746f175b2698f4e6c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py @@ -0,0 +1,90 @@ +# Copyright (c) The DeepSpeed Contributors +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +This is a slimmed-down version of parallel_state.py (mpu) from Megatron-Deepspeed +""" + +from deepspeed import comm as dist + +# Sequence parallel groups to handle both data and sequence parallelisms. +# These groups are used to reduce gradients and shard parameters and optimizer stages for ZeRO. +_SEQUENCE_PARALLEL_GROUP = None +_SEQUENCE_DATA_PARALLEL_GROUP = None + + +def initialize_sequence_parallel(sequence_parallel_size: int) -> None: + """Initialize sequence parallel groups.""" + + assert dist.is_initialized() + world_size: int = dist.get_world_size() + + if world_size < sequence_parallel_size: + raise RuntimeError(f"world_size ({world_size}) is less than sequence_parallel_size {sequence_parallel_size}") + + if sequence_parallel_size <= 1: + raise ValueError(f"sequence_parallel_size must be greater than 1, got {sequence_parallel_size}") + + if world_size % sequence_parallel_size != 0: + raise RuntimeError( + f"world_size ({world_size}) is not divisible by sequence_parallel_size {sequence_parallel_size})") + + data_parallel_size: int = world_size // sequence_parallel_size + sequence_data_parallel_size: int = sequence_parallel_size * data_parallel_size + num_sequence_parallel_groups: int = world_size // sequence_parallel_size + num_sequence_data_parallel_groups: int = world_size // sequence_parallel_size // data_parallel_size + + rank = dist.get_rank() + + # Build the sequence parallel groups. + global _SEQUENCE_PARALLEL_GROUP + assert _SEQUENCE_PARALLEL_GROUP is None, "sequence parallel group is already initialized" + for i in range(num_sequence_parallel_groups): + ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + _SEQUENCE_PARALLEL_GROUP = group + + # Build the sequence data parallel groups. + global _SEQUENCE_DATA_PARALLEL_GROUP + assert _SEQUENCE_DATA_PARALLEL_GROUP is None, "sequence data parallel group is already initialized" + all_data_sequence_parallel_group_ranks = [] + for i in range(num_sequence_data_parallel_groups): + ranks = range(i * sequence_data_parallel_size, (i + 1) * sequence_data_parallel_size) + group = dist.new_group(ranks) + all_data_sequence_parallel_group_ranks.append(list(ranks)) + if rank in ranks: + _SEQUENCE_DATA_PARALLEL_GROUP = group + + +def get_sequence_parallel_group(): + """Get the sequence parallel group the caller rank belongs to.""" + assert _SEQUENCE_PARALLEL_GROUP is not None, "sequence parallel group is not initialized" + return _SEQUENCE_PARALLEL_GROUP + + +def get_sequence_data_parallel_group(): + """Get the sequence parallel group the caller rank belongs to.""" + assert _SEQUENCE_DATA_PARALLEL_GROUP is not None, "sequence data parallel group is not initialized" + return _SEQUENCE_DATA_PARALLEL_GROUP + + +def get_sequence_parallel_world_size(): + """Return world size for the sequence parallel group.""" + return dist.get_world_size(group=get_sequence_parallel_group()) + + +def get_sequence_data_parallel_world_size(): + """Return world size for the sequence parallel group.""" + return dist.get_world_size(group=get_sequence_data_parallel_group()) + + +def get_sequence_parallel_rank(): + """Return my rank for the sequence parallel group.""" + return dist.get_rank(group=get_sequence_parallel_group()) + + +def get_sequence_data_parallel_rank(): + """Return my rank for the sequence data parallel group.""" + return dist.get_rank(group=get_sequence_data_parallel_group()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py new file mode 100644 index 0000000000000000000000000000000000000000..3b1becc166df386723e06b21d26f067d9545d544 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py @@ -0,0 +1,1240 @@ +# Copyright (c) The DeepSpeed Contributors +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +*** Arctic Long Sequence Training (ALST) components *** + +1. Ulysses Sequence Parallelism for HF Transformers implements an efficient way of training on long sequences by employing sequence parallelism and attention head parallelism. +2. ALST enables even longer sequence lengths using a bag of tricks: +- Activation checkpoint offload to CPU +- Tiled MLP compute +- Liger-kernel +- PYTORCH_CUDA_ALLOC_CONF + +ALST features found in this module: + +- `UlyssesSPAttentionHF` - port of UlyssesAttention from Megatron-Deepspeed plus modern MHA-variations +- `UlyssesSPDataLoaderAdapter` - DL adapter to shard the normal DL batches to be used by `UlyssesSPAttentionHF` +- `SequenceTiledCompute` - generic autograd function to perform compute after tiling on the sequence dimension +- `TiledMLP` - a specific autograd function to perform tiled MLP (it's much easier to understand before trying to grok `SequenceTiledCompute`) + +This module implements Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences: https://arxiv.org/abs/2506.13996 + +For integration docs see: https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-pallellism/ + +The other ALST features live inside +https://github.com/snowflakedb/ArcticTraining/blob/main/projects/sequence-parallelism/ + +""" + +from collections import defaultdict +from deepspeed.runtime.utils import see_memory_usage +from deepspeed.sequence.layer import _DimZeroAllToAll +from einops import rearrange +from packaging import version +from torch import Tensor +from torch.utils.data import DataLoader +from typing import Any +from typing import Tuple +import deepspeed.comm as dist +import importlib.metadata +import math +import torch +import torch.distributed.nn + + +class UlyssesSPAttentionHF(torch.nn.Module): + """Re-Implementation of deepspeed.sequence.layer.DistributedAttention. This implementation enforces the input shape + to be standard [sl, bs, hc, hs] form. Any deviation from this shape will raise an error. + + The primary reason for the re-implementation is to make this less error prone, and remove what seemed like bugs in scenarios where batch size > 1 and when using different versions of + flash attention each of which takes different input shape. Those should be handled by + the actual attn implementation, and not by this module. + + This class then has been further adapted to work with HF Transformers' supported attention mechanism. + + Dimension annotation: + bs = bs + hc = head count + hc_l = head count local + hs = head_size + sl = seqlen + sl_l = seqlen local + ws = world_size + em = embedding (hidden size) + em_l = embedding (hidden size) local + + Arguments: + attn: normal attention implementation from transformers.modeling_utils.ALL_ATTENTION_FUNCTIONS + local_seq_length (int): local sequence length per GPU + global_seq_length (int): actual sequence length + batch_size (int): batch size + attn_head_size (int): size of each attention head + attn_head_count (int): total number of attention heads + kv_head_count (int): total number of kv heads + num_hidden_layers (int): total number of layers + process_group (dist.ProcessGroup): Ulysses process group + seq_length_is_variable (bool): whether global seqlen may change between batches + + + Extras: + - set self.skip_all_but_last_attention_debug_mode to True to enable fast debug which will skip calling all core attn layers but the last one, it will produce garbage of course quality-wise. + """ + + def __init__( + self, + attn, + local_seq_length: int, + global_seq_length: int, + batch_size: int, + attn_head_count: int, + attn_head_size: int, + kv_head_count: int, + num_hidden_layers: int, + process_group: dist.ProcessGroup, + seq_length_is_variable: bool = False, + ) -> None: + super().__init__() + self.attn = attn + self.process_group = process_group + self.world_size = dist.get_world_size(process_group) + self.sp_rank = dist.get_rank(process_group) + + self.local_seq_length = local_seq_length + self.global_seq_length = global_seq_length + self.batch_size = batch_size + self.seq_length_is_variable = seq_length_is_variable + + self.attn_head_size = attn_head_size + self.attn_head_count = attn_head_count + self.global_kv_head_count = kv_head_count + + self.num_hidden_layers = num_hidden_layers + self.skip_all_but_last_attention_debug_mode = False + self.rotating_layer_counter = 0 # used for dev work + + self.local_q_head_count = attn_head_count // self.world_size + + # if we have 4 kv heads and sp 8, we need to replicate kv heads 2x + self.kv_replication_factor = self.world_size // kv_head_count + if self.kv_replication_factor > 1: + self.local_kv_head_count = 1 + else: + self.local_kv_head_count = kv_head_count // self.world_size + + transformers_version_min = "4.51.3" + transformers_version_have = importlib.metadata.version("transformers") + if version.parse(transformers_version_have) < version.parse(transformers_version_min): + raise ValueError( + f"transformers>={transformers_version_min} is required, but you have transformers=={transformers_version_have}" + ) + + if self.attn_head_count % self.world_size != 0: + raise ValueError(f"Attention head count {attn_head_count} is not divisible by SP size {self.world_size}") + if not (self.global_kv_head_count % self.world_size == 0 or self.world_size % self.global_kv_head_count == 0): + raise ValueError( + f"KV attention head count {self.global_kv_head_count} is not divisible by SP size {self.world_size} or" + " vice versa") + + # [sl_l bs hc hs] + self.required_query_shape = torch.Size([local_seq_length, batch_size, attn_head_count, attn_head_size]) + self.required_key_value_shape = torch.Size([local_seq_length, batch_size, kv_head_count, attn_head_size]) + + # [sl bs em_l] + self.required_context_shape = torch.Size( + [global_seq_length, batch_size, attn_head_size * attn_head_count // self.world_size]) + + def _combine_local_sequences(self, query, key, value) -> Tuple[Tensor, Tensor, Tensor]: + + def combine_sequence(input, head_type): + """ + expects inputs in shape: [sl_l bs hc hs] + returns output in shape: [sl bs hc_l hs] + + local_head_count could be different for k,v vs q if it's not an MHA situation + """ + if head_type == "q": + local_head_count = self.local_q_head_count + else: # kv + local_head_count = self.local_kv_head_count + + # MQA and some GQA cases: + if self.kv_replication_factor > 1: + # local_head_count *= self.kv_replication_factor + # replicate heads to the kv_replication_factor on hc dimension [sl_l bs hc hs] - so dim=2 + input = input.repeat_interleave(self.kv_replication_factor, dim=2) + + # [sl_l bs hc hs] -> [sl_l bs ws hc_l hs] + input = input.reshape( + [self.local_seq_length, self.batch_size, self.world_size, local_head_count, self.attn_head_size]) + + input = rearrange(input, "sl_l bs ws hc_l hs -> ws sl_l bs hc_l hs").contiguous() + + output = _DimZeroAllToAll.apply(self.process_group, input) + + # [ws sl_l bs hc_l hs] -> [sl bs hc_l hs] + output = output.reshape([self.global_seq_length, *output.shape[2:]]).contiguous() + + # [sl bs hc_l hs] + return output + + return ( + combine_sequence(query, head_type="q"), + combine_sequence(key, head_type="kv"), + combine_sequence(value, head_type="kv"), + ) + + def _partition_global_sequence(self, input) -> Tensor: + """ + expects input in shape: [sl bs em_l] + returns output in shape: [sl_l bs em] + """ + + # [sl bs em_l] -> [ws sl_l bs em_l] + input = input.reshape([ + self.world_size, + self.local_seq_length, + self.batch_size, + self.attn_head_size * self.attn_head_count // self.world_size, + ]).contiguous() + + output = _DimZeroAllToAll.apply(self.process_group, input) + output = rearrange(output, "ws sl_l bs em_l -> sl_l bs ws em_l") + + # [sl_l bs ws em_l] -> [sl_l bs em] + output = output.reshape([*output.shape[:2], -1]).contiguous() + + # [sl_l bs em] + return output + + def forward( + self, + module: torch.nn.Module, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Tensor, + *args: Any, + **kwargs: Any, + ) -> Tensor: + """forward + + Arguments: + query (Tensor): query input to the layer + key (Tensor): key input to the layer + value (Tensor): value input to the layer + attention_mask (Tensor): Attention mask + args: other args + + Returns: + * output (Tensor): context output + """ + # HF incoming shapes are: + # [batch_size, num_heads, seqlen, head_size] + # UlyssesSPAttentionHF expects: + # [seqlen, batch_size, num_heads, head_size] + # print_rank0(f"{query.shape=}") + # print_rank0(f"{key.shape=}") + # print_rank0(f"{value.shape=}") + # print_rank0(f"{self.required_input_shape=}") + current_local_seq_length = query.shape[2] + if self.seq_length_is_variable and current_local_seq_length != self.required_query_shape[0]: + self.local_seq_length = current_local_seq_length + self.global_seq_length = current_local_seq_length * self.world_size + # update the required seqlen shapes + self.required_query_shape = torch.Size([self.local_seq_length] + list(self.required_query_shape)[1:]) + self.required_key_value_shape = torch.Size([self.local_seq_length] + + list(self.required_key_value_shape)[1:]) + self.required_context_shape = torch.Size([self.global_seq_length] + list(self.required_context_shape)[1:]) + + # make the blocks contiguous as early as possible to minimize fragmentation + query = rearrange(query, "bs hc sl hs -> sl bs hc hs") # .contiguous() + key = rearrange(key, "bs hc sl hs -> sl bs hc hs") # .contiguous() + value = rearrange(value, "bs hc sl hs -> sl bs hc hs") # .contiguous() + + # core attn like FA2 expects an unsharded `position_ids` - without which packed samples + # will return loss=nan. + # + # XXX: need to figure out if we can do the same for SDPA - as it doesn't require this and + # wants an attention mask, so possibly doing this for FA2 only? + # + # Ideally we would passing the original unsharded position_ids - but we have no way to pass + # it here as HF Transformers drops unexpected keys in `batch` - so either we need to stash + # it somewhere in UlyssesSPDataLoaderAdapter and retrieve it here or we could gather it once + # per batch and stash it inside `module` arg - I already have a machinery to figure out + # which layer number is being called below in the skip_all_but_last_attention_debug_mode + # code where rotating_layer_counter is used - so we could calculate it on the first layer + # and re-use on the remaining layers + if "position_ids" in kwargs: + position_ids_list = [torch.empty_like(kwargs["position_ids"]) for _ in range(self.world_size)] + dist.all_gather(position_ids_list, kwargs["position_ids"], group=self.process_group) + kwargs["position_ids"] = torch.cat(position_ids_list, dim=1) + + # please don't remove the white-space vertical alignment in the error message + assert query.shape == self.required_query_shape, ( + f"[{dist.get_rank()}]: query input tensor does not match the required shape\n " + f" {self.required_query_shape}:\n {query.shape=}\n {key.shape=}\n {value.shape=}") + assert key.shape == value.shape == self.required_key_value_shape, ( + f"[{dist.get_rank()}]: key or value input tensor does not match the required shape\n " + f" {self.required_key_value_shape}:\n {query.shape=}\n {key.shape=}\n {value.shape=}") + + # expects: [sl_l bs hc hs] + query_layer, key_layer, value_layer = self._combine_local_sequences(query, key, value) + # returns: [sl bs hc_l hs] + + query_layer = rearrange(query_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() + key_layer = rearrange(key_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() + value_layer = rearrange(value_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() + + # crucial in the case of MQA and some GQA cases we need to fix `module.num_key_value_groups` + # XXX: could move this somewhere to do it only once per run + if self.kv_replication_factor > 1: + module.num_key_value_groups = query_layer.size(-3) // key_layer.size(-3) + + if not self.skip_all_but_last_attention_debug_mode: + # expects: [bs hc_l sl hs] + context_layer, attn_weights = self.attn(module, query_layer, key_layer, value_layer, attention_mask, *args, + **kwargs) + # returns [bs sl hc_l hs] + else: + # we need this hack during development in order to be able to check memory fitting w/o + # waiting for 3h to compute 1.5M seqlen attention, because it's quadratic in dense + # attention, so we skip all but the last core attention call - we want the last one to + # still get the memory usage approximately close to the real memory usage. of course + # the loss will be wrong when we do that. + self.rotating_layer_counter = (self.rotating_layer_counter + 1) % self.num_hidden_layers + # we detect the last layer by module counting since we know how many layers there are + if self.rotating_layer_counter % self.num_hidden_layers == 0: + # do the real pass + context_layer, attn_weights = self.attn(module, query_layer, key_layer, value_layer, attention_mask, + *args, **kwargs) + else: + # this feeds bogus data of the right shape - good enough for quick debug + context_layer = rearrange(query_layer, "bs hc_l sl ... -> bs sl hc_l ...") + attn_weights = None + + # [bs sl hc_l hs] -> [sl bs hc_l hs]' + context_layer = rearrange(context_layer, "bs sl ... -> sl bs ...") + context_layer = context_layer.reshape([*context_layer.shape[:2], -1]) + + assert ( + context_layer.shape == self.required_context_shape + ), f"The context shape {context_layer.shape} is not of the expected shape {self.required_context_shape}" + + # expects: [sl bs em_l] + output = self._partition_global_sequence(context_layer) + # returns: [sl_l bs em] + + output = rearrange(output, "sl_l bs ... -> bs sl_l ...") + + output = output.reshape([*output.shape[:2], -1]) + + # expects [bs sl em] + return output, attn_weights + + @classmethod + def register_with_transformers( + cls, + model_name_or_path, + core_attn_implementation, + sequence_parallel_size, + max_length, + micro_batch_size, + seq_length_is_variable=True, + ): + """ + Register "ulysses" attn_implementation with HF transformers and return mpu (Megatron-LM-style parallel state object). + If sequence_parallel_size==1 do nothng and return None. + + """ + if sequence_parallel_size == 1: + return None + + from transformers import AutoConfig + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + import deepspeed.runtime.sequence_parallel.parallel_state_sp as mpu + + mpu.initialize_sequence_parallel(sequence_parallel_size=sequence_parallel_size) + + # we don't have the model yet at this stage + hf_model_config = AutoConfig.from_pretrained(model_name_or_path) + if core_attn_implementation not in ["flash_attention_2", "sdpa"]: + # notes on the excluded ones: + # - eager: The problem is that `eager` wants an attention_mask and it creates the wrong attention mask it seems if we don't provide one - it's possible that we could somehow solve this, but it's also unlikely someone will want to use the slow eager attention with sequence parallelism + # - flex_attention: haven't tried + + raise ValueError( + f"{core_attn_implementation} attn_implementation isn't currently supported by Ulysses sequence" + " parallelism. Set core_attn_implementation arg to either 'flash_attention_2' or 'sdpa'.") + + if core_attn_implementation not in ALL_ATTENTION_FUNCTIONS: + raise ValueError( + f"{core_attn_implementation} is not a valid attn_implementation. The choices are {ALL_ATTENTION_FUNCTIONS.valid_keys()}" + ) + core_attn_function = ALL_ATTENTION_FUNCTIONS[core_attn_implementation] + uattn = UlyssesSPAttentionHF( + attn=core_attn_function, + local_seq_length=max_length // mpu.get_sequence_parallel_world_size(), + global_seq_length=max_length, + batch_size=micro_batch_size, + attn_head_count=hf_model_config.num_attention_heads, + attn_head_size=getattr(hf_model_config, "head_dim", + hf_model_config.hidden_size // hf_model_config.num_attention_heads), + kv_head_count=hf_model_config.num_key_value_heads, + num_hidden_layers=hf_model_config.num_hidden_layers, + process_group=mpu.get_sequence_parallel_group(), + seq_length_is_variable=seq_length_is_variable, + ) + + def uattn_wrapper( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + *args, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + + # We are relaying on position_ids for SP to work so attention_mask has to be None + # the problem is that HF currently doesn't know anything about ALL_ATTENTION_FUNCTIONS["ulysses"] so it doesn't make a special case like for "flash_attention_2" and "sdpa" and it creates an attention mask on the fly and it breaks things. + attention_mask = None + + attn_output, attn_weights = uattn( + module, + query, + key, + value, + attention_mask, + # XXX: fixme + *args, + **kwargs, + ) + return attn_output, attn_weights + + # We don't do: ALL_ATTENTION_FUNCTIONS.register("ulysses", uattn_wrapper) + # The problem with this approach is that we are missing on all the special use cases in HF Transformers that do things like: if self.config._attn_implementation == "flash_attention_2": ... + # So instead we hack `ALL_ATTENTION_FUNCTIONS` to override all existing keys with our implementation, since it only gets used at the point of calling the attention and that's what we want, all other code branches relying on the original core `attn_implementation` will still be executed. This is what we called "Being John Malkovich" + for key in ALL_ATTENTION_FUNCTIONS.keys(): + ALL_ATTENTION_FUNCTIONS[key] = uattn_wrapper + + return mpu + + +class UlyssesSPDataLoaderAdapter: + + def __init__( + self, + dl: DataLoader, + sp_rank: int, + sp_group, + sp_world_size, + device, + ): + """ + This a DataLoader adapter which wraps around any existing DataLoader. It is used in conjunction with Ulysses to perform batch sharding on the sequence dimension. + + It gathers 1 sample from each participating rank, using the DL it wraps, then shards each of them and sends back to the ranks. So that when dl->iter->next is called, we end up with: + - rank 0: getting batch 0 shard 0 + - rank 1: getting batch 0 shard 1 + ... + - rank n: getting batch 0 shard n + which is used to compute the batch (from rank0) using all SP ranks. + + When the next iteration starts and dl->iter->next is called, we end up with: + - rank 0: getting batch 1 shard 0 + - rank 1: getting batch 1 shard 1 + ... + - rank n: getting batch 1 shard n + which is used to compute a second batch (from rank1) using all SP ranks. + + This continues until SP iterations are performed. At this point we need to get more data and so the above repeats. + + The key thing to understand is that all SP ranks participate in processing a single DL sample. So instead of normal DataParallel we perform a sort of SP over DP. + + When SP number of iterations is completed it's an equivalent of performing a single iteration with normal DP. + + If more tokens need to be consumed per step use the gradient accumulation feature. + + Arguments: + - `dl`: an existing DataLoader object to wrap + - `sp_rank`: SP rank + - `sp_group`: SP group + - `sp_world_size`: SP world size + - `device`: cuda device + + Returns: + Another DataLoader object + + Here are the current assumptions on the inputs fetched by dl->iter->next + - the batch is a dict with at least the keys: `input_ids`, `labels`, `position_ids` - but can have any additional keys necessary. + - the tensor values get sharded, the non-tensor values are passed along as is + """ + + self.dl = dl + self.sp_rank = sp_rank + self.sp_group = sp_group + self.sp_world_size = sp_world_size + self.device = device + + self.iter = iter(dl) + self.micro_batches: list[Any] = [] + + def __len__(self): + return len(self.dl) * self.sp_world_size + + def __iter__(self): + return self + + def __next__(self): + if len(self.micro_batches) == 0: + self.refill() + + return self.micro_batches.pop(0) + + def refill(self): + # this will raise StopIteration when empty + batch = next(self.iter) + micro_batches = defaultdict(dict) + # XXX: replace with more efficient all-to-all? + + # we have batches of variable seqlen so in order to do all_gather on batches - we need to know the exact length of each tensor on each rank + seqlen = torch.tensor(batch["input_ids"].shape[1], dtype=torch.int64, device=self.device) + seqlens = [torch.zeros(1, dtype=torch.int64, device=self.device) for _ in range(self.sp_world_size)] + dist.all_gather(seqlens, seqlen, group=self.sp_group) + seqlens = [x[0].item() for x in seqlens] + + for k in batch.keys(): + if torch.is_tensor(batch[k]): + batch[k] = batch[k].to(self.device) + with torch.no_grad(): + tensor_list = [ + torch.zeros((batch[k].shape[0], seqlens[i]), dtype=batch[k].dtype, device=batch[k].device) + for i in range(self.sp_world_size) + ] + dist.all_gather(tensor_list, batch[k], group=self.sp_group) + else: + tensor_list = [None for _ in range(self.sp_world_size)] + dist.all_gather_object(tensor_list, batch[k], group=self.sp_group) + + for rank, tensor in enumerate(tensor_list): + micro_batches[rank][k] = tensor + + del tensor_list + del batch + + for batch in micro_batches.values(): + seq_length = len(batch["input_ids"][0]) + + if seq_length % self.sp_world_size != 0: + raise ValueError(f"batch's seqlen={seq_length} isn't divisible by sp-size={self.sp_world_size}") + chunk_len = seq_length // self.sp_world_size + + # because we have to gather logits from all sp ranks we have to do the loss function ourselves + # therefore remove labels to avoid an attempt to calculate loss by transformers + labels = batch.pop("labels") + labels = torch.nn.functional.pad(labels, (0, 1), value=-100) + batch["shift_labels"] = labels[..., 1:].contiguous() + # free up temp memory + del labels + + # batch sharding + for k in batch.keys(): + # leave non-tensors alone + if not torch.is_tensor(batch[k]): + continue + # at seqlen>10M and 32+ gpus this can take GBs of memory so keep the prefill buffer on cpu + batch[k] = batch[k][:, chunk_len * self.sp_rank:chunk_len * (self.sp_rank + 1)].cpu() + + self.micro_batches.append(batch) + + +def sequence_tiled_compute( + fn, + seqlen, + shards, + kwargs_to_shard, + kwargs_to_pass, + grad_requiring_tensor_key, + compute_params=None, + output_unshard_dimension=1, + output_reduction="mean", +): + """ + This is a wrapper for SequenceTiledCompute which we need since torch.autograd.Function can't work with dicts of tensors (in backward it has to return a grad value and not a dict that may have a non-None grad value). It's also useful for setting default values which we can't do either in torch.autograd.Function. + + Args: + - `fn`: the function to call on sharded inputs + - `seqlen`: total seqlen of the seqlen dimension + - `shards`: how many shards to use + - `kwargs_to_shard`: this dict will be passed to `fn` as `**kwargs` after sharding on seqlen dimension + - `kwargs_to_pass`: this dict will be passed to `fn` as is, as `**kwargs` + - `grad_requiring_tensor_key`: which main key requires grads + - `compute_params`: a list of weights engaged in the compute. Default: `None` (only needed when using DeepSpeed ZeRO) + - `output_reduction`: None, "mean" or "sum": Default: "mean" + - `output_unshard_dimension`: the dimension to concat the outputs on: Default: 1 (seqlen dim) + + Returns: + - unsharded output with an optional reduction applied, depending on the `output_reduction` value: + `None` - return the unsharded output tensor + `"mean"` - apply mean + `"sum"` - apply sum + + Please note that this implementation doesn't require DeepSpeed and can work without it. `compute_params` can remain `None` in such a case. + + """ + args_to_shard = kwargs_to_shard.values() + keys_to_shard = list(kwargs_to_shard.keys()) + args_to_pass = kwargs_to_pass.values() + keys_to_pass = list(kwargs_to_pass.keys()) + + return SequenceTiledCompute.apply( + fn, + seqlen, + shards, + keys_to_shard, + keys_to_pass, + grad_requiring_tensor_key, + compute_params, + output_unshard_dimension, + output_reduction, + *args_to_shard, + *args_to_pass, + ) + + +class SequenceTiledCompute(torch.autograd.Function): + """ + A generic autograd function to perform a tiled compute. + + Please note that this implementation doesn't require DeepSpeed and can work without it. `compute_params` can remain `None` in such a case. + + For an easier to understand example see TiledMLP - which is the same as this autograd function but without the generalization code. + """ + + @staticmethod + def forward( + ctx, + fn, + seqlen, + shards, + keys_to_shard, + keys_to_pass, + grad_requiring_tensor_key, + compute_params, + output_unshard_dimension, + output_reduction, + *args, + ) -> torch.Tensor: + """ + for args and return values see `sequence_tiled_compute`'s doc + + Currently we assume that all kwargs_to_shard values have a shape of `[bs, seqlen, ...]` and we shard on seqlen dimension + """ + ctx.fn = fn + ctx.seqlen = seqlen + ctx.shards = shards + ctx.grad_requiring_tensor_key = grad_requiring_tensor_key + ctx.compute_params = [p for p in compute_params if p.requires_grad] + ctx.output_unshard_dimension = output_unshard_dimension + + with torch.no_grad(): + args = list(args) + ctx.total_args = len(args) + ctx.grad_requiring_tensor_key_index = (keys_to_shard + keys_to_pass).index(grad_requiring_tensor_key) + + kwargs_to_shard = {k: args.pop(0) for k in keys_to_shard} + kwargs_to_pass = {k: args.pop(0) for k in keys_to_pass} + ctx.kwargs_to_shard = kwargs_to_shard + ctx.kwargs_to_pass = kwargs_to_pass + + with torch.no_grad(): + shard_step = math.ceil(seqlen / shards) + output_shards = [] + + for i in range(shards): + output = fn( + **{ + k: v[:, i * shard_step:(i + 1) * shard_step] + for k, v in kwargs_to_shard.items() + }, + **kwargs_to_pass, + ) + output_shards.append(output) + + if output_unshard_dimension == 0: + # this is just the shape=[1] loss use-case, not sure if it's generic enough + output_unsharded = torch.cat([l.unsqueeze(0) for l in output_shards], dim=output_unshard_dimension) + else: + output_unsharded = torch.cat(output_shards, dim=output_unshard_dimension) # .clone().detach() + + if output_reduction is None: + return output_unsharded + elif output_reduction == "mean": + return output_unsharded.mean() + elif output_reduction == "sum": + return output_unsharded.sum() + else: + raise ValueError(f"unknown value {output_reduction}: valid values are: none/mean/sum") + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + fn = ctx.fn + shards = ctx.shards + kwargs_to_shard = ctx.kwargs_to_shard + kwargs_to_pass = ctx.kwargs_to_pass + + grad_requiring_tensor_key = ctx.grad_requiring_tensor_key + grad_requiring_tensor_key_index = ctx.grad_requiring_tensor_key_index + compute_params = ctx.compute_params + output_unshard_dimension = ctx.output_unshard_dimension + grad_requiring_tensor = kwargs_to_shard[grad_requiring_tensor_key] + + grad_requiring_tensor_requires_grad = grad_requiring_tensor.requires_grad + grad_requiring_tensor = grad_requiring_tensor.detach() + # detach() unsets `grad_requiring_tensor.requires_grad`, so restore it + grad_requiring_tensor.requires_grad_(grad_requiring_tensor_requires_grad) + + incoming_grad = grads[0] + grad_requiring_tensor_grad = torch.zeros_like(grad_requiring_tensor) + + kwargs_to_shard_shards = { + k: list(torch.chunk(kwargs_to_shard[k], chunks=shards, dim=1)) + for k in kwargs_to_shard.keys() + } + + # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step + shard_step = kwargs_to_shard_shards[grad_requiring_tensor_key][0].numel() + for i in range(shards): + + # when fn involves one or more model weights deepspeed will normally push a grad to + # reduce per sub-module call, so since we only want it to add a grad for the last + # shard's call , we signal to zero not to add new gradients to reduce until the last + # shard when all gradients have been accumulated. An example for such a call is + # `model.lm_head(hidden_states)` + if compute_params is not None: + if i + 1 < shards: + for param in compute_params: + param.ds_grad_is_ready = False + else: + # last shard, can add the grad + for param in compute_params: + param.ds_grad_is_ready = True + + kwargs_to_shard_shard = {k: kwargs_to_shard_shards[k].pop(0) for k in kwargs_to_shard_shards.keys()} + grad_requiring_tensor_shard = kwargs_to_shard_shard[grad_requiring_tensor_key] + + grad_requiring_tensor_shard.requires_grad_(grad_requiring_tensor_requires_grad) + + shard_offset = i * shard_step + # this will enable gradual population of the pre-allocated + # `grad_requiring_tensor_shard.grad` during `torch.autograd.backward` calls + grad_requiring_tensor_shard.grad = (grad_requiring_tensor_grad.view(-1).narrow( + 0, shard_offset, grad_requiring_tensor_shard.numel()).view_as(grad_requiring_tensor_shard)) + + with torch.enable_grad(): + output = fn(**kwargs_to_shard_shard, **kwargs_to_pass) + + if output_unshard_dimension == 0: + # loss use-case + torch.autograd.backward(output, incoming_grad) + else: + incoming_grad_shard = (incoming_grad.view(-1).narrow( + 0, shard_offset, grad_requiring_tensor_shard.numel()).view_as(grad_requiring_tensor_shard)) + torch.autograd.backward(output, incoming_grad_shard) + + # positional args + grad_outputs = [None] * 9 + # inject the grad for the position of forward input that is grad-requiring + arg_outputs = [None] * ctx.total_args + arg_outputs[grad_requiring_tensor_key_index] = grad_requiring_tensor_grad + + return tuple(grad_outputs + arg_outputs) + + +class TiledMLP(torch.autograd.Function): + """ + Perform a tiled MLP computation to massively reduce memory usage needed to compute MLP when using very long sequence lengths + + For a general tiled compute implementation that can handle any `forward` see `SequenceTiledCompute` + + Args: + - fn: the function to call on sharded inputs + - `self`: the MLP nn.Module object + - `x`: the input to MLP.forward (`hidden_states`) + - `shards`: how many shards to use + - compute_params: a list of weights engaged in the compute Default: `None` (only needed when using DeepSpeed ZeRO) + + Returns: + - the computed `hidden_states` + + Here is an example that monkey patches HF Transformers' LLamaMLP: + + def tiled_mlp_forward(self, x): + bs, seqlen, hidden = x.shape + num_shards = math.ceil(seqlen / hidden) + # to avoid deadlocks get all ranks to agree on the same num_shards by using the max value + tensor = torch.tensor(num_shards, device=x.device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + num_shards = tensor.item() + compute_params = [self.down_proj.weight, self.gate_proj.weight, self.up_proj.weight] + + def mlp_forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return TiledMLP.apply( + mlp_forward, + self, + x, + num_shards, + compute_params, + ) + + # this needs to be done before the model is instantiated + from transformers.models.llama import modeling_llama + modeling_llama.LlamaMLP.forward = tiled_mlp_forward + """ + + @staticmethod + def forward( + ctx, + fn, + self, + x, + shards, + compute_params, + ) -> torch.Tensor: + ctx.fn = fn + ctx.self = self + ctx.shards = shards + ctx.compute_params = [p for p in compute_params if p.requires_grad] + ctx.save_for_backward(x) + + x_shards = list(torch.chunk(x, chunks=shards, dim=1)) + with torch.no_grad(): + output_shards = [fn(self, x_shard) for x_shard in x_shards] + output_unsharded = torch.cat(output_shards, dim=1) + + return output_unsharded + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + fn = ctx.fn + (x, ) = ctx.saved_tensors + self = ctx.self + shards = ctx.shards + compute_params = ctx.compute_params + + x_requires_grad = x.requires_grad + x = x.detach() + # detach() unsets `x.requires_grad`, so restore it + x.requires_grad_(x_requires_grad) + + incoming_grad = grads[0] + x_grad = torch.zeros_like(x) + x_shards = list(torch.chunk(x, chunks=shards, dim=1)) + + shard_step = x_shards[0].numel() + for i, x_shard in enumerate(x_shards): + + # Tell deepspeed not to add a new grad to its ipg bucket until the last shard is run + if compute_params is not None: + if i + 1 < shards: + for param in compute_params: + param.ds_grad_is_ready = False + else: + # last shard, can add the grad + for param in compute_params: + param.ds_grad_is_ready = True + + x_shard.requires_grad_(x_requires_grad) + + shard_offset = i * shard_step + x_shard.grad = x_grad.view(-1).narrow(0, shard_offset, x_shard.numel()).view_as(x_shard) + incoming_grad_shard = incoming_grad.view(-1).narrow(0, shard_offset, x_shard.numel()).view_as(x_shard) + with torch.enable_grad(): + output = fn(self, x_shard) + torch.autograd.backward(output, incoming_grad_shard) + + return (None, None, x_grad, None, None) + + +class AutogradComputeMLP(torch.autograd.Function): + """ + This is a simplified example to override the normal MLP via an autograd function - then tiling can be added - this simplified version was useful to detect a leak in Deepspeed, so let's keep it. + + Here is an example of performing the monkey patching on LlamaMLP + + def mlp_forward_new(self, x): + + def mlp_forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return AutogradComputeMLP.apply(mlp_forward, self, x) + + from transformers.models.llama import modeling_llama + modeling_llama.LlamaMLP.forward = mlp_forward_new + """ + + @staticmethod + def forward( + ctx, + fn, + self, + x, + ) -> torch.Tensor: + ctx.fn = fn + ctx.self = self + ctx.save_for_backward(x) + + with torch.no_grad(): + return fn(self, x) + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + fn = ctx.fn + (x, ) = ctx.saved_tensors + self = ctx.self + + x1 = x.detach() + x1.requires_grad = x.requires_grad + with torch.enable_grad(): + output = fn(self, x1) + + torch.autograd.backward(output, grads[0]) + return (None, None, x1.grad, None) + + +########################################################### +### below are older versions that some might still want ### +########################################################### + + +class TiledLoss(torch.autograd.Function): + + @staticmethod + def forward(ctx, loss_fn, logits, vocab_size, shift_labels, shards) -> torch.Tensor: + """ + + This is a memory efficient loss autograd function that takes the existing logits and performs loss calculation in shards. + + This one is an SFT-aware version, therefore it takes care of special cases where the whole shard is made of -100 labels and which requires then a special care. + + Note: logits seqlen dimension doesn't have to be divisible by shards, the last shard will be shorter than the rest. The calculating of the number of shards is in the example. + + Here is an example of using it: + + def loss(self, batch) -> torch.Tensor: + batch = to_device(batch, self.device) + shift_labels = batch.pop("shift_labels") + outputs = self.model(**batch, use_cache=False) + logits = outputs.logits + + if all((shift_labels == -100).squeeze()): + # this is the case where all labels in a micro-batch are -100 (very common for SFT if the seqlen is short) - CE returns `nan` in this case, so we don't want to call loss and instead create a differentiable loss `0` which will also set all the grads to `0` in `backward` - the effect of this is akin to a perfect score where the model needs no adjustment since grads will be all zeros. + loss = (logits.sum() * 0.0).float() + + num_shards: Any = "auto" + if num_shards == "auto": + # parameterize to about 1GB fp32 logits shards + slice_size_in_gb = 1 + size_in_gb = logits.numel() * 4 / 2**30 # fp32 + # the sp shard's seqlen sp shard can be easily not divisible by the derived number of chunked loss shards, so we use the uppper ceiling and allow the last chunk to be shorter than the rest + num_shards = math.ceil(size_in_gb / slice_size_in_gb) + # print(f"derived {num_shards} shards for size {size_in_gb}GB") + if num_shards > 1: + # if shards == 1 this will lead to a higher memory usage then calling the normal loss function, so don't do that. + loss = TiledLoss.apply( + self.model_unwrapped.loss_function, + logits, + self.model_unwrapped.config.vocab_size, + shift_labels, + num_shards, + ) + else: + loss = self.model_unwrapped.loss_function( + logits=logits, + labels=None, + vocab_size=self.model_unwrapped.config.vocab_size, + shift_labels=shift_labels, + ) + + return loss + + + """ + ctx.save_for_backward(logits, shift_labels) + ctx.loss_fn = loss_fn + ctx.vocab_size = vocab_size + ctx.shards = shards + + with torch.no_grad(): + seqlen = shift_labels.shape[1] + shard_step = math.ceil(seqlen / shards) + loss_shards = [] + total_good_items = 0 + + # since -100s are ignored we have to perform a weighted average on each loss slice as each slice may contribute a different number of non- -100 labels + # if seqlen / shards != 0 - the last chunk is just shorter than the rest but no data is ignored + for i in range(shards): + # XXX: here and everywhere don't make a copy, pass the slice or perhaps narrow/view? + shift_labels_shard = shift_labels[:, i * shard_step:(i + 1) * shard_step] + if all((shift_labels_shard == -100).squeeze()): + continue # ignore this shard + loss_shard = loss_fn( + logits=logits[:, i * shard_step:(i + 1) * shard_step, :], + labels=None, + vocab_size=vocab_size, + shift_labels=shift_labels_shard, + ) + good_items = sum((shift_labels_shard != -100).squeeze()) + loss_shards.append(loss_shard * good_items) + total_good_items += good_items + total_loss = torch.cat([l.unsqueeze(0) for l in loss_shards], dim=0).sum() + weighted_loss = total_loss / total_good_items + + return weighted_loss + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + logits, shift_labels = ctx.saved_tensors + loss_fn = ctx.loss_fn + vocab_size = ctx.vocab_size + shards = ctx.shards + + grad = grads[0] + logits_grad = torch.zeros_like(logits) + logits_shards = list(torch.chunk(logits, chunks=shards, dim=1)) + shift_labels_shards = list(torch.chunk(shift_labels, chunks=shards, dim=1)) + + # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step + shard_step = logits_shards[0].numel() + for i in range(shards): + logits_shard = logits_shards.pop(0) + shift_labels_shard = shift_labels_shards.pop(0) + + shard_offset = i * shard_step + # this will enable gradual population of the pre-allocated `logits_shard.grad` during `torch.autograd.backward` calls + logits_shard.grad = (logits_grad.view(-1).narrow(0, shard_offset, + logits_shard.numel()).view_as(logits_shard)) + + with torch.enable_grad(): + if all((shift_labels_shard == -100).squeeze()): + # fake loss calculation, since CE will return nan, but grads will be set + # a normal loss_fn upcasts logits to float so match it + loss_shard = (logits_shard.sum() * 0.0).float() + else: + loss_shard = loss_fn( + logits=logits_shard.requires_grad_(), + labels=None, + vocab_size=vocab_size, + shift_labels=shift_labels_shard, + ) + + torch.autograd.backward(loss_shard, grad) + + logits_grad /= shards + + # only logits (2nd arg) needs grads + return None, logits_grad, None, None, None + + +# This is the original implementation/integration of UlyssesSP into the training loop, which was superseded by using UlyssesSPDataLoaderAdapter which did all the sharding and pull the shards from the DL +# +# There are 2 issues with this implementation: +# - it's complex and difficult to integrate into various training scenarios +# - it could lead to a huge number of tokens per step - e.g. 32 ranks of 15M seqlen -> 0.5B token step - which is very wasteful +# +# Therefore if you want to use UlyssesSP via UlyssesSPFwdLossBwdWithLogits with its fwd/loss/bwd for those don't want to use UlyssesSPDataLoaderAdapter - here is how it should be installed into the sub-trainer class: +# class SFTTrainer(Trainer): +# def sp_fwd_loss_bwd(self, batch) -> torch.Tensor: +# batch = to_device(batch, self.device) +# +# from arctic_training.trainer.trainer import UlyssesAttentionHFFwdLossBwdWithLogits +# ulysses = UlyssesAttentionHFFwdLossBwdWithLogits( +# model=self.model, +# model_unwrapped=self.model_unwrapped, +# device=self.device, +# num_loss_logit_shards="auto", +# ) +# return ulysses.sp_fwd_loss_bwd(batch) + + +class UlyssesSPFwdLossBwdWithLogits: + + def __init__(self, model, model_unwrapped, device, num_loss_logit_shards="auto", **kwargs): + + self.model = model + self.model_unwrapped = model_unwrapped + self.device = device + self.num_loss_logit_shards = num_loss_logit_shards + self.kwargs = kwargs + + from deepspeed.utils import groups + + self.sp_group = groups._get_sequence_parallel_group() + self.sp_world_size = groups._get_sequence_parallel_world_size() + self.sp_rank = groups._get_sequence_parallel_rank() + + def sp_fwd_loss_bwd(self, batch) -> torch.Tensor: + + see_memory_usage(f"entered sp_fwd_loss_bwd", force=True) + + # ensure shapes are correct + if not (batch["input_ids"].shape == batch["position_ids"].shape == batch["labels"].shape): + raise ValueError( + f'Borked batch {batch["input_ids"].shape=} != {batch["position_ids"].shape=} !=' + f' {batch["labels"].shape=}) in DataLoader->iter->next, cannot continue with Ulysses Sequence' + " parallelism") + + # gather DL batches into super-batches + # Important: DL doesn't always yield max_length batches. Different ranks may have different seqlen and each could be <= max_length (but always divisible by 256) + + micro_batches: list[Any] = defaultdict(dict) + # Efficient gathering of batch inputs across ranks: + # The problem is that our DL doesn't guarantee the same seqlen on all ranks and may give, 3x 1024 and 1x 768 on 4 gpus for max_length 1024. so 3 options we have to be able to gather batches are: + # 1. use all_gather_object - which allows different shapes - but potentially introducing an undesired overhead - 2x pickle calls + # 2. use all_gather and change DL pad to make sure that all ranks always get the same input shape - this creates its own overhead since if we say have ranks with seqlen 512, 768, 1024, 1024 - now we will need to process 4x 1024 seqlens + # 3. use all_gather and post gathering truncate tensors to their intended length - another overhead of allocating and truncating tensors + # using approach (1) for now but might want to benchmark later the other 2 approaches + + # XXX: if using all_gather_object we can gather the whole batch at once and not per-key! so can drop the loop for that approach + + # we have batches of variable seqlen so in order to do all_gather on batches - we need to know the exact length of each tensor on each rank + seqlen = torch.tensor(batch["input_ids"].shape[1], dtype=torch.int64, device=self.device) + # print(seqlen) + seqlens = [torch.zeros(1, dtype=torch.int64, device=self.device) for _ in range(self.sp_world_size)] + dist.all_gather(seqlens, seqlen, group=self.sp_group) + seqlens = [x[0].item() for x in seqlens] + + for k in batch.keys(): + batch[k] = batch[k].to(self.device) + with torch.no_grad(): + tensor_list = [ + torch.zeros((batch[k].shape[0], seqlens[i]), dtype=batch[k].dtype, device=batch[k].device) + for i in range(self.sp_world_size) + ] + dist.all_gather(tensor_list, batch[k], group=self.sp_group) + + # gathering on the data dimension + # will be concatenating and later splitting again for the more general case + # batch[k] = torch.cat(tensor_list, dim=1) + for rank, tensor in enumerate(tensor_list): + micro_batches[rank][k] = tensor + + del tensor_list + del batch + + # we need to chunk twice - each time on SP size level + # - the first time is because we artificially made the seqlen SP-times longer + # - the second time is because of the Ulysses algorithm + + see_memory_usage("after gathering", force=False) + + self.model.set_gradient_accumulation_boundary(False) + + losses = [] + for sub_step_id in range(self.sp_world_size): + batch = micro_batches[sub_step_id] + seq_length = len(batch["input_ids"][0]) + + if seq_length % self.sp_world_size != 0: + raise ValueError( + f"{sub_step_id=}: batch's seqlen={seq_length} isn't divisible by sp-size={self.sp_world_size}") + chunk_len = int(seq_length / self.sp_world_size) + + # to enable the correct mean calculation across shards before sharding the micro batch: + # 1. count the number of non- `-100`` elements per shard + # 2. and subtract one more element because of label shifting + non_skipped_items = {} + for rank in range(self.sp_world_size): + non_skipped = (batch["labels"][:, chunk_len * rank:chunk_len * (rank + 1)] != -100).sum().item() + if non_skipped > 1: + non_skipped -= 1 + non_skipped_items[rank] = non_skipped + + # because we have to gather logits from all sp ranks we have to do the loss function ourselves + # therefore remove labels to avoid an attempt to calculate loss by transformers + labels = batch.pop("labels") + labels = torch.nn.functional.pad(labels, (0, 1), value=-100) + batch["shift_labels"] = labels[..., 1:].contiguous() + # free up temp memory + del labels + + # batch sharding + for k in batch.keys(): + batch[k] = batch[k][:, chunk_len * self.sp_rank:chunk_len * (self.sp_rank + 1)].to(self.device) + + shift_labels = batch.pop("shift_labels") + + outputs = self.forward(batch) + loss = self.compute_loss(labels=None, shift_labels=shift_labels) + + # free up temp mem (e.g. outputs.logits are huge) + del outputs + + # differentiable loss aggregation across ranks + losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=self.sp_group) + + # since each shard may have a variable number of skipped elemented - need to calculate a weighted mean depending on each rank's contribution - this will also take care of loss=0 when all elements are -100 in a shard + # XXX: not expecting a total of 0-non-skipped items for div + loss = sum(losses_per_rank[rank] * non_skipped_items[rank] + for rank in range(self.sp_world_size)) / sum(non_skipped_items.values()) + + self.backward() + + losses.append(loss.detach().item()) + + self.model.set_gradient_accumulation_boundary(True) + + # for per-iteration reporting + if len(losses) == 0: + loss = float("nan") + else: + loss = sum(losses) / len(losses) + + return loss + + def forward(self, batch): + # critical: the labels shouldn't be in batch + outputs = self.model(**batch, use_cache=False) + self.logits = outputs.logits + return outputs + + def compute_loss(self, labels, shift_labels): + if all((shift_labels == -100).squeeze()): + # this is the case where all labels in a micro-batch are -100 (very common for SFT) - CE returns `nan` in this case, so we don't want to call loss and instead create a differentiable loss `0` which will also set all the grads to `0` in `backward` - the effect of this is akin to a perfect score where the model needs no adjustment since grads will be all zeros. + # XXX: should this be float and not the original dtype? + loss = (self.logits.sum() * 0.0).float() + else: + if self.num_loss_logit_shards == "auto": + # parameterize to about 1GB fp32 logits shards + slice_size_in_gb = 1 # XXX: make configurable? + size_in_gb = self.logits.numel() * 4 / 2**30 # fp32 + # the sp shard's seqlen sp shard can be easily not divisible by the derived number of chunked loss shards, so we use the uppper ceiling and allow the last chunk to be shorter than the rest + self.num_loss_logit_shards = math.ceil(size_in_gb / slice_size_in_gb) + # print(f"derived {self.num_loss_logit_shards} shards for size {size_in_gb}GB") + if self.num_loss_logit_shards > 1: + loss = TiledLoss.apply( + self.model_unwrapped.loss_function, + self.logits, + self.model_unwrapped.config.vocab_size, + shift_labels, + self.num_loss_logit_shards, + ) + else: + # XXX: for some reason this fails with zero1 + loss = self.model_unwrapped.loss_function( + logits=self.logits, + labels=None, + vocab_size=self.model_unwrapped.config.vocab_size, + shift_labels=shift_labels, + ) + + self.loss = loss + return loss + + def backward(self): + self.model.backward(self.loss) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sparse_tensor.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sparse_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..291ba5f0c78668b793bea12749bd0a188b680ece --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/sparse_tensor.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Implementation of a compressed sparse tensor. Similar in +functionality to TensorFlow's IndexedSlices implementation. +""" + +import torch + + +class SparseTensor(object): + """ Compressed Sparse Tensor """ + + def __init__(self, dense_tensor=None): + self.orig_dense_tensor = dense_tensor + self.dtype = self.orig_dense_tensor.dtype + self.is_sparse = dense_tensor.is_sparse + if dense_tensor is not None: + if dense_tensor.is_sparse: + dense_tensor = dense_tensor.coalesce() + self.indices = dense_tensor.indices().flatten() + self.values = dense_tensor.values() + else: + result = torch.sum(dense_tensor, dim=1) + self.indices = result.nonzero().flatten() + self.values = dense_tensor[self.indices] + self.dense_size = list(dense_tensor.size()) + else: + self.indices = None + self.values = None + self.dense_size = None + + def to_coo_tensor(self): + return torch.sparse_coo_tensor(self.indices.unsqueeze(0), self.values, self.dense_size) + + @staticmethod + def type(): + return "deepspeed.SparseTensor" + + def to_dense(self): + it = self.indices.unsqueeze(1) + full_indices = torch.cat([it for _ in range(self.dense_size[1])], dim=1) + return self.values.new_zeros(self.dense_size).scatter_add_(0, full_indices, self.values) + + def sparse_size(self): + index_size = list(self.indices.size()) + index_size = index_size[0] + value_size = list(self.values.size()) + value_size = value_size[0] * value_size[1] + dense_size = self.dense_size[0] * self.dense_size[1] + return index_size + value_size, dense_size + + def add(self, b): + assert self.dense_size == b.dense_size + self.indices = torch.cat([self.indices, b.indices]) + self.values = torch.cat([self.values, b.values]) + + def __str__(self): + sparse_size, dense_size = self.sparse_size() + return "DeepSpeed.SparseTensor(indices_size={}, values_size={}, " \ + "dense_size={}, device={}, reduction_factor={})".format( + self.indices.size(), self.values.size(), self.dense_size, + self.indices.get_device(), dense_size / sparse_size + ) + + def __repr__(self): + return self.__str__() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..c2db85d1ba58f16dbe59b1dcd3d3066ab11ff6d2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py @@ -0,0 +1,427 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import os +import copy +import collections +import json +from abc import ABC, abstractmethod + +from deepspeed.utils import logger +from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine + +from .weight_quantizer import WeightQuantization + +AUTO_MODULE_KEY = 'auto' + + +class SDLoaderFactory: + + @staticmethod + def get_sd_loader_json(json_file, checkpoint_engine): + if isinstance(json_file, str): + with open(json_file) as f: + data = json.load(f) + else: + assert isinstance(json_file, dict) + data = json_file + sd_type = data['type'] + ckpt_list = data['checkpoints'] + version = data['version'] + ckpt_type = data.get('parallelization', 'pp') + mp_size = data.get('mp_size', 0) + if sd_type.lower() in ['bloom', 'ds_model']: + return data + return SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine, sd_type, version) + + @staticmethod + def get_sd_loader(ckpt_list, checkpoint_engine, sd_type='Megatron', version=None): + if sd_type == 'Megatron': + return MegatronSDLoader(ckpt_list, version, checkpoint_engine) + else: + assert False, '{} checkpoint type is not supported'.format(sd_type) + + +class SDLoaderBase(ABC): + + def __init__(self, ckpt_list, version, checkpoint_engine): + self.module_key = None + self.ckpt_list = ckpt_list + self.version = version + self.checkpoint_engine = TorchCheckpointEngine() if checkpoint_engine is None else checkpoint_engine + self.check_ckpt_list() + + def load(self, + mp_world_size, + mp_rank, + module_key=AUTO_MODULE_KEY, + is_pipe_parallel=False, + quantize=False, + quantize_bits=8, + quantize_groups=64, + mlp_extra_grouping=True): + self.module_key = module_key + num_ckpt = len(self.ckpt_list) + idx = mp_rank * num_ckpt // mp_world_size + """ We have multiple cases to handle here for both training and inference: + 1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None + a. if no mp_size/pp_size resizing occurs, for both training & inference, loading + the mp_rank related checkpoint directly. + b. if has mp_size/pp_size resizing, only Megatron model inference is supported, + in this case each mp_rank_*.pt have same content, we will load the first checkpoint + file (idx=0), to avoid idx exceeding file list boundary. + + 2. PipeModule loading layer_*.pt files, is_pipe_parallel=True, module_key is None + a. if no mp_size resizing occurs, for both training & inference, loading + the mp_rank related checkpoint directly. + b. if has mp_size resizing, only Megatron model inference is supported, + checkpoint file(s) will be merged/split according to mp_rank, mp_world_size and + checkpoint file list. + + 3. Non-PipeModule loading mp_rank_*.pt files, is_pipe_parallel=False + Same with case (2). + """ + if is_pipe_parallel and module_key is not None and mp_world_size != num_ckpt: + mp_world_size = num_ckpt + idx = 0 + + load_path = self.ckpt_list[idx] + + merge_count = 1 + if num_ckpt == mp_world_size: + assert os.path.exists(load_path) + #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}') + sd = self.checkpoint_engine.load(load_path, map_location=lambda storage, \ + loc: storage) + + if quantize: + quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) + sd_module, all_scales = quantizer.sd_quantize_megatron(self.get_module(sd), quantize_bits, + quantize_groups) + self.set_module(sd, sd_module) + else: + all_scales = None + elif num_ckpt > mp_world_size: + sd, all_scales, merge_count = self.merge_state_dict(mp_world_size, mp_rank, quantize, \ + quantize_bits, quantize_groups, mlp_extra_grouping) + else: + sd, all_scales = self.split_state_dict(mp_world_size, mp_rank, quantize, quantize_bits, \ + quantize_groups, mlp_extra_grouping) + return load_path, sd, (all_scales, merge_count) + + def get_merge_state_dicts(self, mp_world_size, mp_rank): + num_ckpt = len(self.ckpt_list) + assert num_ckpt % mp_world_size == 0, 'Invalid checkpoints and world size for sd merge' + + num_to_merge = num_ckpt // mp_world_size + ckpt_list = [self.ckpt_list[i] for i in range(num_to_merge * mp_rank, num_to_merge * (mp_rank + 1))] + + logger.info(f"mp_rank: {mp_rank}, ckpt_list: {ckpt_list}") + sd_list = [self.checkpoint_engine.load(ckpt, map_location=lambda storage, loc: storage) for ckpt in ckpt_list] + return sd_list + + def get_split_state_dict(self, mp_world_size, mp_rank): + num_ckpt = len(self.ckpt_list) + assert mp_world_size % num_ckpt == 0, 'Invalid checkpoints and world size for sd split' + + num_to_split = mp_world_size // num_ckpt + ckpt_index = mp_rank // num_to_split + ckpt_offset = mp_rank % num_to_split + + logger.info(f"mp_rank: {mp_rank}, ckpt_list: {self.ckpt_list[ckpt_index]}, offset: {ckpt_offset}") + + sd = self.checkpoint_engine.load(self.ckpt_list[ckpt_index], map_location=lambda storage, loc: storage) + + return sd, num_to_split, ckpt_offset + + def _choose_module_key(self, sd): + assert not ('module' in sd + and 'model' in sd), "checkpoint has both 'model' and 'module' keys, not sure how to proceed" + assert 'module' in sd or 'model' in sd, "checkpoint contains neither 'model' or 'module' keys, not sure how to proceed" + if 'module' in sd: + return 'module' + elif 'model' in sd: + return 'model' + + def get_module(self, sd): + if self.module_key is None: + return sd + elif self.module_key == AUTO_MODULE_KEY: + return sd[self._choose_module_key(sd)] + else: + return sd[self.module_key] + + def set_module(self, sd, module): + if self.module_key is None: + sd = module + elif self.module_key == AUTO_MODULE_KEY: + sd[self._choose_module_key(sd)] = module + else: + sd[self.module_key] = module + return sd + + def check_ckpt_list(self): + #logger.info(f'checkpoint file list: {self.ckpt_list}') + assert len(self.ckpt_list) > 0 + + sd = self.checkpoint_engine.load(self.ckpt_list[0], map_location=lambda storage, loc: storage) + + # check checkpoint count is same with saved mp_world_size + if 'mp_world_size' in sd.keys(): + assert len(self.ckpt_list) == sd[ + 'mp_world_size'], f"checkpoint count {len(self.ckpt_list)} is different from saved mp_world_size {sd['mp_world_size']}" + + @abstractmethod + def merge_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping): + pass + + @abstractmethod + def split_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping): + pass + + @abstractmethod + def sanity_check(self, ckpt_file_name): + pass + + +class MegatronSDLoader(SDLoaderBase): + + def __init__(self, ckpt_list, version, checkpoint_engine): + super().__init__(ckpt_list, version, checkpoint_engine) + """ + ## Q/K/V data need special processing + key: transformer.layers.0.attention.query_key_value.weight, shape: torch.Size([3192, 4256]) + key: transformer.layers.0.attention.query_key_value.bias, shape: torch.Size([3192]) + + ## merge or split on axis=0 + key: word_embeddings.weight, shape: torch.Size([12672, 4256]) + key: transformer.layers.0.mlp.dense_h_to_4h.bias, shape: torch.Size([4256]) + key: transformer.layers.0.mlp.dense_h_to_4h.weight, shape: torch.Size([4256, 4256]) + + ## merge or split on axis=1 + key: transformer.layers.0.attention.dense.weight, shape: torch.Size([4256, 1064]) + key: transformer.layers.0.mlp.dense_4h_to_h.weight, shape: torch.Size([4256, 4256]) + + ## no change required + key: transformer.layers.0.mlp.dense_4h_to_h.bias, shape: torch.Size([4256]) + key: transformer.final_layernorm.weight, shape: torch.Size([4256]) + key: transformer.final_layernorm.bias, shape: torch.Size([4256]) + key: transformer.layers.0.attention.dense.bias, shape: torch.Size([4256]) + key: transformer.layers.0.post_attention_layernorm.weight, shape: torch.Size([4256]) + key: transformer.layers.0.post_attention_layernorm.bias, shape: torch.Size([4256]) + key: transformer.layers.0.input_layernorm.weight, shape: torch.Size([4256]) + key: transformer.layers.0.input_layernorm.bias, shape: torch.Size([4256]) + key: position_embeddings.weight, shape: torch.Size([1024, 4256]) + """ + + def merge_query_key_value(self, param_list, ckpt_ver): + """ + Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions: + + 1. version 0, there is no version information saved in checkpoint. + format: [(3 * np * hn), h] + 2. version 1.0 + format: [(np * hn * 3), h] + 3. version 2.0 + format: [(np * 3 * hn), h] + + h: hidden size + n: number of attention heads + p: number of model parallel partitions + np: n/p + hn: h/n + """ + + new_qkv = None + if ckpt_ver == 0: + # [(3 * np * hn), h] + assert param_list[0].shape[0] % 3 == 0 + size_qkv = param_list[0].shape[0] // 3 + split_tensors = [torch.split(param, size_qkv, dim=0) for param in param_list] + + tensors = [] + for i in range(3): + tensor_tuple = [t[i] for t in split_tensors] + tensors.append(torch.cat(tensor_tuple, axis=0)) + new_qkv = torch.cat(tensors, axis=0) + elif ckpt_ver == 1.0 or ckpt_ver == 2.0: + # [(np * hn * 3), h] or [(np * 3 * hn), h] + new_qkv = torch.cat(param_list, axis=0) + else: + assert False, f'checkpoint version: {ckpt_ver} is not supported' + + return new_qkv + + def split_query_key_value(self, param, num_to_split, offset, ckpt_ver): + """ + Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions: + + 1. version 0, there is no version information saved in checkpoint. + format: [(3 * np * hn), h] + 2. version 1.0 + format: [(np * hn * 3), h] + 3. version 2.0 + format: [(np * 3 * hn), h] + + h: hidden size + n: number of attention heads + p: number of model parallel partitions + np: n/p + hn: h/n + """ + + new_qkv = None + if ckpt_ver == 0: + # [(3 * np * hn), h] + assert param.shape[0] % 3 == 0 + size_qkv = param.shape[0] // 3 + split_tensors = torch.split(param, size_qkv, dim=0) + + assert split_tensors[0].shape[0] % num_to_split == 0 + split_size = split_tensors[0].shape[0] // num_to_split + + tensors = [] + for i in range(3): + tensors.append(torch.split(split_tensors[i], split_size, dim=0)[offset]) + new_qkv = torch.cat(tensors, axis=0) + elif ckpt_ver == 1.0 or ckpt_ver == 2.0: + # [(np * hn * 3), h] or [(np * 3 * hn), h] + assert param.shape[0] % num_to_split == 0 + size_qkv = param.shape[0] // num_to_split + split_tensors = torch.split(param, size_qkv, dim=0) + new_qkv = split_tensors[offset] + else: + assert False, f'checkpoint version: {ckpt_ver} is not supported' + + return new_qkv + + def merge_state_dict(self, + mp_world_size, + mp_rank, + quantize=False, + quantize_bits=8, + groups=64, + mlp_extra_grouping=True): + self.sanity_check(self.ckpt_list[0]) + + sd_list = self.get_merge_state_dicts(mp_world_size, mp_rank) + ds_sd = copy.deepcopy(sd_list[0]) + new_client_sd = collections.OrderedDict() + + client_sd_list = [self.get_module(sd) for sd in sd_list] + keys = client_sd_list[0].keys() + + ckpt_ver = self.get_checkpoint_version(ds_sd) + logger.info(f"checkpoint version: {ckpt_ver}") + if quantize: + quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) + + for key in keys: + value_list = [sd[key] for sd in client_sd_list] + + if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key: + if quantize: + value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key, merge_dim=1) + new_client_sd[key] = torch.cat(value_list, axis=1) + elif "attention.query_key_value" in key: + if quantize and "attention.query_key_value.weight" in key: + value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key) + new_client_sd[key] = torch.cat(value_list, axis=0) + else: + if quantize: + new_client_sd[key] = torch.cat(value_list, axis=0) + else: + new_client_sd[key] = self.merge_query_key_value(value_list, ckpt_ver) + elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key: + if quantize and "mlp.dense_h_to_4h.weight" in key: + value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key) + new_client_sd[key] = torch.cat(value_list, axis=0) + else: + new_client_sd[key] = value_list[0] + if quantize: + all_scales = quantizer.merge_scales() + ds_sd = self.set_module(ds_sd, new_client_sd) + + return ds_sd, (all_scales if quantize else None), len(client_sd_list) + + def split_state_dict(self, + mp_world_size, + mp_rank, + quantize=False, + quantize_bits=8, + groups=64, + mlp_extra_grouping=True): + #self.sanity_check(self.ckpt_list[0]) + + sd, num_to_split, ckpt_offset = self.get_split_state_dict(mp_world_size, mp_rank) + ds_sd = copy.deepcopy(sd) + new_client_sd = collections.OrderedDict() + + client_sd = self.get_module(sd) + + ckpt_ver = self.get_checkpoint_version(ds_sd) + logger.info(f"checkpoint version: {ckpt_ver}") + + if quantize: + quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) + + for key in client_sd.keys(): + value = client_sd[key] + + if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key: + assert value.shape[1] % num_to_split == 0 + split_size = value.shape[1] // num_to_split + if quantize: + q_vals = quantizer.Quantize([value], quantize_bits, groups, key) + value = q_vals[0] + new_client_sd[key] = torch.split(value, split_size, dim=1)[ckpt_offset] + elif "attention.query_key_value" in key: + if quantize and "attention.query_key_value.weight" in key: + q_vals = quantizer.Quantize([value], quantize_bits, groups, key) + value = q_vals[0] + new_client_sd[key] = self.split_query_key_value(value, num_to_split, ckpt_offset, ckpt_ver) + elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key or "final_linear.weight" in key: + assert value.shape[0] % num_to_split == 0 + split_size = value.shape[0] // num_to_split + if quantize and "mlp.dense_h_to_4h.weight" in key: + q_vals = quantizer.Quantize([value], quantize_bits, groups, key) + value = q_vals[0] + new_client_sd[key] = torch.split(value, split_size, dim=0)[ckpt_offset] + else: + new_client_sd[key] = value + + if quantize: + all_scales = quantizer.merge_scales_split(num_to_split) + + ds_sd = self.set_module(ds_sd, new_client_sd) + + return ds_sd, (all_scales if quantize else None) + + def sanity_check(self, ckpt_file_name): + keys_to_check = [ + "attention.dense.weight", "mlp.dense_4h_to_h.weight", "attention.query_key_value", + "mlp.dense_h_to_4h.weight", "mlp.dense_h_to_4h.bias" + ] + + sd = self.checkpoint_engine.load(ckpt_file_name, map_location=lambda storage, loc: storage) + + # partial_key is a sub-string of one key in the sd + def check_key_exist(partial_key, sd): + keys = sd.keys() + found = False + for k in keys: + if partial_key in k: + found = True + break + return found + + for key in keys_to_check: + assert check_key_exist(key, + self.get_module(sd)), f'key: {key} is not found in the checkpoint {ckpt_file_name}' + + def get_checkpoint_version(self, state_dict): + # Use 0 if version info doesn't exist + return self.version if self.version is not None else state_dict.get('checkpoint_version', 0) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..006dfd6dcbc6f7675733ba10938552cdc13f7bcb --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .utils import MIN_SWAPPABLE_BYTES diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2293668ed88341df3f4604eac7d7cdf0ca1c5056 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2263b50560c511b5296f3cdb55c6ded2c3062de9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..262f5bd2bae2b7b124a4196b191ba501f58dfcca Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..992b88f494df544aca5df256d8044ece7ff7484f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5966fd5b9962585be989925b335d63279e85f64d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c4c0fb6bb5dcd81c69334166b0cc8cdb3d5f376 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f32fa9b9daf3f4ca3c6ff531dd7f75a1b08bccd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de18b50786787b23c393e2fe9b059378a53219e4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0b8c185cf79821179ca005914c3e87e6e94084b Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/aio_config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/aio_config.py new file mode 100644 index 0000000000000000000000000000000000000000..be6c7d93c86a4f5e268391f717217db1a687d0dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/aio_config.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.accelerator import get_accelerator + +AIO_DEFAULT_DICT = { + AIO_BLOCK_SIZE: AIO_BLOCK_SIZE_DEFAULT, + AIO_QUEUE_DEPTH: AIO_QUEUE_DEPTH_DEFAULT, + AIO_INTRA_OP_PARALLELISM: AIO_INTRA_OP_PARALLELISM_DEFAULT, + AIO_SINGLE_SUBMIT: AIO_SINGLE_SUBMIT_DEFAULT, + AIO_OVERLAP_EVENTS: AIO_OVERLAP_EVENTS_DEFAULT, + AIO_USE_GDS: AIO_USE_GDS_DEFAULT +} + + +def get_aio_config(param_dict): + if AIO in param_dict.keys() and param_dict[AIO] is not None: + aio_dict = param_dict[AIO] + aio_config = { + AIO_BLOCK_SIZE: + get_scalar_param(aio_dict, AIO_BLOCK_SIZE, AIO_BLOCK_SIZE_DEFAULT), + AIO_QUEUE_DEPTH: + get_scalar_param(aio_dict, AIO_QUEUE_DEPTH, AIO_QUEUE_DEPTH_DEFAULT), + AIO_INTRA_OP_PARALLELISM: + get_scalar_param(aio_dict, AIO_INTRA_OP_PARALLELISM, AIO_INTRA_OP_PARALLELISM_DEFAULT), + AIO_SINGLE_SUBMIT: + get_scalar_param(aio_dict, AIO_SINGLE_SUBMIT, AIO_SINGLE_SUBMIT_DEFAULT), + AIO_OVERLAP_EVENTS: + get_scalar_param(aio_dict, AIO_OVERLAP_EVENTS, AIO_OVERLAP_EVENTS_DEFAULT), + AIO_USE_GDS: + get_scalar_param(aio_dict, AIO_USE_GDS, AIO_USE_GDS_DEFAULT) + } + + if aio_config[AIO_USE_GDS]: + assert get_accelerator().device_name() == 'cuda', 'GDS currently only supported for CUDA accelerator' + + return aio_config + + return AIO_DEFAULT_DICT diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..b808721537fef5e3907bd3b8819b2e2d0b2f6936 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py @@ -0,0 +1,175 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" +import torch + +from deepspeed import comm as dist +from deepspeed.utils.logging import logger +from deepspeed.runtime.swap_tensor.utils import swap_out_tensors, SwapBuffer +from deepspeed.accelerator import get_accelerator + +INVALID_BUFFER_INDEX = -1 +ASYNC_SWAPPER_WAIT_TIMER = 'async_swap_gradient_wait' + + +class AsyncTensorSwapper(object): + + def __init__(self, aio_handle, numel_alignment, timers): + self.free_buffer_index = [] + self.swapping_buffer_index = [] + self.ready_buffer_index = [] + self.current_buffer_index = INVALID_BUFFER_INDEX + self.all_buffers = [] + self.aio_handle = aio_handle + self.numel_alignment = numel_alignment + self.max_numel = 0 + self.num_pending_swaps = 0 + self.timers = timers + self.timer_names = set() + self.num_elements_swapped = 0 + self.dtype = None + + def has_buffers(self): + return len(self.all_buffers) > 0 + + def add_buffers(self, buffer_list): + assert len(self.all_buffers) == 0 + assert all([get_accelerator().is_pinned(buffer) for buffer in buffer_list]) + dtype = buffer_list[0].dtype + assert all([buffer.dtype == dtype for buffer in buffer_list]) + + self.dtype = dtype + self.all_buffers = [SwapBuffer(buffer) for buffer in buffer_list] + self.free_buffer_index += [i for i in range(len(self.all_buffers))] + self.max_numel = max([buffer.numel() for buffer in buffer_list]) + self.timer_names = set() + + def get_timer_names(self): + return list(self.timer_names) + + def release_buffers(self): + self._report_statistics('Swapped out[Before flush]') + self._flush_buffers_until_complete() + self._report_statistics('Swapped out[After flush]') + + pinned_buffers = [buf.buffer for buf in self.all_buffers] + self.all_buffers = [] + self.free_buffer_index = [] + self.current_buffer_index = INVALID_BUFFER_INDEX + self.num_elements_swapped = 0 + self.dtype = None + + return pinned_buffers + + def swap_out_tensors(self, tensor_list, path_list): + for tensor, swap_path in zip(tensor_list, path_list): + self._swap_out_tensor(tensor, swap_path) + + def _report_statistics(self, message): + if dist.get_rank() == 0: + element_size = torch.tensor([], dtype=self.dtype).element_size() + swapped_GB = (self.num_elements_swapped * element_size) / (1024**3) + logger.debug(f'{message} num_elems = {self.num_elements_swapped}, {swapped_GB:5.2f} GB') + + def _swap_out_tensor(self, tensor, swap_path): + assert len(self.all_buffers) > 0 + + aligned_numel = self._io_aligned_numel(tensor.numel()) + assert aligned_numel <= self.max_numel + + self._make_swap_space(aligned_numel) + assert self.current_buffer_index != INVALID_BUFFER_INDEX + + swap_buffer = self._get_current_buffer() + swap_buffer.insert_tensor(tensor, swap_path, aligned_numel) + + def _make_swap_space(self, numel): + if self.current_buffer_index == INVALID_BUFFER_INDEX: + self._allocate_buffer() + return + + if not self._get_current_buffer().has_space(numel): + if len(self.free_buffer_index) > 0: + self._flush_ready_buffers() + else: + self._flush_buffers_until_complete() + self._allocate_buffer() + + def _io_aligned_numel(self, numel): + remainder = numel % self.numel_alignment + return numel if remainder == 0 else (numel + self.numel_alignment - remainder) + + def _allocate_buffer(self): + assert self.current_buffer_index == INVALID_BUFFER_INDEX + assert len(self.all_buffers) > 0 + assert len(self.free_buffer_index) > 0 + self.current_buffer_index = self.free_buffer_index[-1] + self.free_buffer_index = self.free_buffer_index[:-1] + + def _flush_ready_buffers(self): + if self.current_buffer_index != INVALID_BUFFER_INDEX: + self.ready_buffer_index.append(self.current_buffer_index) + self.current_buffer_index = INVALID_BUFFER_INDEX + + self._swap_out_ready_buffers() + + def _flush_buffers_until_complete(self): + self._flush_ready_buffers() + assert len(self.ready_buffer_index) == 0 + + self._wait_for_swap_complete() + assert len(self.swapping_buffer_index) == 0 + assert len(self.free_buffer_index) == len(self.all_buffers) + + def _swap_out_ready_buffers(self): + for buffer_index in self.ready_buffer_index: + buffer = self._get_buffer(buffer_index) + swap_tensors = buffer.get_swap_tensors() + swap_paths = buffer.get_swap_paths() + self.num_pending_swaps += len(swap_tensors) + swap_out_tensors(self.aio_handle, swap_tensors, swap_paths) + + self.swapping_buffer_index += self.ready_buffer_index + self.ready_buffer_index = [] + + def _wait_for_swap_complete(self): + assert len(self.swapping_buffer_index) > 0 + + self._start_timer(ASYNC_SWAPPER_WAIT_TIMER) + assert self.aio_handle.wait() == self.num_pending_swaps + self._stop_timer(ASYNC_SWAPPER_WAIT_TIMER) + self.timer_names.add(ASYNC_SWAPPER_WAIT_TIMER) + + self.num_pending_swaps = 0 + + for buffer_index in self.swapping_buffer_index: + buffer = self._get_buffer(buffer_index) + self.num_elements_swapped += buffer.get_num_elem() + buffer.reset() + + self.free_buffer_index += self.swapping_buffer_index + assert len(self.free_buffer_index) <= len(self.all_buffers) + self.swapping_buffer_index = [] + + def _get_buffer(self, index): + assert index != INVALID_BUFFER_INDEX + return self.all_buffers[index] + + def _get_current_buffer(self): + return self._get_buffer(self.current_buffer_index) + + def _start_timer(self, name): + if self.timers: + self.timers(name).start() + + def _stop_timer(self, name): + if self.timers: + self.timers(name).stop() + + def _log_timers(self, name_list, force=False): + if self.timers and force: + self.timers.log(name_list) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..c1207749eac6915b37c9d3d01f93f2e06e99250e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/constants.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +AIO +""" +AIO_FORMAT = ''' +"aio": { + "block_size": 1048576, + "queue_depth": 8, + "intra_op_parallelism": 1, + "single_submit": false, + "overlap_events": true, + "use_gds": false +} +''' +AIO = "aio" +AIO_BLOCK_SIZE = "block_size" +AIO_BLOCK_SIZE_DEFAULT = 1048576 +AIO_QUEUE_DEPTH = "queue_depth" +AIO_QUEUE_DEPTH_DEFAULT = 8 +AIO_INTRA_OP_PARALLELISM = "intra_op_parallelism" +AIO_INTRA_OP_PARALLELISM_DEFAULT = 1 +AIO_SINGLE_SUBMIT = "single_submit" +AIO_SINGLE_SUBMIT_DEFAULT = False +AIO_OVERLAP_EVENTS = "overlap_events" +AIO_OVERLAP_EVENTS_DEFAULT = True +AIO_USE_GDS = "use_gds" +AIO_USE_GDS_DEFAULT = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6729fd28b8fc8acd33dd552d6e70d86b326d2ac2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py @@ -0,0 +1,527 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" + +import os +import torch + +from deepspeed import comm as dist +from deepspeed.utils.logging import logger +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, \ + MIN_AIO_BYTES, AIO_ALIGNED_BYTES, get_sized_buffers +from deepspeed.runtime.swap_tensor.utils import SwapBufferManager, SwapBufferPool +from deepspeed.accelerator import get_accelerator + + +class FlattenedTensorSwapInfo(object): + + def __init__(self, path, length, offset): + self.path = path + self.offset = offset + self.length = length + + +class SwapTensorContext(object): + + def __init__(self, tensor, swap_folder): + self.compute_tensor = tensor + self.swap_tensor = torch.Tensor() + self.swap_path = os.path.join(swap_folder, f'{OptimizerSwapper.parameter_id(tensor)}.tensor.swp') + + def release_memory(self): + self.compute_tensor.data = torch.Tensor() + self.swap_tensor.data = torch.Tensor() + + def set_buffers(self, compute_buffer, swap_buffer): + self.compute_tensor.data = compute_buffer.data + self.swap_tensor.data = swap_buffer.data + + +class OptimizerStateSwapInfo(object): + + def __init__(self, parameter, numel, base_folder): + self.tensors = [] + self.param_id = OptimizerSwapper.parameter_id(parameter) + self.swap_folder = base_folder + self.swapped_gradients = {} + self.unswapped_gradients = {} + self.tensor_numel = numel + self.tensor_dtype = parameter.dtype + self.tensor_device = parameter.device + self.has_state_tensors = False + self.swap_buffers = [] + self._add_tensors([parameter]) + + def numel(self): + return self.tensor_numel + + def has_gradients(self): + return bool(self.swapped_gradients) or bool(self.unswapped_gradients) + + def _add_tensors(self, tensor_list): + for t in tensor_list: + self.tensors.append(SwapTensorContext(t, self.swap_folder)) + + def add_state_tensors(self, tensor_list): + self.has_state_tensors = True + self._add_tensors(tensor_list) + + def num_tensors(self): + return len(self.tensors) + + def device(self): + return self.tensor_device + + def dtype(self): + return self.tensor_dtype + + def release_memory(self): + for t in self.tensors: + t.release_memory() + + def get_compute_tensors(self): + return [t.compute_tensor for t in self.tensors] + + def get_swap_paths(self): + return [t.swap_path for t in self.tensors] + + def get_swap_buffers_and_paths(self, pinned): + swap_buffers = [] + swap_paths = [] + select_tensors = [t for t in self.tensors if get_accelerator().is_pinned(t.compute_tensor) == pinned] + for t in select_tensors: + swap_buffers.append(t.swap_tensor if pinned else t.compute_tensor) + swap_paths.append(t.swap_path) + return swap_buffers, swap_paths + + def get_or_create_gradient_paths(self, offsets, lengths): + gradient_paths = [] + for offset, length in zip(offsets, lengths): + if not offset in self.swapped_gradients.keys(): + path = os.path.join(self.swap_folder, f'{self.param_id}_gradient_{offset}_{length}.tensor.swp') + self.swapped_gradients[offset] = FlattenedTensorSwapInfo(path, length, offset) + + gradient_paths.append(self.swapped_gradients[offset].path) + + return gradient_paths + + def set_swap_buffers(self, buffers, aligned_numel): + num_tensors = len(self.tensors) + compute_lengths = [self.numel()] * num_tensors + compute_buffers = get_sized_buffers(buffers, compute_lengths) + swap_lengths = [aligned_numel] * num_tensors + swap_buffers = get_sized_buffers(buffers, swap_lengths) + + for i, t in enumerate(self.tensors): + t.set_buffers(compute_buffer=compute_buffers[i], swap_buffer=swap_buffers[i]) + + def get_swap_gradient_buffers(self, swap_buffer): + assert self.numel() <= swap_buffer.numel() + return [swap_buffer.narrow(0, grad.offset, grad.length) for grad in self.swapped_gradients.values()] + + def get_swap_gradient_paths(self): + return [grad.path for grad in self.swapped_gradients.values()] + + def get_unpinned_state_tensors(self): + return [t.compute_tensor for t in self.tensors if not get_accelerator().is_pinned(t.compute_tensor)] + + def read_unswapped_gradients(self, dest_buffer): + num_elem_count = 0 + for offset, grad_partition in self.unswapped_gradients.items(): + dst_tensor = dest_buffer.narrow(0, offset, grad_partition.numel()) + dst_tensor.data.copy_(grad_partition.data) + num_elem_count += grad_partition.numel() + + return num_elem_count + + def write_unswapped_gradients(self, src_buffer): + num_elem_count = 0 + for offset, grad_partition in self.unswapped_gradients.items(): + src_tensor = src_buffer.narrow(0, offset, grad_partition.numel()) + grad_partition.data.copy_(src_tensor.data) + num_elem_count += grad_partition.numel() + + return num_elem_count + + def release_unswapped_gradients(self): + self.unswapped_gradients = {} + + +SWAPPER_DEBUG_MODE = False +SWAP_OUT_GRADIENT_TIMER = 'swap_out_gradient' + + +class OptimizerSwapper(object): + + @staticmethod + def parameter_id(param): + return param.ds_id + + def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): + self.swap_config = swap_config + self.aio_config = aio_config + + # NVMe swap management + self.swap_params_info = {} + self.swap_element_size = torch.tensor([], dtype=dtype).element_size() + self.swap_folder = os.path.join(base_folder, 'optimizer', f'rank{dist.get_rank()}') + os.makedirs(self.swap_folder, exist_ok=True) + + self.optimizer = optimizer + + # Read/Write alignment for each thread during Intra-request parallelism + self.min_aio_bytes = max(MIN_AIO_BYTES, aio_config[AIO_BLOCK_SIZE]) + self.aligned_bytes = AIO_ALIGNED_BYTES * aio_config[AIO_INTRA_OP_PARALLELISM] + self.numel_alignment = self.aligned_bytes // self.swap_element_size + + # Swap buffer management + self.largest_numel = self._io_aligned_numel(largest_numel) + self.dtype = dtype + self.swap_buffer_manager = SwapBufferManager(num_elems=self.largest_numel, + count=swap_config.buffer_count, + dtype=dtype) + + # Timers + self.timers = timers + self.timer_names = set() + + # Print exclusion list + self.print_exclude_list = [ + 'optimizer', + 'swap_buffer_manager', + 'swap_params_info', + 'timers', + 'timer_names', + ] + + def purge_state(self): + for swap_info in self.swap_params_info.values(): + swap_info.tensors = [swap_info.tensors[0]] + swap_info.has_state_tensors = False + + def is_swappable_tensor(self, tensor=None, numel=None): + assert tensor is not None or numel is not None, "Either tensor or numel must be provided" + if tensor is not None: + return self.min_aio_bytes <= (tensor.numel() * self.swap_element_size) + return self.min_aio_bytes <= (numel * self.swap_element_size) + + def init_timers(self): + self.timer_names = set() + + def log_timers(self): + if self.timer_names: + self._log_timers(list(self.timer_names), force=True) + + def pre_backward(self): + self.init_timers() + + def post_backward(self): + pass + + def _flush_gradient_swapper(self, gradient_swapper): + if gradient_swapper.has_buffers(): + self._start_timer(SWAP_OUT_GRADIENT_TIMER) + pinned_buffers = gradient_swapper.release_buffers() + self.swap_buffer_manager.free(pinned_buffers) + self._stop_timer(SWAP_OUT_GRADIENT_TIMER) + self.timer_names.add(SWAP_OUT_GRADIENT_TIMER) + self.timer_names.update(gradient_swapper.get_timer_names()) + + def _swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors, gradient_swapper): + if not OptimizerSwapper.parameter_id(parameter) in self.swap_params_info.keys(): + return + + swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)] + + swappable_tensors = [] + swappable_offsets = [] + swappable_lengths = [] + + aligned_gradients, aligned_offsets = self._adjust_for_misaligned_lengths(tensors=gradient_tensors, + offsets=gradient_offsets) + + self._start_timer(SWAP_OUT_GRADIENT_TIMER) + for tensor, offset in zip(aligned_gradients, aligned_offsets): + if not self.is_swappable_tensor(tensor=tensor): + swap_info.unswapped_gradients[offset] = tensor + continue + + swappable_tensors.append(tensor) + swappable_offsets.append(offset) + swappable_lengths.append(tensor.numel()) + + if len(swappable_tensors) > 0: + if not gradient_swapper.has_buffers(): + pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + + gradient_swapper.add_buffers(pinned_buffers) + + swappable_paths = swap_info.get_or_create_gradient_paths(swappable_offsets, swappable_lengths) + + gradient_swapper.swap_out_tensors(tensor_list=swappable_tensors, path_list=swappable_paths) + + self._stop_timer(SWAP_OUT_GRADIENT_TIMER) + self.timer_names.add(SWAP_OUT_GRADIENT_TIMER) + + def _initialize_from_swapped_fp16_params(self, aio_handle, fp16_partitions_info, fp16_num_elems, + fp16_pinned_buffers, fp32_parameters): + assert len(fp32_parameters) == len(fp16_partitions_info) + assert len(fp32_parameters) == len(fp16_num_elems) + assert all([get_accelerator().is_pinned(buffer) for buffer in fp16_pinned_buffers]) + + fp32_swap_paths = self._get_swap_paths(parameters=fp32_parameters, num_elems=fp16_num_elems) + + fp32_pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + + fp16_buffer_numel = [buf.numel() for buf in fp16_pinned_buffers] + assert all([numel >= self.largest_numel for numel in fp16_buffer_numel]), \ + f"numel of fp16 buffers {fp16_buffer_numel} is too small for initializing fp32 params {self.largest_numel}" + + fp32_swap_buffers = SwapBufferPool(fp32_pinned_buffers) + fp16_swap_buffers = SwapBufferPool(fp16_pinned_buffers) + + curr_index = 0 + while curr_index < len(fp32_parameters): + fp16_pinned_tensors = self._swap_in_fp16_params(aio_handle=aio_handle, + fp16_num_elems=fp16_num_elems[curr_index:], + fp16_partitions_info=fp16_partitions_info[curr_index:], + fp16_swap_buffers=fp16_swap_buffers) + + if dist.get_rank() == 0 and SWAPPER_DEBUG_MODE: + for i, tensor in enumerate(fp16_pinned_tensors): + true_index = curr_index + i + logger.info( + f'swap_in_fp16_param: fp32_id = {OptimizerSwapper.parameter_id(fp32_parameters[true_index])} index = {true_index} orig_num_elem = {fp16_num_elems[true_index]}, swap_num_elem = {fp16_pinned_tensors[i].numel()}' + ) + + swap_out_count = self._swap_out_fp16_params(aio_handle=aio_handle, + fp32_swap_paths=fp32_swap_paths[curr_index:], + fp32_swap_buffers=fp32_swap_buffers, + fp16_pinned_tensors=fp16_pinned_tensors) + assert swap_out_count == len(fp16_pinned_tensors), \ + f"{swap_out_count} does not match {len(fp16_pinned_tensors)}" + + fp16_swap_buffers.reset() + fp32_swap_buffers.reset() + curr_index += swap_out_count + + self.swap_buffer_manager.free(fp32_pinned_buffers) + + def _swap_in_fp16_params(self, aio_handle, fp16_num_elems, fp16_partitions_info, fp16_swap_buffers): + assert len(fp16_num_elems) > 0 + + swapped_fp16_tensors = [] + swap_tensors = [] + swap_paths = [] + unswapped_srcs = [] + unswapped_dsts = [] + + for i, numel in enumerate(fp16_num_elems): + pinned_tensor, _ = fp16_swap_buffers.allocate_tensor(numel, None, numel) + if pinned_tensor is None: + break + + swapped_fp16_tensors.append(pinned_tensor) + offset = 0 + for tensor, partition_numel, partition_path in fp16_partitions_info[i]: + dst_tensor = pinned_tensor.narrow(0, offset, partition_numel) + if partition_path is None: + unswapped_srcs.append(tensor) + unswapped_dsts.append(dst_tensor) + else: + swap_paths.append(partition_path) + swap_tensors.append(dst_tensor) + offset += partition_numel + + assert len(swapped_fp16_tensors) + len(unswapped_srcs) > 0 + ret = swap_in_tensors(aio_handle, swap_tensors, swap_paths) + for src, dst in zip(unswapped_srcs, unswapped_dsts): + dst.data.copy_(src.data) + + assert len(swap_tensors) == aio_handle.wait() + + return swapped_fp16_tensors + + def _swap_out_fp16_params(self, aio_handle, fp32_swap_paths, fp32_swap_buffers, fp16_pinned_tensors): + + assert len(fp16_pinned_tensors) <= len(fp32_swap_paths) + swap_out_count = 0 + for i, fp16_tensor in enumerate(fp16_pinned_tensors): + if not fp32_swap_buffers.has_space(fp16_tensor.numel()): + fp32_swap_buffers.swap_out(aio_handle) + fp32_swap_buffers.reset() + + pinned_tensor, _ = fp32_swap_buffers.insert_tensor(fp16_tensor, fp32_swap_paths[i], + self._io_aligned_numel(fp16_tensor.numel())) + assert pinned_tensor is not None + swap_out_count += 1 + + if len(fp32_swap_buffers.get_swap_tensors()) > 0: + fp32_swap_buffers.swap_out(aio_handle) + + return swap_out_count + + def _initialize_parameters(self, parameters, src_tensors, aio_handle): + assert len(parameters) == len(src_tensors) + + swap_paths = self._get_swap_paths(parameters=parameters, num_elems=[src.numel() for src in src_tensors]) + + SWAP_INIT_TIMER = "swap_init_write" + self._start_timer(SWAP_INIT_TIMER) + + pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + assert pinned_buffers is not None + + self._swap_out_unpinned_tensors(aio_handle=aio_handle, + unpinned_tensors=src_tensors, + dest_paths=swap_paths, + pinned_buffers=pinned_buffers) + + if dist.get_rank() == 0 and SWAPPER_DEBUG_MODE: + for i, tensor in enumerate(src_tensors): + logger.info( + f'copy_in_fp16_param: fp32_id = {OptimizerSwapper.parameter_id(parameters[i])} index = {i}, swap_num_elem = {src_tensors[i].numel()}' + ) + + self.swap_buffer_manager.free(pinned_buffers) + + self._stop_timer(SWAP_INIT_TIMER) + self._log_timers([SWAP_INIT_TIMER]) + + def _get_swap_paths(self, parameters, num_elems): + swap_info_list = [ + self._create_param_swap_info(parameter=p, + numel=numel) \ + for p, numel in zip(parameters, num_elems) + ] + assert len(swap_info_list) == len(num_elems) + + swap_paths = [info.tensors[0].swap_path for info in swap_info_list] + return swap_paths + + def _swap_out_unpinned_tensors(self, aio_handle, unpinned_tensors, dest_paths, pinned_buffers): + + swap_buffer_count = len(pinned_buffers) + unpinned_tensor_count = len(unpinned_tensors) + + for i in range(0, unpinned_tensor_count, swap_buffer_count): + swap_tensor_count = min((unpinned_tensor_count - i), swap_buffer_count) + + src_tensors = unpinned_tensors[i:(i + swap_tensor_count)] + compute_lengths = [t.numel() for t in src_tensors] + compute_buffers = get_sized_buffers(pinned_buffers, compute_lengths) + + for dst, src in zip(compute_buffers, src_tensors): + dst.data.copy_(src.data) + + swap_lengths = [self._io_aligned_numel(t.numel()) for t in src_tensors] + swap_buffers = get_sized_buffers(pinned_buffers, swap_lengths) + + swap_paths = dest_paths[i:(i + swap_tensor_count)] + swap_out_tensors(aio_handle, swap_buffers, swap_paths) + + assert aio_handle.wait() == swap_tensor_count + + def _adjust_for_misaligned_lengths(self, tensors, offsets): + new_tensors = [] + new_offsets = [] + + for orig_tensor, orig_offset in zip(tensors, offsets): + if not self.is_swappable_tensor(tensor=orig_tensor): + new_tensors.append(orig_tensor) + new_offsets.append(orig_offset) + continue + + remainder = orig_tensor.numel() % self.numel_alignment + if remainder == 0: + new_tensors.append(orig_tensor) + new_offsets.append(orig_offset) + continue + + # Split into two by making remainder a tensor + aligned_length = (orig_tensor.numel() // self.numel_alignment) * self.numel_alignment + new_tensors.append(orig_tensor.narrow(0, 0, aligned_length)) + new_offsets.append(orig_offset) + + # remainder tensor + new_tensors.append(orig_tensor.narrow(0, aligned_length, remainder)) + new_offsets.append(orig_offset + aligned_length) + + return new_tensors, new_offsets + + def _retrieve_unswapped_grad_partitions(self, swap_info, dest_buffer): + UNSWAPPED_READ_GRADIENTS = 'unswapped_read_gradients' + self._start_timer(UNSWAPPED_READ_GRADIENTS) + tensor_count = len(swap_info.unswapped_gradients) + num_elem_count = swap_info.read_unswapped_gradients(dest_buffer) + self._stop_timer(UNSWAPPED_READ_GRADIENTS) + self._log_timers([UNSWAPPED_READ_GRADIENTS]) + + # It should be safe to discard unswapped gradient partitions + swap_info.release_unswapped_gradients() + + if SWAPPER_DEBUG_MODE: + logger.info( + f'optimizer_retrieve_unswapped_gradients: param={swap_info.param_id} tensor_count={tensor_count} elem_count={num_elem_count}' + ) + + def _get_state_tensors(self, parameter): + if not parameter in self.optimizer.state: + return [] + + tensor_list = [] + for state_name, value in self.optimizer.state[parameter].items(): + if torch.is_tensor(value) and self.is_swappable_tensor(tensor=value): + value.ds_id = state_name + '-' + parameter.ds_id + tensor_list.append(value) + + return tensor_list + + def _update_param_state_info(self, swap_info, parameter): + if not swap_info.has_state_tensors: + state_tensors = self._get_state_tensors(parameter) + if state_tensors: + swap_info.add_state_tensors(state_tensors) + + def _create_param_swap_info(self, parameter, numel): + param_id = OptimizerSwapper.parameter_id(parameter) + assert not param_id in self.swap_params_info + + self.swap_params_info[param_id] = OptimizerStateSwapInfo(parameter=parameter, + numel=numel, + base_folder=self.swap_folder) + swap_info = self.swap_params_info[param_id] + + self._update_param_state_info(swap_info, parameter) + + return swap_info + + def _get_param_swap_info(self, parameter): + param_id = OptimizerSwapper.parameter_id(parameter) + swap_info = self.swap_params_info.get(param_id, None) + + if swap_info is not None: + self._update_param_state_info(swap_info, parameter) + + return swap_info + + def _start_timer(self, name): + if self.timers: + self.timers(name).start() + + def _stop_timer(self, name): + if self.timers: + self.timers(name).stop() + + def _log_timers(self, name_list, force=False): + if self.timers and (SWAPPER_DEBUG_MODE or force): + self.timers.log(name_list) + + def _io_aligned_numel(self, numel): + remainder = numel % self.numel_alignment + return numel if remainder == 0 else (numel + self.numel_alignment - remainder) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..52b873ba58a1f7816aae4f79527773eb26870a34 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py @@ -0,0 +1,229 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +from deepspeed.utils.logging import logger +from deepspeed.ops.op_builder import AsyncIOBuilder +from deepspeed import comm as dist + +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object, \ + get_sized_buffers +from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper +from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper +from deepspeed.accelerator import get_accelerator + +DEBUG_MODE = False + +SWAP_IN_PARAM_TIMER = 'swap_in_param' +SWAP_OUT_PARAM_TIMER = 'swap_out_param' +SWAP_IN_GRADIENT_TIMER = 'swap_in_gradient' + + +class PartitionedOptimizerSwapper(OptimizerSwapper): + + def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): + super(PartitionedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer, + largest_numel, device, dtype, timers) + + aio_op = AsyncIOBuilder().load() + self.aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + # Overlap swapping out + self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.aio_handle, + numel_alignment=self.numel_alignment, + timers=self.timers) + + self.print_exclude_list += ['aio_handle', 'gradient_swapper', 'print_exclude_list'] + + if dist.get_rank() == 0: + print_object(obj=self, name='PartitionedOptimizerSwapper', exclude_list=self.print_exclude_list) + + def initialize_parameters(self, parameters, src_tensors): + self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.aio_handle) + + def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers, + fp32_parameters): + self._initialize_from_swapped_fp16_params(aio_handle=self.aio_handle, + fp16_partitions_info=fp16_partitions_info, + fp16_num_elems=fp16_num_elems, + fp16_pinned_buffers=fp16_pinned_buffers, + fp32_parameters=fp32_parameters) + + def flush_gradients(self): + self._flush_gradient_swapper(self.gradient_swapper) + + def release_swap_buffers(self, parameter): + swap_info = self._get_param_swap_info(parameter) + if swap_info is None: + return + swap_info.release_memory() + + self.swap_buffer_manager.free(swap_info.swap_buffers) + swap_info.swap_buffers = [] + + def swap_in_optimizer_state(self, parameter, async_parameter=None): + swap_info = self._get_param_swap_info(parameter) + if swap_info is None: + return + + self._flush_gradient_swapper(self.gradient_swapper) + + required_buffer_count = swap_info.num_tensors() + (1 if swap_info.has_gradients() else 0) + aligned_numel = self._io_aligned_numel(swap_info.numel()) + pinned_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel, + count=required_buffer_count, + dtype=parameter.dtype) + assert pinned_buffers is not None + swap_info.swap_buffers = pinned_buffers.copy() + + self._start_timer(SWAP_IN_PARAM_TIMER) + self._swap_in_parameter(aio_handle=self.aio_handle, + parameter=parameter, + dest_buffers=pinned_buffers[:swap_info.num_tensors()]) + self._stop_timer(SWAP_IN_PARAM_TIMER) + self.timer_names.add(SWAP_IN_PARAM_TIMER) + + if swap_info.has_gradients(): + self._start_timer(SWAP_IN_GRADIENT_TIMER) + self._swap_in_gradients(aio_handle=self.aio_handle, parameter=parameter, dest_buffer=pinned_buffers[-1]) + self._stop_timer(SWAP_IN_GRADIENT_TIMER) + self.timer_names.add(SWAP_IN_GRADIENT_TIMER) + + def _swap_out_optimizer_state(self, swap_info): + pinned_tensors, pinned_paths = swap_info.get_swap_buffers_and_paths(True) + WRITE_TIMER = 'swap_submit_write' + self._start_timer(WRITE_TIMER) + + swap_out_tensors(self.aio_handle, pinned_tensors, pinned_paths) + assert self.aio_handle.wait() == len(pinned_tensors) + + unpinned_tensors, unpinned_paths = swap_info.get_swap_buffers_and_paths(False) + if len(unpinned_tensors) > 0: + pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + self._swap_out_unpinned_tensors(aio_handle=self.aio_handle, + unpinned_tensors=unpinned_tensors, + dest_paths=unpinned_paths, + pinned_buffers=pinned_buffers) + swap_info.swap_buffers += pinned_buffers.copy() + + self._stop_timer(WRITE_TIMER) + self._log_timers([WRITE_TIMER]) + + def writeback_optimizer_state_and_gradients(self, parameter, write_opt_state, write_gradients): + swap_info = self._get_param_swap_info(parameter=parameter) + + if swap_info is None: + return + + if write_opt_state: + self._swap_out_optimizer_state(swap_info) + + if write_gradients and swap_info.has_gradients(): + param_gradients = swap_info.swapped_gradients.values() + swap_buffers = [parameter.grad.narrow(0, grad.offset, grad.length) for grad in param_gradients] + swap_paths = [grad.path for grad in param_gradients] + swap_out_tensors(self.aio_handle, swap_buffers, swap_paths) + assert len(swap_buffers) == self.aio_handle.wait() + if swap_info.unswapped_gradients: + swap_info.write_unswapped_gradients(src_buffer=parameter.grad) + + self.release_swap_buffers(parameter) + + def swap_out_optimizer_state(self, parameter, async_swap=False): + swap_info = self._get_param_swap_info(parameter=parameter) + + if swap_info is None: + return + + swap_bytes = sum( + [self._io_aligned_numel(t.numel()) * t.element_size() for t in swap_info.get_compute_tensors()]) + + self._start_timer(SWAP_OUT_PARAM_TIMER) + self._swap_out_optimizer_state(swap_info) + self.release_swap_buffers(parameter) + self._stop_timer(SWAP_OUT_PARAM_TIMER) + self.timer_names.add(SWAP_OUT_PARAM_TIMER) + + if DEBUG_MODE and dist.get_rank() == 0: + logger.info(f'optimizer_param_swap_out: {(swap_bytes/(1024**3)):5.2f} GB') + + def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): + self._swap_out_gradients(parameter=parameter, + gradient_offsets=gradient_offsets, + gradient_tensors=gradient_tensors, + gradient_swapper=self.gradient_swapper) + + def _swap_in_parameter(self, aio_handle, parameter, dest_buffers): + swap_info = self._get_param_swap_info(parameter) + if swap_info is None: + return + + num_swap_tensors = swap_info.num_tensors() + assert num_swap_tensors <= len(dest_buffers) + + swap_lengths = [self._io_aligned_numel(swap_info.numel())] * num_swap_tensors + swap_buffers = get_sized_buffers(dest_buffers, swap_lengths) + + compute_lengths = [swap_info.numel()] * num_swap_tensors + compute_buffers = get_sized_buffers(dest_buffers, compute_lengths) + + READ_TIMER = 'swap_submit_read_param' + WAIT_TIMER = 'swap_wait_read_param' + + self._start_timer(READ_TIMER) + swap_in_tensors(aio_handle, swap_buffers, swap_info.get_swap_paths()) + self._stop_timer(READ_TIMER) + + swap_bytes = sum([buffer.numel() * buffer.element_size() for buffer in swap_buffers]) + + self._start_timer(WAIT_TIMER) + aio_handle.wait() + self._stop_timer(WAIT_TIMER) + + swap_info.set_swap_buffers(dest_buffers, self._io_aligned_numel(swap_info.numel())) + + self._log_timers([READ_TIMER, WAIT_TIMER]) + if DEBUG_MODE and dist.get_rank() == 0: + logger.info(f'optimizer_param_swap_in: {(swap_bytes/(1024**3)):5.2f} GB') + + def _swap_in_pinned_gradients(self, aio_handle, parameter, gradient_tensor): + swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)] + param_gradients = swap_info.swapped_gradients.values() + swap_buffers = [gradient_tensor.narrow(0, grad.offset, grad.length) for grad in param_gradients] + swap_paths = [grad.path for grad in param_gradients] + SWAP_READ_GRADIENTS = 'swap_submit_read_gradient' + SWAP_WAIT_GRADIENTS = 'swap_submit_wait_gradient' + self._start_timer(SWAP_READ_GRADIENTS) + swap_in_tensors(aio_handle, swap_buffers, swap_paths) + self._stop_timer(SWAP_READ_GRADIENTS) + + self._start_timer(SWAP_WAIT_GRADIENTS) + assert len(swap_buffers) == aio_handle.wait() + self._stop_timer(SWAP_WAIT_GRADIENTS) + + self._log_timers([SWAP_READ_GRADIENTS, SWAP_WAIT_GRADIENTS]) + + def _swap_in_gradients(self, aio_handle, parameter, dest_buffer): + swap_info = self.swap_params_info.get(OptimizerSwapper.parameter_id(parameter), None) + if not (swap_info and swap_info.has_gradients()): + return + + assert get_accelerator().is_pinned(dest_buffer) + assert parameter.numel() <= dest_buffer.numel() + + parameter.grad = dest_buffer.narrow(0, 0, parameter.numel()) + + if swap_info.swapped_gradients: + self._swap_in_pinned_gradients(aio_handle, parameter, parameter.grad) + + if swap_info.unswapped_gradients: + self._retrieve_unswapped_grad_partitions(swap_info=swap_info, dest_buffer=parameter.grad) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..3e00a1860b0b20522e90b9c53437207c2052c92a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py @@ -0,0 +1,419 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" + +import os +import shutil +from enum import Enum +import torch +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import AsyncIOBuilder +from deepspeed.ops.op_builder import GDSBuilder +from .constants import * +from .utils import swap_in_tensors, swap_out_tensors, MIN_AIO_BYTES, AIO_ALIGNED_BYTES, print_object, SwapBufferPool + + +def print_rank_0(message, debug=False, force=False): + if dist.get_rank() == 0 and (debug or force): + print(message) + + +class PartitionedParamStatus(Enum): + # Partitioned parameters are present and ready for use + AVAILABLE = 1 + + # partitioned params are in some non-memory device + NOT_AVAILABLE = 2 + + # partitioned params are being read from some non-memory device. + INFLIGHT = 3 + + +class AsyncPartitionedParameterSwapper(object): + + def __init__(self, ds_config, model_dtype): + + self.dtype = model_dtype + + #set swap buffers, create aio handles + self._configure_aio(ds_config) + + #mapping from param id to path + self.id_to_path = {} + + #mapping from pram_id to buffer id + self.param_id_to_buffer_id = {} + + # mapping from param_id to swap buffer + self.param_id_to_swap_buffer = {} + + #number of elements in the param + self.param_id_to_numel = {} + + self.pending_writes = 0 + self.pending_reads = 0 + + #keep track of async swap in params and buffers + self.inflight_params = [] + self.inflight_swap_in_buffers = [] + self.inflight_numel = 0 + + #keep track of available params + self.available_params = set() + self.available_numel = 0 + + # for swapping out from partitioned fp32 params + self.partitioned_swap_buffer = None + self.partitioned_swap_pool = None + + self.invalid_buffer = torch.tensor(1).half() + + if dist.get_rank() == 0: + exclude_list = ['aio_read_handle', 'aio_write_handle', 'buffers'] + print_object(obj=self, name='AsyncPartitionedParameterSwapper', exclude_list=exclude_list) + + def available_swap_in_buffers(self): + return len(self.available_buffer_ids) + + def _configure_aio(self, ds_config): + self.swap_config = ds_config.zero_config.offload_param + torch_dtype_string = str(self.dtype).split(".")[1] + self.swap_folder = os.path.join(self.swap_config.nvme_path, 'zero_stage_3', f'{torch_dtype_string}params', + f'rank{dist.get_rank()}') + shutil.rmtree(self.swap_folder, ignore_errors=True) + os.makedirs(self.swap_folder, exist_ok=True) + + self.swap_element_size = torch.tensor([], dtype=self.dtype).element_size() + + self.aio_config = ds_config.aio_config + + self.use_gds = self.aio_config[AIO_USE_GDS] + self.aio_handle = GDSBuilder().load(verbose=False).gds_handle if self.use_gds else AsyncIOBuilder().load( + verbose=False).aio_handle + + # Read/Write alignment for each thread during Intra-request parallelism + self.min_aio_bytes = max(MIN_AIO_BYTES, self.aio_config[AIO_BLOCK_SIZE]) + self.aligned_bytes = AIO_ALIGNED_BYTES * self.aio_config[AIO_INTRA_OP_PARALLELISM] + self.numel_alignment = self.aligned_bytes // self.swap_element_size + + self.elements_per_buffer = self.swap_config.buffer_size + self.aligned_elements_per_buffer = self._io_aligned_numel(self.elements_per_buffer) + self.param_buffer_count = self.swap_config.buffer_count + + self.available_buffer_ids = [i for i in range(self.param_buffer_count)] + self.reserved_buffer_ids = [] + + self.aio_read_handle = self.aio_handle(block_size=self.aio_config[AIO_BLOCK_SIZE], + queue_depth=self.aio_config[AIO_QUEUE_DEPTH], + single_submit=self.aio_config[AIO_SINGLE_SUBMIT], + overlap_events=self.aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=self.aio_config[AIO_INTRA_OP_PARALLELISM]) + + self.aio_write_handle = self.aio_handle(block_size=self.aio_config[AIO_BLOCK_SIZE], + queue_depth=self.aio_config[AIO_QUEUE_DEPTH], + single_submit=self.aio_config[AIO_SINGLE_SUBMIT], + overlap_events=self.aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=self.aio_config[AIO_INTRA_OP_PARALLELISM]) + + buffer_device = get_accelerator().device_name() if self.use_gds else "cpu" + self.buffers = torch.empty(int(self.aligned_elements_per_buffer * self.param_buffer_count), + dtype=self.dtype, + device=buffer_device, + requires_grad=False) + if self.use_gds: + self.aio_read_handle.pin_device_tensor(self.buffers) + else: + self.buffers = get_accelerator().pin_memory(self.buffers, align_bytes=0) + + self.swap_out_params = [] + + #Check if partitioned param or numel in a tensor is swappable or not + def swappable_tensor(self, param=None, numel=None): + if param is not None: + assert numel is None, "Both parma and numel cannot be provided" + numel = param.ds_tensor.ds_numel + if numel is not None: + return self.min_aio_bytes <= numel * self.swap_element_size + assert False, "Either param or numel must be provided" + + def get_path(self, param, must_exist=False): + paths = self._get_swap_paths([param], must_exist=must_exist) + return paths[0] + + def _get_swap_paths(self, params, must_exist=False): + paths = [] + for param in params: + param_id = param.ds_id + if param_id in self.id_to_path.keys(): + param_path = self.id_to_path[param_id] + else: + assert not must_exist, f"Path for param id {param_id} does not exist" + param_path = os.path.join(self.swap_folder, f'{param_id}_param.tensor.swp') + + self.id_to_path[param_id] = param_path + paths.append(param_path) + + return paths + + def _get_swap_buffers(self, params): + buffers = [] + for param in params: + param_id = param.ds_id + assert param_id in self.param_id_to_swap_buffer.keys(), \ + f'param {param_id} has not been assigned a swap buffer' + buffers.append(self.param_id_to_swap_buffer[param_id]) + + return buffers + + def _track_numel(self, params): + for param in params: + assert param.ds_tensor is not None, "Partitioned tensor is None" + self.param_id_to_numel[param.ds_id] = param.ds_tensor.ds_numel + + def _allocate_and_return_buffers_for_swap_in(self, params): + compute_buffers = [] + swap_buffers = [] + + for param in params: + param_id = param.ds_id + assert param_id in self.param_id_to_numel.keys(), f" Number of elements in param {param_id} is unknown" + assert param_id not in self.param_id_to_buffer_id.keys( + ), f"param {param_id} already assigned swap buffer id {self.param_id_to_buffer_id[param_id]}" + assert param_id not in self.param_id_to_swap_buffer.keys( + ), f"param {param_id} has already been assigned a swap buffer" + + buffer_id = self.available_buffer_ids.pop() + print_rank_0(f"param {param.ds_id} is assigned swap in buffer id {buffer_id} ") + self.param_id_to_buffer_id[param_id] = buffer_id + aligned_swap_numel = self._io_aligned_numel(self.param_id_to_numel[param_id]) + swap_buffer = self.buffers.narrow(0, int(buffer_id * self.aligned_elements_per_buffer), aligned_swap_numel) + + self.param_id_to_swap_buffer[param_id] = swap_buffer + compute_buffer = swap_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) + compute_buffers.append(compute_buffer) + swap_buffers.append(swap_buffer) + + return compute_buffers, swap_buffers + + #waits for inflight nvme write to complete + def synchronize_writes(self): + if self.pending_writes == 0: + return + assert self.pending_writes == self.aio_write_handle.wait() + self.pending_writes = 0 + self.remove_partition_and_release_buffers(self.swap_out_params) + self.swap_out_params = [] + + #waits for inflight nvme reads to complete + def synchronize_reads(self): + if self.pending_reads == 0: + return + + assert self.pending_reads == self.aio_read_handle.wait() + + self.pending_reads = 0 + + for param, swap_in_buffer in zip(self.inflight_params, self.inflight_swap_in_buffers): + param_id = param.ds_id + compute_buffer = swap_in_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) + param.ds_tensor.data = compute_buffer.data + param.ds_tensor.status = PartitionedParamStatus.AVAILABLE + + self.available_params.update([param.ds_id for param in self.inflight_params]) + self.available_numel += self.inflight_numel + + self.inflight_params = [] + self.inflight_swap_in_buffers = [] + self.inflight_numel = 0 + + #Removes the memory assignment and releases the buffers + #Should only be executed after swapping out the tensors + def remove_partition_and_release_buffers(self, params): + for param in params: + param_id = param.ds_id + + if param_id in self.param_id_to_buffer_id.keys(): + + buffer_id = self.param_id_to_buffer_id[param_id] + + assert buffer_id is not None, "Missing buffer id for releasing" + + self.available_buffer_ids.append(buffer_id) + del self.param_id_to_buffer_id[param_id] + del self.param_id_to_swap_buffer[param_id] + print_rank_0(f"param {param.ds_id} releases buffer id {buffer_id} ") + + if param_id in self.available_params: + self.available_params.remove(param_id) + self.available_numel -= self.param_id_to_numel[param_id] + + param.ds_tensor.data = self.invalid_buffer.data + param.ds_tensor.status = PartitionedParamStatus.NOT_AVAILABLE + + #writes from in memory to nvme. Does not release the buffers + def _swap_out(self, params, async_op=True): + + swap_out_paths = self._get_swap_paths(params) + swap_out_params = self._get_swap_buffers(params) + self._track_numel(params) + + swap_out_tensors(self.aio_write_handle, swap_out_params, swap_out_paths) + + self.pending_writes += len(swap_out_params) + self.swap_out_params += params + + if not async_op: + self.synchronize_writes() + + #blocking swap out followed by releasing the memory buffers + def swap_out_and_release(self, params, async_op=False, force_buffer_release=False): + if async_op: + assert force_buffer_release, "Should not release preallocated buffers without completing the swap out. Set force_buffer_release to True to do it anyways" + self._swap_out(params, async_op=async_op) + + # book keeping function for inflight swap in + def _update_inflight_swap_in(self, params, swap_in_buffers, inflight_numel): + self.inflight_params.extend(params) + self.inflight_swap_in_buffers.extend(swap_in_buffers) + self.inflight_numel += inflight_numel + + for param in params: + param.ds_tensor.status = PartitionedParamStatus.INFLIGHT + + self.pending_reads += len(params) + + #assigns an in memory buffer and swaps in from nvme + def swap_in(self, params, async_op=True, swap_in_buffers=None): + + assert all([param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE + for param in params]), "Some params are already available or in flight" + swap_in_paths = self._get_swap_paths(params) + + if swap_in_buffers is None: + if len(self.available_buffer_ids) < len(swap_in_paths): + ids = [p.ds_id for p in params] + print_rank_0( + f'Not enough swap in buffers {len(self.available_buffer_ids)} for {len(swap_in_paths)} params, ids = {ids}', + force=True) + print_rank_0( + f'Num inflight: params {len(self.inflight_params)}, buffers {len(self.inflight_swap_in_buffers)}, numel = {self.inflight_numel}', + force=True) + print_rank_0( + f'Num available params: count = {len(self.available_params)}, ids = {self.available_params}, numel = {self.available_numel}', + force=True) + + assert len(swap_in_paths) <= len( + self.available_buffer_ids + ), f"Not enough buffers {len(self.available_buffer_ids)} for swapping {len(swap_in_paths)}" + compute_buffers, swap_in_buffers = self._allocate_and_return_buffers_for_swap_in(params) + inflight_numel = sum([t.numel() for t in compute_buffers]) + else: + inflight_numel = sum([t.numel() for t in swap_in_buffers]) + + swap_in_tensors(self.aio_read_handle, swap_in_buffers, swap_in_paths) + + self._update_inflight_swap_in(params, swap_in_buffers, inflight_numel) + + if not async_op: + self.synchronize_reads() + + # Enables swapping into buffer that is out the control of swapper. This is always synchronous + def swap_into_buffer(self, param, dest_buffer): + assert param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE, f"param {param.ds_id} is already available or inflight" + + require_swap_buffer = not (get_accelerator().is_pinned(dest_buffer) + and self._is_io_aligned(dest_buffer.numel())) + + if require_swap_buffer: + assert len(self.available_buffer_ids) > 0, f"No buffer available to swap param {param.ds_id}." + compute_buffers, swap_in_buffers = self._allocate_and_return_buffers_for_swap_in([param]) + inflight_numel = compute_buffers[0].numel() + else: + swap_in_buffers = [dest_buffer] + inflight_numel = dest_buffer.numel() + + swap_in_paths = self._get_swap_paths([param]) + + swap_in_tensors(self.aio_read_handle, swap_in_buffers, swap_in_paths) + self._update_inflight_swap_in([param], swap_in_buffers, inflight_numel) + self.synchronize_reads() + + if require_swap_buffer: + dest_buffer.data.copy_(param.ds_tensor.data) + # Release swap buffer memory assignment. Note, this will mark the parameter not available. + self.remove_partition_and_release_buffers([param]) + + #assign a buffer to a param and return the buffer + def get_buffer(self, param, numel): + param_id = param.ds_id + + assert self.available_swap_in_buffers( + ) > 0, f"No swap buffers to allocate for fp16 param {param_id} of numel = {numel}" + assert numel <= self.elements_per_buffer, f"More elements {numel} than buffer size {self.elements_per_buffer}" + + self.param_id_to_numel[param_id] = numel + buffer_id = self.available_buffer_ids.pop() + self.param_id_to_buffer_id[param_id] = buffer_id + aligned_swap_numel = self._io_aligned_numel(self.param_id_to_numel[param_id]) + swap_buffer = self.buffers.narrow(0, int(buffer_id * self.aligned_elements_per_buffer), aligned_swap_numel) + + self.param_id_to_swap_buffer[param_id] = swap_buffer + compute_buffer = swap_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) + print_rank_0(f"param {param.ds_id} is assigned swap in buffer id {buffer_id}") + return compute_buffer + + def reserve_available_buffers(self): + buffers = [] + for id in self.available_buffer_ids: + buffers.append( + self.buffers.narrow(0, int(id * self.aligned_elements_per_buffer), + int(self.aligned_elements_per_buffer))) + self.reserved_buffer_ids.append(id) + + self.available_buffer_ids = [] + return buffers + + def release_reserved_buffers(self): + for id in self.reserved_buffer_ids: + self.available_buffer_ids.append(id) + self.reserved_buffer_ids = [] + + def _io_aligned_numel(self, numel): + remainder = numel % self.numel_alignment + return numel if remainder == 0 else (numel + self.numel_alignment - remainder) + + def _is_io_aligned(self, numel): + return (numel % self.numel_alignment) == 0 + + def reserve_partitioned_swap_space(self, partition_num_elems): + aligned_numel = sum([self._io_aligned_numel(numel) for numel in partition_num_elems]) + self.partitioned_swap_buffer = get_accelerator().pin_memory(torch.zeros(aligned_numel, + device='cpu', + dtype=self.dtype), + align_bytes=0) + self.partitioned_swap_pool = SwapBufferPool([self.partitioned_swap_buffer]) + + def swap_out_partitioned_params(self, dst_fp16_params, src_fp32_params): + assert self.partitioned_swap_buffer is not None, f'partitioned swap buffers for fp16 params not initialized' + assert self.partitioned_swap_pool is not None, f'partitioned swap pool for fp16 params not initialized' + assert len(dst_fp16_params) == len(src_fp32_params), \ + f'mismatch in number of fp16 params {len(dst_fp16_params)} and fp32 params {len(src_fp32_params)}' + + fp16_swap_paths = self._get_swap_paths(dst_fp16_params, must_exist=True) + self.synchronize_writes() + self.partitioned_swap_pool.reset() + for i, fp32_tensor in enumerate(src_fp32_params): + swap_tensor, _ = self.partitioned_swap_pool.insert_tensor(fp32_tensor, fp16_swap_paths[i], + self._io_aligned_numel(fp32_tensor.numel())) + assert swap_tensor is not None + dst_fp16_params[i].ds_tensor.status = PartitionedParamStatus.AVAILABLE + + self.partitioned_swap_pool.swap_out(self.aio_write_handle) + + for param in dst_fp16_params: + param.ds_tensor.status = PartitionedParamStatus.NOT_AVAILABLE diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..1ff570ed3bb9dda06c68de5be8a7c105b98694d7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py @@ -0,0 +1,242 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +from deepspeed.ops.op_builder import AsyncIOBuilder +from deepspeed import comm as dist +import torch + +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object +from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper +from deepspeed.runtime.swap_tensor.utils import get_sized_buffer +from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper + + +class OptimizerSwapOp(object): + + def __init__(self, aio_handle, read_op, param_info, allocated_buffers, state_buffers, num_ops): + self.aio_handle = aio_handle + self.read_op = read_op + self.param_info = param_info + self.allocated_buffers = allocated_buffers + self.state_buffers = state_buffers + self.wait_required = True + self.num_ops = num_ops + + def is_parameter(self, parameter): + return OptimizerSwapper.parameter_id(parameter) == self.param_info.param_id + + def wait(self): + assert self.wait_required + assert self.aio_handle.wait() == self.num_ops + self.wait_required = False + + +SYNC_SWAP_IN = 'sync_swap_in' +ASYNC_SWAP_IN = 'async_swap_in' +SYNC_SWAP_OUT = 'sync_swap_out' +ASYNC_SWAP_OUT = 'async_swap_out' + +SWAP_IN_STATE_TIMER = 'swap_in_state' +SWAP_OUT_STATE_TIMER = 'swap_out_state' +SWAP_OUT_GRADIENT_TIMER = 'swap_out_gradient' +ASYNC_SWAP_IN_STATE_TIMER = "async_swap_in_state" +ASYNC_SWAP_OUT_STATE_TIMER = 'async_swap_out_state' + + +class PipelinedOptimizerSwapper(OptimizerSwapper): + + def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): + super(PipelinedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer, largest_numel, + device, dtype, timers) + + aio_op = AsyncIOBuilder().load() + self.write_aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + self.read_aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + # Overlap gradient swap out + self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.write_aio_handle, + numel_alignment=self.numel_alignment, + timers=self.timers) + + self.async_swap_in = swap_config.pipeline_read + self.async_swap_out = swap_config.pipeline_write + + self.swap_ops = {SYNC_SWAP_IN: None, ASYNC_SWAP_IN: None, SYNC_SWAP_OUT: None, ASYNC_SWAP_OUT: None} + + self.print_exclude_list += [ + 'gradient_swapper', 'read_aio_handle', 'write_aio_handle', 'swap_ops', 'print_exclude_list' + ] + + if dist.get_rank() == 0: + print_object(obj=self, name='PipelinedOptimizerSwapper', exclude_list=self.print_exclude_list) + + def initialize_parameters(self, parameters, src_tensors): + self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.write_aio_handle) + + def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers, + fp32_parameters): + self._initialize_from_swapped_fp16_params(aio_handle=self.write_aio_handle, + fp16_partitions_info=fp16_partitions_info, + fp16_num_elems=fp16_num_elems, + fp16_pinned_buffers=fp16_pinned_buffers, + fp32_parameters=fp32_parameters) + + def flush_gradients(self): + self._flush_gradient_swapper(self.gradient_swapper) + + def swap_in_optimizer_state(self, parameter, async_parameter): + assert parameter is not None + assert self.swap_ops[SYNC_SWAP_IN] is None + + self._flush_gradient_swapper(self.gradient_swapper) + + self._start_timer(SWAP_IN_STATE_TIMER) + + if self.swap_ops[ASYNC_SWAP_IN]: + assert self.swap_ops[ASYNC_SWAP_IN].is_parameter(parameter) + self.swap_ops[SYNC_SWAP_IN] = self.swap_ops[ASYNC_SWAP_IN] + self.swap_ops[ASYNC_SWAP_IN] = None + else: + self.swap_ops[SYNC_SWAP_IN] = self._swap_in_optimizer_state(aio_handle=self.read_aio_handle, + parameter=parameter) + + if self.swap_ops[SYNC_SWAP_IN]: + self.swap_ops[SYNC_SWAP_IN].wait() + + if self.async_swap_in and async_parameter is not None: + assert self.swap_ops[ASYNC_SWAP_IN] is None + self.swap_ops[ASYNC_SWAP_IN] = self._swap_in_optimizer_state(aio_handle=self.read_aio_handle, + parameter=async_parameter) + + self._stop_timer(SWAP_IN_STATE_TIMER) + self.timer_names.add(SWAP_IN_STATE_TIMER) + + def swap_out_optimizer_state(self, parameter, async_swap): + self._start_timer(SWAP_OUT_STATE_TIMER) + + if self.swap_ops[ASYNC_SWAP_OUT]: + self._start_timer(ASYNC_SWAP_OUT_STATE_TIMER) + self._complete_swap_out(ASYNC_SWAP_OUT) + self._stop_timer(ASYNC_SWAP_OUT_STATE_TIMER) + self.timer_names.add(ASYNC_SWAP_OUT_STATE_TIMER) + + assert self.swap_ops[SYNC_SWAP_IN] is not None + assert not self.swap_ops[SYNC_SWAP_IN].wait_required + swap_op = self._swap_out_optimizer_state(aio_handle=self.write_aio_handle, + parameter=parameter, + swap_in_op=self.swap_ops[SYNC_SWAP_IN]) + self.swap_ops[SYNC_SWAP_IN] = None + + if self.async_swap_out and async_swap: + self.swap_ops[ASYNC_SWAP_OUT] = swap_op + else: + self.swap_ops[SYNC_SWAP_OUT] = swap_op + self._complete_swap_out(SYNC_SWAP_OUT) + + self._stop_timer(SWAP_OUT_STATE_TIMER) + self.timer_names.add(SWAP_OUT_STATE_TIMER) + + def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): + self._swap_out_gradients(parameter=parameter, + gradient_offsets=gradient_offsets, + gradient_tensors=gradient_tensors, + gradient_swapper=self.gradient_swapper) + + def _complete_swap_out(self, swap_out_type): + self.swap_ops[swap_out_type].wait() + for buffer in self.swap_ops[swap_out_type].state_buffers: + buffer = torch.Tensor() + self.swap_buffer_manager.free(self.swap_ops[swap_out_type].allocated_buffers) + self.swap_ops[swap_out_type] = None + + def _swap_out_optimizer_state(self, aio_handle, parameter, swap_in_op): + assert swap_in_op.is_parameter(parameter) + + allocated_buffers = swap_in_op.allocated_buffers.copy() + swap_buffers = swap_in_op.state_buffers.copy() + + param_info = swap_in_op.param_info + self._update_param_state_info(param_info, parameter) + unpinned_tensors = param_info.get_unpinned_state_tensors() + + if len(unpinned_tensors) > 0: + new_alloc_buffers = self.swap_buffer_manager.allocate(num_elems=self._io_aligned_numel(param_info.numel()), + count=len(unpinned_tensors), + dtype=param_info.dtype()) + assert new_alloc_buffers is not None + + allocated_buffers += new_alloc_buffers + swap_buffers += new_alloc_buffers + + for pinned_dst, unpinned_src in zip(new_alloc_buffers, unpinned_tensors): + dst = get_sized_buffer(pinned_dst, unpinned_src.numel()) + dst.data.copy_(unpinned_src.data) + + swap_paths = param_info.get_swap_paths() + assert len(swap_paths) == len(swap_buffers) + + swap_out_tensors(aio_handle, swap_buffers, swap_paths) + + swap_out_op = OptimizerSwapOp(aio_handle=aio_handle, + param_info=param_info, + read_op=False, + allocated_buffers=allocated_buffers, + state_buffers=swap_buffers, + num_ops=len(swap_buffers)) + + return swap_out_op + + def _swap_in_optimizer_state(self, aio_handle, parameter): + param_info = self._get_param_swap_info(parameter) + if param_info is None: + return None + + num_swap_tensors = param_info.num_tensors() + required_buffer_count = num_swap_tensors + (1 if param_info.has_gradients() else 0) + aligned_numel = self._io_aligned_numel(param_info.numel()) + allocated_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel, + count=required_buffer_count, + dtype=parameter.dtype) + assert allocated_buffers is not None, \ + f"PipelinedOptimizerSwapper ran out of swap buffers, try increasing 'buffer_count'" + + state_buffers = allocated_buffers[:num_swap_tensors] + param_info.set_swap_buffers(state_buffers, aligned_numel) + + swap_buffers = state_buffers.copy() + swap_paths = param_info.get_swap_paths() + + if param_info.has_gradients(): + parameter.grad = allocated_buffers[-1].narrow(0, 0, param_info.numel()) + if param_info.swapped_gradients: + swap_buffers += param_info.get_swap_gradient_buffers(parameter.grad) + swap_paths += param_info.get_swap_gradient_paths() + + swap_in_tensors(aio_handle, swap_buffers, swap_paths) + + if param_info.unswapped_gradients: + self._retrieve_unswapped_grad_partitions(swap_info=param_info, dest_buffer=parameter.grad) + + swap_in_op = OptimizerSwapOp(aio_handle=aio_handle, + param_info=param_info, + read_op=True, + allocated_buffers=allocated_buffers, + state_buffers=state_buffers, + num_ops=len(swap_buffers)) + + return swap_in_op diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..60a4004381d9d9080023bc843c1eb425afc51377 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/utils.py @@ -0,0 +1,241 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" + +import torch +from deepspeed.utils.logging import logger +from deepspeed.accelerator import get_accelerator + +from deepspeed import comm as dist + +MIN_AIO_BYTES = 1024**2 +AIO_ALIGNED_BYTES = 1024 +MIN_SWAPPABLE_BYTES = MIN_AIO_BYTES + + +def swap_in_tensors(swap_handle, tensor_buffers, swap_paths): + for buffer, path in zip(tensor_buffers, swap_paths): + assert (swap_handle.async_pread(buffer, path, 0) == 0) + + +def swap_out_tensors(swap_handle, tensor_buffers, swap_paths): + for buffer, path in zip(tensor_buffers, swap_paths): + assert (swap_handle.async_pwrite(buffer, path, 0) == 0) + + +def print_object(obj, name, exclude_list=[]): + logger.info('{}:'.format(name)) + for arg in sorted(vars(obj)): + if not arg in exclude_list: + dots = '.' * (29 - len(arg)) + logger.info(' {} {} {}'.format(arg, dots, getattr(obj, arg))) + + +class SwapBuffer(object): + + def __init__(self, buffer): + self.buffer = buffer + self.reset() + + def reset(self): + self.offset = 0 + self.swap_tensors = {} + self.compute_tensors = {} + self.swap_paths = {} + self.num_elem = 0 + + def insert_tensor(self, tensor, swap_path, aligned_numel): + swap_tensor, compute_tensor = self.allocate_tensor(swap_path, tensor.numel(), aligned_numel) + compute_tensor.data.copy_(tensor.data) + return swap_tensor, compute_tensor + + def allocate_tensor(self, swap_path, numel, aligned_numel): + assert self.has_space(aligned_numel) + assert not self.offset in self.swap_tensors + + allocate_offset = self.offset + swap_tensor = self.buffer.narrow(0, allocate_offset, aligned_numel) + dest_tensor = swap_tensor.narrow(0, 0, numel) + + self.swap_tensors[allocate_offset] = swap_tensor + self.compute_tensors[allocate_offset] = dest_tensor + self.swap_paths[allocate_offset] = swap_path + self.offset += aligned_numel + self.num_elem += numel + + return self.swap_tensors[allocate_offset], self.compute_tensors[allocate_offset] + + def has_space(self, numel): + return (self.offset + numel) <= self.buffer.numel() + + def get_swap_tensors(self): + return [tensor for tensor in self.swap_tensors.values()] + + def get_swap_paths(self): + return [path for path in self.swap_paths.values()] + + def get_compute_tensors(self): + return [tensor for tensor in self.compute_tensors.values()] + + def get_num_elem(self): + return self.num_elem + + def get_swap_tensor(self, offset): + return self.swap_tensors.get(offset, None) + + def get_compute_tensor(self, offset): + return self.compute_tensors.get(offset, None) + + def get_swap_path(self, offset): + return self.swap_paths(offset, None) + + +class SwapBufferPool(object): + + def __init__(self, buffers): + assert all([get_accelerator().is_pinned(buf) for buf in buffers]) + self.buffers = [SwapBuffer(buf) for buf in buffers] + self.current_index = 0 + + def reset(self): + self.current_index = 0 + for buffer in self.buffers: + buffer.reset() + + def allocate_tensor(self, numel, swap_path, aligned_numel): + if self.has_space(aligned_numel): + swap_tensor, compute_tensor = self._get_current_buffer().allocate_tensor(swap_path, numel, aligned_numel) + return swap_tensor, compute_tensor + + return None, None + + def insert_tensor(self, tensor, swap_path, aligned_numel): + if self.has_space(aligned_numel): + swap_tensor, compute_tensor = self._get_current_buffer().insert_tensor(tensor, swap_path, aligned_numel) + return swap_tensor, compute_tensor + + return None, None + + def get_swap_tensors(self): + swap_tensors = [] + for buffer in self._get_used_buffers(): + swap_tensors += buffer.get_swap_tensors() + + return swap_tensors + + def get_swap_paths(self): + swap_paths = [] + for buffer in self._get_used_buffers(): + swap_paths += buffer.get_swap_paths() + + return swap_paths + + def get_compute_tensors(self): + compute_tensors = [] + for buffer in self._get_used_buffers(): + compute_tensors += buffer.get_compute_tensors() + + return compute_tensors + + def has_space(self, numel): + if self._get_current_buffer().has_space(numel): + return True + + if self.current_index == len(self.buffers) - 1: + return False + + self.current_index += 1 + return self._get_current_buffer().has_space(numel) + + def swap_out(self, aio_handle, async_op=False): + swap_tensors = self.get_swap_tensors() + swap_paths = self.get_swap_paths() + assert all([p is not None for p in swap_paths]) + + swap_out_tensors(aio_handle, swap_tensors, swap_paths) + + if not async_op: + assert len(swap_tensors) == aio_handle.wait() + + def swap_in(self, aio_handle, async_op=False): + swap_tensors = self.get_swap_tensors() + swap_paths = self.get_swap_paths() + assert all([p is not None for p in swap_paths]) + + swap_in_tensors(aio_handle, swap_tensors, swap_paths) + + if not async_op: + assert len(swap_tensors) == aio_handle.wait() + + def _get_current_buffer(self): + return self.buffers[self.current_index] + + def _get_used_buffers(self): + return self.buffers[:self.current_index + 1] + + +class SwapBufferManager(object): + + def __init__(self, num_elems, count, dtype): + self.num_elems = num_elems + self.count = count + self.dtype = dtype + self.all_buffers = [ + get_accelerator().pin_memory(torch.zeros(num_elems, device='cpu', dtype=dtype), align_bytes=0) + for _ in range(count) + ] + self.free_buffer_index = [i for i in range(count)] + self.used_buffer_index = {} + self.gigabytes = (self.all_buffers[0].element_size() * num_elems * count) / (1024**3) + + if dist.get_rank() == 0: + exclude_list = ['all_buffers'] + print_object(obj=self, name='SwapBufferManager', exclude_list=exclude_list) + + def allocate(self, num_elems, count, dtype): + assert dtype == self.dtype + assert num_elems <= self.num_elems + if count > len(self.free_buffer_index): + return None + + used_indices = self.free_buffer_index[-count:] + self.free_buffer_index = self.free_buffer_index[:-count] + + buffers = [] + for i in used_indices: + tmp_buffer = self.all_buffers[i].narrow(0, 0, num_elems) + buffers.append(tmp_buffer) + self.used_buffer_index[id(tmp_buffer)] = i + return buffers + + def allocate_all(self, num_elems, dtype): + return self.allocate(num_elems=num_elems, count=len(self.free_buffer_index), dtype=dtype) + + def free(self, buffers): + buffer_ids = [] + for buf in buffers: + buffer_ids.append(id(buf)) + + assert all([b_id in self.used_buffer_index for b_id in buffer_ids]) + + for b_id in buffer_ids: + self.free_buffer_index.append(self.used_buffer_index[b_id]) + del (self.used_buffer_index[b_id]) + + +def get_sized_buffer(buffer, num_elems): + assert num_elems <= buffer.numel(), \ + f'num_elems {num_elems} > buffer {buffer.numel()}' + return buffer.narrow(0, 0, num_elems) if num_elems < buffer.numel() else buffer + + +def get_sized_buffers(buffer_list, num_elems_list): + swap_buffers = [ + get_sized_buffer(buffer, num_elems) \ + for buffer, num_elems in zip(buffer_list, num_elems_list) + ] + return swap_buffers diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3882393453512ac915a9786bbfcad349c249fb4a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .config import AUTOTP_MODE, get_tensor_parallel_config +from .tp_manager import TpTrainingManager diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e8235a66fd97e18bfef37dde0e26d1a8f65eea9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ecbd9932baa7120398c4114eb11c69c1178b88bd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1317aeeb609f8359f3a906d85cb9700ca9846760 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/config.py new file mode 100644 index 0000000000000000000000000000000000000000..957984e9f8b339c195e94c685a9b1aa10e26cfbf --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/config.py @@ -0,0 +1,84 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import Enum +from deepspeed.runtime.config_utils import DeepSpeedConfigModel +import torch +from pydantic import Field +from typing import Optional + + +class AUTOTP_MODE(Enum): + TRAINING = "TRAINING" + INFERENCE = "INFERENCE" + + +class TPConfig(DeepSpeedConfigModel): + """ Configure tensor parallelism settings """ + + tp_size: int = 1 + """ Number of devices to split the model across using tensor parallelism. """ + + tp_grain_size: int = 1 + "The variable required by the autoTP parser has not been activated in training yet" + "as it depends on the gather logic that supports uneven partitioning. " + "Desired MLP/lm_head tp size granularity. DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size." + + mpu: object = None + """ + A model parallelism unit object that implements + ``get_{model,data}_parallel_{rank,group,world_size}()``. + """ + + tp_group: object = None + + +class TPTrainingConfig(DeepSpeedConfigModel): + + dtype: torch.dtype = torch.float16 + """ + Desired model data type, will convert model to this type. + """ + + autotp_size: int = 0 + """ + In automatic tensor-parallelism training, 'tensor_parallel_size' + When set to 0, indicates that it is disabled. + """ + tp_overlap_comm: bool = False + """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ + + tensor_parallel: TPConfig = Field({}, alias="tp") + """ + Configuration for tensor parallelism used to split the model across several + GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`. + """ + + injection_policy_tuple: Optional[tuple] = None + #The following parameters are required by autoTP parser. + ######################################## + keep_module_on_host: bool = False + """ + When loading checkpoints to model parameters, they are moved to the device. In very large models + this might fill the device and cause OOM. Setting this flag to true, will keep checkpoints on + host and not move them directly to the device (giving an option to quantize checkpoint data before + moving it to the device for example). + """ + + replace_with_kernel_inject: bool = Field(False, alias="kernel_inject") + """ + Set to true to inject inference kernels for models such as, Bert, GPT2, + GPT-Neo and GPT-J. Otherwise, the injection_dict provides the names of two + linear layers as a tuple: + `(attention_output projection, transformer output projection)` + """ + ######################################## + + +def get_tensor_parallel_config(ds_config): + + if 'tensor_parallel' in ds_config: + return TPTrainingConfig(**ds_config['tensor_parallel']) + return TPTrainingConfig() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..cf0b5a75c92aa3d426747917735835d1d11afe5f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from .config import TPTrainingConfig, TPConfig +from deepspeed.utils import groups +import deepspeed.comm as dist + + +class TpTrainingManager(): + + def __init__(self, model, tp_size, dtype): + self.module = model + self.config = self._initialize_config(dtype) + + from deepspeed.module_inject.auto_tp import AutoTP + from deepspeed import get_accelerator + + # Parse model configuration + parser_dict = AutoTP.tp_parser(model) + print("AutoTP: ", parser_dict) + + # Initialize TP configuration and model + self._initialize_tp_config(tp_size) + self._get_model_config_generate() + + # Synchronize random number generator state across devices + _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) + dist.broadcast(_rng_state, groups.get_tensor_model_parallel_src_rank(), self.tp_config.tp_group) + get_accelerator().set_rng_state(_rng_state.cpu()) + + # Apply injection policies + self._apply_policies(parser_dict) + + def _initialize_config(self, dtype): + """Initialize and return the DeepSpeed TP training configuration.""" + config = TPTrainingConfig() + config.dtype = dtype + return config + + def _apply_policies(self, parser_dict): + """Apply injection policies to the parsed modules.""" + for client_module, injection_policy in parser_dict: + self.config.injection_policy_tuple = injection_policy + self._apply_injection_policy(self.config, client_module) + + def _apply_injection_policy(self, config, client_module=None): + from deepspeed.module_inject import replace_transformer_layer + """Apply the given injection policy to a client module.""" + if isinstance(self.module, torch.nn.Module): + replace_transformer_layer(client_module, self.module, None, self.config, self.model_config) + + def _initialize_tp_config(self, tp_size): + """Perform TP configuration initialization.""" + self.tp_config = TPConfig() + self.tp_config.tp_size = tp_size + + groups._init_tp_mesh_device(tp_size) + self.tp_config.tp_group = groups.get_tensor_model_parallel_group() + self.config.tensor_parallel = self.tp_config + + def _get_model_config_generate(self): + """Generate and apply HF model configuration.""" + self.model_config = getattr(self.module, 'config', None) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/torch_autocast.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/torch_autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..faaa457e33bbd4d33a1816ec3e2ecbad1499f05b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/torch_autocast.py @@ -0,0 +1,100 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Iterable, Set, List, Union +import importlib + +import torch +from deepspeed.utils import logger + +LOWER_PRECISION_SAFE_MODULES = [ + torch.nn.Linear, + torch.nn.Conv1d, + torch.nn.Conv2d, + torch.nn.Conv3d, +] + +TORCH_AUTOCAST_INITIALIZED = False +_WARNED_NESTED_AUTOCAST = False + + +def _validate_auto_cast_settings(engine): + + assert not engine.fp16_enabled(), "Cannot enable both torch autocast and fp16" + assert not engine.bfloat16_enabled(), "Cannot enable both torch autocast and bfloat16" + assert not engine.zero_quantized_weights(), "Cannot enable both torch autocast and zero quantized weights" + + assert all(p.dtype == torch.float32 + for p in engine.parameters()), "All parameters must be float32 for torch autocast" + assert engine.communication_data_type == torch.float32, "Communication data type must be float32 for torch autocast" + + +def init_autocast_params(engine, dtype: torch.dtype, + torch_autocast_lower_precision_safe_modules: Union[None, List[str]]) -> None: + + _validate_auto_cast_settings(engine) + model = engine.module + + if torch_autocast_lower_precision_safe_modules is None: + lower_precision_safe_module_classes = LOWER_PRECISION_SAFE_MODULES + else: + lower_precision_safe_module_classes = [] + for module_name in torch_autocast_lower_precision_safe_modules: + try: + package_name, class_name = module_name.rsplit('.', 1) + module = importlib.import_module(package_name) + class_ = getattr(module, class_name) + lower_precision_safe_module_classes.append(class_) + except Exception as e: + raise ValueError(f"Failed to import lower precision safe module {module_name}: {e}") + + for module in model.modules(): + if module.__class__ in lower_precision_safe_module_classes: + for p in module.parameters(recurse=False): + p.autocast_dtype = dtype + + global TORCH_AUTOCAST_INITIALIZED + TORCH_AUTOCAST_INITIALIZED = True + + +def is_autocast_initialized() -> bool: + return TORCH_AUTOCAST_INITIALIZED + + +def get_default_autocast_lower_precision_modules() -> List[str]: + return [f"{cls.__module__}.{cls.__name__}" for cls in LOWER_PRECISION_SAFE_MODULES] + + +def get_autocast_dtype(param: torch.nn.Parameter) -> torch.dtype: + return param.autocast_dtype if hasattr(param, "autocast_dtype") else param.dtype + + +def has_autocast_dtype(param: torch.nn.Parameter) -> bool: + return hasattr(param, "autocast_dtype") + + +def get_all_autocast_dtypes(params: Iterable) -> Set[torch.dtype]: + return {get_autocast_dtype(p) for p in params} + + +def sort_dtypes(dtypes: List[torch.dtype]) -> List[torch.dtype]: + return sorted(dtypes, key=str) + + +def validate_nested_autocast(engine): + global _WARNED_NESTED_AUTOCAST + + if torch.is_autocast_enabled(): + if engine.torch_autocast_enabled(): + if not _WARNED_NESTED_AUTOCAST: + logger.warning( + "DeepSpeed detected torch.autocast context outside the engine. " + "This is unnecessary when torch.autocast is already enabled through the DeepSpeed config.") + _WARNED_NESTED_AUTOCAST = True + else: + raise AssertionError( + "torch.autocast is enabled outside DeepSpeed, but not in the DeepSpeed config. " + "Please enable torch.autocast through the DeepSpeed config to ensure the correct communication dtype is used." + ) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fa65b9041faf95447385637325dbcc03c1cf58bd --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/utils.py @@ -0,0 +1,1185 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/Megatron + +Helper functions and classes from multiple sources. +""" + +from collections.abc import Iterable +import os +import psutil +import gc +from math import sqrt + +from numpy import prod + +import torch +from torch.nn import functional as F +try: + from torch._six import inf +except ModuleNotFoundError: + from torch import inf +from typing import Union, List, Dict +from deepspeed import comm as dist +from deepspeed.moe.utils import is_moe_param +from deepspeed.utils import groups, logger +from deepspeed.utils.bwc import (bwc_tensor_model_parallel_rank, bwc_pipeline_parallel_world_size, + bwc_pipeline_parallel_group) +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.policy import transpose + +torch_memory_reserved = get_accelerator().memory_reserved +torch_max_memory_reserved = get_accelerator().max_memory_reserved + + +class DummyOptim(): + """ + Dummy optimizer presents model parameters as a param group, this is + primarily used to allow ZeRO-3 without an optimizer + """ + + def __init__(self, params): + self.param_groups = [] + self.param_groups.append({'params': params}) + + +graph_cache = {} + + +def graph_process(replay_first_step, func, *args, **kwargs): + # `func` should only contain operations on the GPU + # Please ensure that the memory address of the data required by 'func' remains constant + if func.__name__ not in graph_cache: + cuda_stream = get_accelerator().Stream() + cuda_stream.wait_stream(get_accelerator().current_stream()) + with get_accelerator().stream(cuda_stream): + func(*args, **kwargs) + get_accelerator().current_stream().wait_stream(cuda_stream) + graph_cache[func.__name__] = get_accelerator().create_graph() + with get_accelerator().capture_to_graph(graph_cache[func.__name__]): + func(*args, **kwargs) + if replay_first_step: + get_accelerator().replay_graph(graph_cache[func.__name__]) + else: + get_accelerator().replay_graph(graph_cache[func.__name__]) + + +def noop_decorator(func): + return func + + +class noop_context(object): + + def __init__(self): + pass + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + +def ensure_directory_exists(filename): + """Create the directory path to ``filename`` if it does not already exist. + + Args: + filename (str): A file path. + """ + dirname = os.path.dirname(filename) + os.makedirs(dirname, exist_ok=True) + + +def set_random_seed(seed): + """Set the random seed for common PRNGs used during training: random, numpy, and torch. + + Args: + seed (int): the seed to use + """ + import numpy + import random + random.seed(seed) + numpy.random.seed(seed) + torch.manual_seed(seed) + + +def is_model_parallel_parameter(p) -> bool: + if hasattr(p, 'model_parallel') and p.model_parallel: + return True + + if hasattr(p, 'tensor_model_parallel') and p.tensor_model_parallel: + return True + + return False + + +def copy_to_device(item, device, criterion_func): + """ + Return a copy of tensor on specified device. + Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. + Parameters: + item: tensor to copy or (possibly nested) container of tensors to copy. + device: target device + criterion_func: Function to restrict copy operation to items meet criterion + + Returns: + None + """ + if criterion_func(item): + return item.to(device) + elif isinstance(item, list): + return [copy_to_device(v, device, criterion_func) for v in item] + elif isinstance(item, tuple): + return tuple([copy_to_device(v, device, criterion_func) for v in item]) + elif isinstance(item, dict): + return {k: copy_to_device(v, device, criterion_func) for k, v in item.items()} + else: + return item + + +def move_to_device(item, device, criterion_func=None): + """ + Move tensor on to specified device by changing the storage. + Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. + Parameters: + item: tensor to move or (possibly nested) container of tensors to move. + device: target device + criterion_func: Function to restrict move operation to items meet criterion, defaults to `None` which is an equivalent to always move + + Returns: + None + """ + if (criterion_func is not None and criterion_func(item)): + device_copy = item.to(device) + item.data = device_copy.data + return item + elif isinstance(item, list): + return [move_to_device(v, device, criterion_func) for v in item] + elif isinstance(item, tuple): + return tuple([move_to_device(v, device, criterion_func) for v in item]) + elif isinstance(item, dict): + return {k: move_to_device(v, device, criterion_func) for k, v in item.items()} + else: + return item.to(device) + + +def get_norm_with_moe_layers_fast(all_groups_norm, group): + # This implementation standardizes the grad_norm across ranks. A more precise implementation can be found in 'get_norm_with_moe_layers'. + # Need to allreduce (avg) the norms across different ranks because moe params will not be synced during allreduce + scaled_norm = all_groups_norm * 1.0 / float(dist.get_world_size(group=group)) + scaled_norm_tensor = torch.tensor(scaled_norm, device=get_accelerator().current_device_name(), dtype=torch.float) + dist.all_reduce(scaled_norm_tensor, group=group) + all_groups_norm = scaled_norm_tensor.item() + #print(f"old = {all_groups_norm_old} and new = {all_groups_norm} at rank: {deepspeed.comm.get_rank()}") + return all_groups_norm + + +class CheckOverflow(object): + '''Checks for overflow in gradient across parallel process''' + + def __init__(self, param_groups=None, mpu=None, zero_reduce_scatter=False, deepspeed=None): + self.mpu = mpu + self.params = [] if param_groups else None + self.zero_reduce_scatter = zero_reduce_scatter + self.deepspeed = deepspeed + self.has_moe_params = False + if param_groups: + for group in param_groups: + for param in group: + self.params.append(param) + if is_moe_param(param): + self.has_moe_params = True + + def check_using_norm(self, norm_group, reduce_overflow=True): + # TODO: I don't think reduce_overflow is needed if mpu is None + overflow = -1 in norm_group + overflow_gpu = get_accelerator().FloatTensor([overflow]) + if self.has_moe_params: + # In this case, we need to do an all_reduce across + # the expert_parallel_group, so that if there was + # an overflow due to expert weights, we detect it + + # Only need to check groups.get_largest_expert_parallel_group() + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) + if self.mpu is not None: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) + elif reduce_overflow: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX) + dist.barrier() + overflow = overflow_gpu[0].item() + return bool(overflow) + + def check(self, param_groups=None): + params = [] + has_moe_params = False + if param_groups is None: + params = self.params + has_moe_params = self.has_moe_params + else: + assert param_groups is not None, \ + "self.params and param_groups both cannot be none" + + for group in param_groups: + for param in group: + params.append(param) + if is_moe_param(param): + has_moe_params = True + + return self.has_overflow(params, has_moe_params=has_moe_params) + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params): + for i, p in enumerate(params): + if p.grad is not None and self._has_inf_or_nan(p.grad.data, i): + return True + return False + + def has_overflow(self, params, has_moe_params=None): + if has_moe_params is None: + has_moe_params = self.has_moe_params + overflow = self.has_overflow_serial(params) + # Since each model parallel GPU carries only part of the model, + # make sure overflow flag is synced across all the model parallel GPUs + overflow_gpu = get_accelerator().ByteTensor([overflow]) + # deepspeed.comm.all_reduce(overflow_gpu, + # op=deepspeed.comm.ReduceOp.MAX, + # group=mpu.get_model_parallel_group()) + if has_moe_params: + # All reduce this across expert_parallel_group, so that if an expert + # overflows, we detect it here + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) + if self.zero_reduce_scatter: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) + elif self.mpu is not None: + if self.deepspeed is not None: + using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + if (using_pipeline and self.deepspeed.pipeline_enable_backward_allreduce + is False) or (not using_pipeline and self.deepspeed.enable_backward_allreduce is False): + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_data_parallel_group()) + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) + elif self.deepspeed is not None and self.deepspeed.enable_backward_allreduce is False: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) + + overflow = overflow_gpu[0].item() + return bool(overflow) + + # `x` is a torch.Tensor + @staticmethod + def _has_inf_or_nan(x, i): + try: + # if x is half, the .float() incurs an additional deep copy, but it's necessary if + # Pytorch's .sum() creates a one-element tensor of the same type as x + # (which is true for some recent version of pytorch). + cpu_sum = float(x.float().sum()) + # More efficient version that can be used if .sum() returns a Python scalar + # cpu_sum = float(x.sum()) + except RuntimeError as instance: + # We want to check if inst is actually an overflow exception. + # RuntimeError could come from a different error. + # If so, we still want the exception to propagate. + if "value cannot be converted" not in instance.args[0]: + raise + return True + else: + if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: + return True + return False + + +def _handle_overflow(cpu_sum, x, i): + import math + rank = dist.get_rank() + if rank == 0: + t_i = -1 + for v_i, v in enumerate(x.data.contiguous().view(-1)): + if not math.isfinite(float(v)): + t_i = v_i + break + logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") + + +def get_global_norm(norm_list): + """ Compute total from a list of norms + """ + total_norm = 0.0 + for norm in norm_list: + total_norm += norm**2.0 + # logger.info(f'norm_list = {norm_list} global = {sqrt(total_norm)}') + return sqrt(total_norm) + + +def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None): + """Clips gradient norm of an iterable of parameters. + + This has been adapted from Nvidia megatron. We add norm averaging + to consider MoE params when calculating norm as they will result + in different norms across different ranks. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float or int): max norm of the gradients + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + norm_type = float(norm_type) + all_norms = [] + if norm_type == inf: + for p in parameters: + all_norms.append(p.grad.data.abs().max().float()) + total_norm = torch.stack(all_norms).max() + total_norm = total_norm.to(get_accelerator().current_device_name()) + # Take max across all GPUs. + if mpu is not None: + dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + else: + total_norm = 0 + for p in parameters: + if mpu is not None: + if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p): + param_norm = p.grad.data.detach().float().norm(norm_type) + all_norms.append(param_norm) + else: + param_norm = p.grad.data.detach().float().norm(norm_type) + all_norms.append(param_norm) + if len(all_norms) > 0: + total_norm = torch.stack(all_norms).square().sum().float() + else: + total_norm = get_accelerator().FloatTensor([0.0]) + total_norm = total_norm.to(get_accelerator().current_device_name()) + # Sum across all model parallel GPUs. + if mpu is not None: + dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_norm = total_norm.pow(1. / norm_type) + + # Need to average total_norm across different GPUs due to the presence of moe params + pg = groups._get_data_parallel_group() + scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg)) + scaled_norm_tensor = scaled_norm + + dist.all_reduce(scaled_norm_tensor, group=pg) + total_norm = scaled_norm_tensor + total_norm = total_norm.to(parameters[0].device) + + max_norm = torch.tensor([float(max_norm)], device=total_norm.device) + clip_coef = max_norm / (total_norm + 1e-6) + tmp_tensor = torch.tensor([1.0], device=clip_coef.device) + clip_coef = torch.min(tmp_tensor, clip_coef) + for p in parameters: + p.grad.data.mul_(clip_coef) + return total_norm + + +def get_flattened_grad_norm(parameters, norm_type=2, mpu=None, grad_norm_mask=None): + """Get grad norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. Taken from Nvidia Megatron. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + grad_norm_mask (List[Tensor]): A list of Tensor, where + each Tensor is a 2D Tensor containing ranges of [start_index, end_index]. + Returns: + Total norm of the parameters (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + + norm_type = float(norm_type) + if norm_type == inf: + total_norm = max(p.grad.data.abs().max() for p in parameters) + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + # Take max across all GPUs. + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item() + else: + total_norm = 0. + for idx, p in enumerate(parameters): + # Use grad_norm_mask to avoid redundant computation of flattened gradient norm + if grad_norm_mask is not None and len(grad_norm_mask[idx]) > 0: + + # A loop-free implementation to create a mask tensor based on a range list + # which is logically equivalent to the following implementation. + # # mask_tensor_ = torch.zeros_like(p, device=p.device, dtype=bool) + # # for mask_idx in grad_norm_mask[idx]: + # # mask_tensor_[mask_idx[0]:mask_idx[1]] = True + cum_sum_pairs = torch.tensor([1, -1], device=get_accelerator().current_device_name(), + dtype=p.dtype).repeat(grad_norm_mask[idx].shape[0], 1) + mask_tensor = torch.zeros(p.shape[0] + 1, + device=get_accelerator().current_device_name(), + dtype=p.dtype) + mask_tensor = mask_tensor.scatter_(0, grad_norm_mask[idx].view(-1), + cum_sum_pairs.view(-1)).cumsum(0).bool()[:-1] + + param_norm = torch.masked_fill(p.grad.data, mask_tensor, 0).float().norm(norm_type) + + else: + param_norm = p.grad.data.float().norm(norm_type) + total_norm += param_norm.item()**norm_type + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item()**(1. / norm_type) + + if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: + total_norm = -1 + + return total_norm + + +def get_grad_zeros(parameters, mpu=None): + """Compute the number of grads with zero values. + + This is adapted from get_grad_norm + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + + Returns: + Total number of params with zero values (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + + total_zeros = 0. + tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) + for p in parameters: + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + + # Filter to avoid over-counting replicated tensors from tensor + # model parallelism + if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): + continue + + count_zeros = p.grad.numel() - torch.count_nonzero(p.grad) + total_zeros += count_zeros.item() + + # Sum across all model parallel GPUs. + total_zeros_cuda = get_accelerator().FloatTensor([float(total_zeros)]) + if mpu is not None: + dist.all_reduce(total_zeros_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_zeros = total_zeros_cuda[0].item() + + return total_zeros + + +def get_weight_norm(parameters, norm_type=2, mpu=None): + """Get norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. Taken from Nvidia Megatron. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + -1 if the norm value is NaN or Inf. + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + + norm_type = float(norm_type) + if norm_type == inf: + total_norm = max(p.data.abs().max() for p in parameters) + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + # Take max across all GPUs. + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item() + else: + total_norm = 0. + tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) + for p in parameters: + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + + # Filter to avoid over-counting replicated tensors from tensor + # model parallelism + if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): + continue + + param_norm = p.data.float().norm(norm_type) + total_norm += param_norm**norm_type + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item()**(1. / norm_type) + + if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: + total_norm = -1 + + return total_norm + + +def prefix_sum_inc(weights): + """ Compute an inclusive prefix sum. + + Example: + >>> prefix_sum_inc([3,4,5]) + [3, 7, 12] + """ + weights_ = [w for w in weights] + for x in range(1, len(weights_)): + weights_[x] += weights_[x - 1] + return weights_ + + +def partition_uniform(num_items, num_parts): + import numpy + parts = [0] * (num_parts + 1) + # First check for the trivial edge case + if num_items <= num_parts: + for p in range(num_parts + 1): + parts[p] = min(p, num_items) + return parts + + chunksize = num_items // num_parts + residual = num_items - (chunksize * num_parts) + + parts = numpy.arange(0, (num_parts + 1) * chunksize, chunksize) + + for i in range(residual): + parts[i + 1:] += 1 + parts = parts.tolist() + + return parts + + +def partition_balanced(weights, num_parts): + """ + use dynamic programming solve `The Linear Partition Problem`. + see https://www8.cs.umu.se/kurser/TDBAfl/VT06/algorithms/BOOK/BOOK2/NODE45.HTM + """ + import numpy as np + n = len(weights) + m = num_parts + + if n <= m: + return partition_uniform(n, m) + + dp_max = np.full((n + 1, m + 1), np.inf) + dp_min = np.full((n + 1, m + 1), np.inf) + dp_cost = np.full((n + 1, m + 1), np.inf) + position = np.zeros((n + 1, m + 1), dtype=int) + prefix_sum = np.zeros((n + 1)) + prefix_sum[1:] = np.cumsum(weights) + + dp_max[0, 0] = 0 + dp_cost[0, 0] = 0 + for i in range(1, n + 1): + for j in range(1, min(i, m) + 1): + for k in range(i): + max_sum = max(dp_max[k, j - 1], prefix_sum[i] - prefix_sum[k]) + min_sum = min(dp_min[k, j - 1], prefix_sum[i] - prefix_sum[k]) + cost = max_sum - min_sum + if dp_cost[i, j] >= cost: + dp_cost[i, j] = cost + dp_max[i, j] = max_sum + dp_min[i, j] = min_sum + position[i, j] = k + + parts = [n] + for i in reversed(range(1, m + 1)): + parts.append(position[parts[-1], i]) + parts.reverse() + + return parts + + +class PartitionedTensor: + + def __init__(self, tensor, group, partition_meta=None): + super().__init__() + + self.group = group + self.num_parts = dist.get_world_size(group=self.group) + self.rank = dist.get_rank(group=self.group) + self.orig_size = list(tensor.size()) + self.orig_device = tensor.device + self.local_data, self.partition = self._partition_tensor(tensor) + self.even_split = tensor.numel() % self.num_parts == 0 + + @classmethod + def from_meta(cls, meta, local_part, group, device=get_accelerator().device_name()): + assert meta.dtype == torch.long + dummy = torch.ones(dist.get_world_size(group=group)) + part_obj = cls(tensor=dummy, group=group) + + meta = meta.tolist() + + # [N, list0, ..., listN-1] + part_obj.orig_size = meta[1:(1 + meta[0])] + meta = meta[1 + meta[0]:] + + part_obj.orig_device = device + part_obj.local_data = local_part.detach() + + part_obj.group = group + + # Partition is encoded like the rowptr of a CSR matrix: + # [num_parts, rank, 0, part_1, ..., part_num_parts] + # TODO: support shuffle between different partition granularities + assert part_obj.num_parts == meta[0] + assert part_obj.rank == meta[1] + part_obj.partition = meta[2:] # length num_parts+1 + + return part_obj + + def _partition_tensor(self, tensor): + partition = partition_uniform(num_items=tensor.numel(), num_parts=self.num_parts) + start = partition[self.rank] + length = partition[self.rank + 1] - start + tensor_part = tensor.detach().contiguous().view(-1).narrow(0, start=start, length=length).clone() + + return tensor_part, partition + + def full(self, device=None): + if device is None: + device = self.orig_device + + # Allocate the full tensor as a flat buffer. + full_numel = prod(self.full_size()) + flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device) + if self.even_split: + # Collect the full tensor + dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group) + else: + for part_id in range(self.num_parts): + part_size = self.partition[part_id + 1] - self.partition[part_id] + buf = flat_tensor.narrow(0, start=self.partition[part_id], length=part_size) + if part_id == self.rank: + buf.copy_(self.local_data) + dist.broadcast(buf, part_id, self.group) + return flat_tensor.view(self.full_size()).clone().detach() + + def to_meta(self): + """Returns a torch.LongTensor that encodes partitioning information. + + Can be used along with ``data()`` to serialize a ``PartitionedTensor`` for + communication. + + Returns: + torch.LongTensor: a tensor encoding the meta-information for the partitioning + """ + meta = [] + meta.append(len(self.orig_size)) + meta += list(self.orig_size) + meta.append(self.num_parts) + meta.append(self.rank) + meta += self.partition + return torch.LongTensor(data=meta).to(self.orig_device) + + def data(self): + return self.local_data + + def local_size(self): + return self.local_data.size() + + def full_size(self): + return self.orig_size + + +mem_alloced = 0 +mem_cached = 0 + + +def memory_status(msg, print_rank=-1, reset_max=False): + global mem_alloced, mem_cached + + rank = dist.get_rank() + if print_rank != -1 and rank != print_rank: + return + + get_accelerator().synchronize() + + if reset_max: + get_accelerator().reset_max_memory_cached() + get_accelerator().reset_max_memory_allocated() + + new_alloced = get_accelerator().memory_allocated() + new_cached = get_accelerator().memory_cached() + + delta_alloced = new_alloced - mem_alloced + delta_cached = new_cached - mem_cached + + mem_cached = new_cached + mem_alloced = new_alloced + + max_alloced = get_accelerator().max_memory_allocated() + max_cached = get_accelerator().max_memory_cached() + + # convert to GB for printing + new_alloced /= 1024**3 + new_cached /= 1024**3 + delta_alloced /= 1024**3 + delta_cached /= 1024**3 + max_alloced /= 1024**3 + max_cached /= 1024**3 + + print( + f'RANK={rank} MEMSTATS', msg, f'device={get_accelerator().current_device_name()} ' + f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' + f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') + + +def get_ma_status(): + if dist.is_initialized() and not dist.get_rank() == 0: + return 0 + return get_accelerator().memory_allocated() + + +def empty_cache(): + get_accelerator().empty_cache() + get_accelerator().reset_peak_memory_stats() + + +def see_memory_usage(message, force=False): + if not force: + return + if dist.is_initialized() and not dist.get_rank() == 0: + return + + # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports + gc.collect() + + # Print message except when distributed but not rank 0 + logger.info(message) + logger.info(f"MA {round(get_accelerator().memory_allocated() / (1024 * 1024 * 1024),2 )} GB \ + Max_MA {round(get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \ + CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \ + Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ") + + vm_stats = psutil.virtual_memory() + used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2) + logger.info(f'CPU Virtual Memory: used = {used_GB} GB, percent = {vm_stats.percent}%') + + # get the peak memory to report correct data, so reset the counter for the next call + get_accelerator().reset_peak_memory_stats() + + +def call_to_str(base, *args, **kwargs): + """Construct a string representation of a call. + + Args: + base (str): name of the call + args (tuple, optional): args to ``base`` + kwargs (dict, optional): kwargs supplied to ``base`` + + Returns: + str: A string representation of base(*args, **kwargs) + """ + name = f'{base}(' + if args: + name += ', '.join(repr(arg) for arg in args) + if kwargs: + name += ', ' + if kwargs: + name += ', '.join(f'{key}={repr(arg)}' for key, arg in kwargs.items()) + name += ')' + return name + + +def get_only_unique_item(items): + item_set = set(items) + if len(item_set) != 1: + raise RuntimeError(f"expected there to be only one unique element in {items}") + unique_item, = item_set + + return unique_item + + +def mask_nan_or_inf_with_val_inplace(input, device=None, val=-1.): + norm_is_inf = input.isinf() + norm_is_nan = input.isnan() + inf_or_nan = norm_is_nan.logical_or(norm_is_inf) + err = torch.tensor(-1.0, device=device, dtype=torch.float) + input.masked_fill_(inf_or_nan, err) + + +def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, use_graph=False, moe_ep_group=None): + """Get norm of an iterable of tensors. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Taken from Nvidia Megatron. + + Arguments: + input_tensors (Iterable[Tensor]): an iterable of Tensors will have norm computed + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the tensors (viewed as a single vector). + """ + assert isinstance(input_tensors, Iterable), f'expected Iterable type not {type(input_tensors)}' + assert all([torch.is_tensor(t) for t in input_tensors]), f'expected list of only tensors' + + norm_type = float(norm_type) + all_norms = [] + if norm_type == inf: + for t in input_tensors: + all_norms.append(t.data.abs().max().float()) + total_norm = torch.stack(all_norms).max() + device_total_norm = total_norm.to(get_accelerator().current_device_name()) + # Max across model parallel + if mpu is not None: + # For MoE grads, max over model parallel only if MoE-TP is enabled + if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + # If MoE grads and MoE-TP disabled, max over pipeline parallel + elif bwc_pipeline_parallel_world_size(mpu) > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=bwc_pipeline_parallel_group(mpu)) + + # MoE grads: max across expert parallel group + if moe_ep_group is not None: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=moe_ep_group) + total_norm = device_total_norm.to(input_tensors[0].device) + else: + + if 'norm_tensors_compute_buffer' not in graph_cache or len( + graph_cache['norm_tensors_compute_buffer']) != len(input_tensors): + graph_cache['norm_tensors_compute_buffer'] = [ + torch.empty([], dtype=torch.float, device=get_accelerator().current_device_name()) + for t in input_tensors + ] + compute_buffer = graph_cache['norm_tensors_compute_buffer'] + + def _norm_tensors(tensor_list, _compute_buffer, _norm_type): + for i, t in enumerate(tensor_list): + _compute_buffer[i].data.copy_(t.data.float().norm(_norm_type)**_norm_type) + if i != 0: + _compute_buffer[0].data.add_(_compute_buffer[i].data) + + if use_graph: + graph_process(False, _norm_tensors, input_tensors, compute_buffer, norm_type) + else: + _norm_tensors(input_tensors, compute_buffer, norm_type) + + device_total_norm = compute_buffer[0].float().detach() + + # Sum across model parallel + if mpu is not None: + # For MoE grads, sum over model parallel only if MoE-TP is enabled + if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + # If MoE grads and MoE-TP disabled, sum over pipeline parallel + elif bwc_pipeline_parallel_world_size(mpu) > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=bwc_pipeline_parallel_group(mpu)) + + # MoE grads: sum across expert parallel group + if moe_ep_group is not None: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=moe_ep_group) + total_norm = device_total_norm.to(input_tensors[0].device).pow(1. / norm_type) + + mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device) + + return total_norm + + +def clip_tensors_by_global_norm(input_tensors, max_norm=1.0, global_norm=None, mpu=None, eps=1e-6, use_graph=False): + """Clip list of tensors by global norm. + Args: + input_tensors: List of tensors to be clipped + global_norm (float, optional): Precomputed norm. Defaults to None. + mpu (optional): model parallelism unit. Defaults to None. + eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6 + Returns: + float: the global norm + """ + if global_norm is None: + global_norm = get_global_norm_of_tensors(input_tensors, mpu=mpu, use_graph=use_graph) + clip_coef = max_norm / (global_norm + eps) + if clip_coef < 1: + if use_graph: + + def clip_tensors(_tensor_list, _clip_coef_tensor): + for t in _tensor_list: + t.detach().mul_(_clip_coef_tensor) + + if 'clip_coef_tensor' not in graph_cache: + # Alloc memory + graph_cache['clip_coef_tensor'] = torch.tensor(clip_coef, + dtype=torch.float32).to(get_accelerator().device_name()) + clip_coef_tensor = graph_cache['clip_coef_tensor'] + clip_coef_tensor.copy_(torch.tensor(clip_coef, dtype=torch.float32)) + graph_process(False, clip_tensors, input_tensors, clip_coef_tensor) + + else: + for t in input_tensors: + t.detach().mul_(clip_coef) + return global_norm + + +def align_dense_tensors(tensor_list, alignment): + num_elements = sum(t.numel() for t in tensor_list) + remaining = num_elements % alignment + + if remaining: + elements_to_add = alignment - remaining + pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) + padded_tensor_list = tensor_list + [pad_tensor] + else: + padded_tensor_list = tensor_list + + return padded_tensor_list + + +def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, dp_process_group): + for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)): + partition_id = dist.get_rank(group=dp_process_group[group_id]) + dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) + if dp_world_size == 1: + # no groups share optimizer states + # pipeline parallel with bf16 will default call this even if dp size = 1. + continue + dist.all_gather_into_tensor(group_flat, partitioned_params[partition_id], dp_process_group[group_id]) + + +def all_gather_dp_groups(groups_flat, partitioned_param_groups, dp_process_group, start_alignment_factor, + allgather_bucket_size): + if dist.has_all_gather_into_tensor(): + return all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, dp_process_group) + + for group_id, partitioned_params in enumerate(partitioned_param_groups): + # Sequential AllGather Best of both worlds + partition_id = dist.get_rank(group=dp_process_group[group_id]) + dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) + + if dp_world_size == 1: + # no groups share optimizer states + # pipeline parallel with bf16 will default call this even if dp size = 1. + continue + num_shards = max(1, partitioned_params[partition_id].numel() * dp_world_size // allgather_bucket_size) + + shard_size = partitioned_params[partition_id].numel() // num_shards + + # Enforce nccl/rccl alignment of start location of each shard + shard_size = shard_size - (shard_size % start_alignment_factor) + + num_elements = shard_size + + assert shard_size * num_shards <= partitioned_params[partition_id].numel() + + for shard_id in range(num_shards): + + if shard_id == (num_shards - 1): + num_elements = partitioned_params[partition_id].numel() - shard_id * shard_size + + shard_list = [] + for dp_id in range(dp_world_size): + curr_shard = partitioned_params[dp_id].narrow(0, shard_id * shard_size, num_elements).detach() + shard_list.append(curr_shard) + + dist.all_gather(shard_list, shard_list[partition_id], dp_process_group[group_id]) + + +def get_tensor_bytes(item): + if torch.is_tensor(item): + return item.numel() * item.element_size() + elif isinstance(item, list): + return sum([get_tensor_bytes(v) for v in item]) + elif isinstance(item, tuple): + return sum([get_tensor_bytes(v) for v in item]) + elif isinstance(item, dict): + return sum([get_tensor_bytes(v) for v in item.values()]) + else: + return 0 + + +def _get_folder_size(folder): + size = 0 + for path, _, files in os.walk(folder): + size += sum([os.path.getsize(os.path.join(path, f)) for f in files]) + return size + + +def get_checkpoint_folder_size(save_dir, tag, local_rank=None): + if local_rank == 0: + folder = os.path.join(save_dir, tag) + size_tensor = torch.tensor(_get_folder_size(folder)).to(get_accelerator().device_name()) + else: + size_tensor = torch.tensor(0).to(get_accelerator().device_name()) + + dist.reduce(tensor=size_tensor, dst=0) + return int(size_tensor) + + +class TLinear(torch.nn.Linear): + + def __init__(self, orig_layer, name=""): + self.name = name + super().__init__(orig_layer.weight.shape[1], orig_layer.weight.shape[0], bias=(orig_layer.bias is not None)) + self.weight.data = transpose(orig_layer.weight.data) + self.bias = orig_layer.bias + self._fwd_func = self._fwd_bias_add if self.bias is not None else self._fwd + + def _fwd(self, input): + return F.linear(input, self.weight) + + def _fwd_bias_add(self, input): + return F.linear(input, self.weight, bias=self.bias) + + def forward(self, input): + return self._fwd_func(input) + + +def get_inactive_params(param_list): + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + return [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + + +def get_norm_with_moe_layers(non_expert_norm, mpu, expert_tensors, norm_type=2): + """ Compute the global norm with MoE experts + + Inputs: + non_expert_norm (float) : the calculated norm of the non-expert params + expert_tensors (Dict[ep_name, List[Tensor]): Dictionary of expert group name to list of grad tensors + norm_type (int): the norm to use + + Returns: + if norm is (-/+) inf, returns -1 + otherwise the global norm (float) + """ + + def to_tensor(v): + return get_accelerator().FloatTensor(float(v)).detach() + + group_norms = [non_expert_norm] + for exp_name, tensors in expert_tensors.items(): + group_norm = get_global_norm_of_tensors(input_tensors=tensors, + mpu=mpu, + norm_type=norm_type, + use_graph=False, + moe_ep_group=groups._get_expert_parallel_group(exp_name)) + group_norms.append(group_norm) + + # check if all norms are valid + group_norms = torch.stack([to_tensor(norm) for norm in group_norms]) + if group_norms.eq(-1).any(): + return -1 + + # combine norms + if norm_type == inf: + total_norm = group_norms.max().item() + else: + total_norm = group_norms.pow(norm_type).sum() + total_norm = total_norm.item()**(1. / norm_type) + if total_norm == float('inf') or total_norm == -float('inf'): + total_norm = -1 + + return total_norm + + +def _make_offload_state_key(key): + return f"{key}_offload_buffer" + + +def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_key(state, key): + offload_buf_key = _make_offload_state_key(key) + if offload_buf_key not in state: + state[offload_buf_key] = torch.empty_like(state[key], device=device) + if pin_memory: + state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) + state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) + state[key].data = state[offload_buf_key] + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_key(state, "exp_avg_sq") + + +def reload_adam_states(optimizer, device, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_back_key(state, key): + state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking) + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_back_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_back_key(state, "exp_avg_sq") + + +def compare_tensors_in_structures(inputs1: Union[List, Dict], inputs2: Union[List, Dict]) -> bool: + """ + Compare two lists or dictionaries for equality, including any tensors they may contain. + + Args: + inputs1: First input, either a list or a dictionary. + inputs2: Second input, either a list or a dictionary. + + Returns: + True if inputs1 and inputs2 are equal; False otherwise. + """ + if type(inputs1) != type(inputs2): # Ensure types match + return False + + if isinstance(inputs1, list) and isinstance(inputs2, list): + if len(inputs1) != len(inputs2): + return False + for val1, val2 in zip(inputs1, inputs2): + if isinstance(val1, torch.Tensor) and isinstance(val2, torch.Tensor): + val1 = val1.to(get_accelerator().current_device()) + val2 = val2.to(get_accelerator().current_device()) + if not torch.equal(val1, val2): + return False + elif val1 != val2: + return False + return True + + elif isinstance(inputs1, dict) and isinstance(inputs2, dict): + if inputs1.keys() != inputs2.keys(): + return False + for key in inputs1: + val1, val2 = inputs1[key], inputs2[key] + if isinstance(val1, torch.Tensor) and isinstance(val2, torch.Tensor): + val1 = val1.to(get_accelerator().current_device()) + val2 = val2.to(get_accelerator().current_device()) + if not torch.equal(val1, val2): + return False + elif val1 != val2: + return False + return True + + return False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/weight_quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/weight_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..95d737614e594ee200d13617654714d6ae9d26f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/weight_quantizer.py @@ -0,0 +1,153 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..module_inject.replace_policy import HFBertLayerPolicy, replace_policies +from deepspeed.accelerator import get_accelerator + + +class WeightQuantization(object): + + def __init__(self, mlp_extra_grouping=True, mp_size=1): + self.dense_scales = [] + self.qkv_scales = [] + self.mlp4hh_scales = [] + self.mlph4h_scales = [] + self.mlp_extra_grouping = mlp_extra_grouping + self.mp_size = mp_size + + def quantize_data(self, data, quantize_bits, groups, key=None): + data_groups = torch.split(data.float().view(-1), data.numel() // groups) + max_d = [max(g.max(), g.min().abs()) for g in data_groups] + data_scale = [float(1 << quantize_bits) / (2 * mx + 1e-5) for mx in max_d] + data_int = [(g * s) for g, s in zip(data_groups, data_scale)] + data_int = [ + di.round().clamp(-(1 << (quantize_bits - 1)), (((1 << (quantize_bits - 1)) - 1))) for di in data_int + ] + data_int = torch.cat(data_int).reshape(data.shape) + data_int = data_int.to(torch.int8) + data_scale = torch.cat([s.unsqueeze(0).unsqueeze(0) for s in data_scale]) + return data_int, data_scale + + def is_mlp(self, data, merge_count=1): + return ((self.mp_size *data.shape[0] * merge_count) / data.shape[1] == 4 or \ + (self.mp_size *data.shape[1] * merge_count) / data.shape[0] == 4) + + def is_qkv(self, data): + return ((self.mp_size * data.shape[0]) / data.shape[1] == 3 or \ + (self.mp_size * data.shape[1]) / data.shape[0] == 3) + + def Quantize(self, value_list, quantize_bits, groups, key, merge_dim=0): + if self.mlp_extra_grouping and self.is_mlp(value_list[0], merge_count=len(value_list)): + groups *= 2 + q_scale = [] + index = 0 + for data in value_list: + data_int, data_scale = self.quantize_data(data, quantize_bits, groups, key) + q_scale.append(data_scale) + value_list[index] = data_int + index += 1 + q_scale = (1 / + torch.cat(q_scale, dim=merge_dim).to(get_accelerator().current_device_name()).view(-1).unsqueeze(0)) + if "mlp.dense_4h_to_h.weight" in key: + self.mlp4hh_scales.append(q_scale) + elif "mlp.dense_h_to_4h.weight" in key: + self.mlph4h_scales.append(q_scale) + elif "attention.query_key_value.weight" in key: + self.qkv_scales.append(q_scale) + else: + self.dense_scales.append(q_scale) + return value_list + + def merge_layer_scales(self, layer_scales): + max_dim = max([s.shape[-1] for s in layer_scales]) + layer_scales = [ + torch.cat((s, torch.zeros((1, max_dim - s.shape[-1]), device=get_accelerator().current_device_name())), + dim=-1) if s.shape[-1] < max_dim else s for s in layer_scales + ] + return torch.cat(layer_scales).unsqueeze(0) + + def merge_scales(self): + all_scales = [] + for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in \ + zip(self.dense_scales, self.qkv_scales, self.mlp4hh_scales, self.mlph4h_scales): + all_scales.append(self.merge_layer_scales([qkv_scale, dense_scale, mh4h_scale, m4hh_scale])) + return torch.cat(all_scales) + + def merge_scales_split(self, split_count): + all_scales = [[] for _ in range(split_count)] + for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in \ + zip(self.dense_scales, self.qkv_scales, self.mlp4hh_scales, self.mlph4h_scales): + dense_scale = torch.split(dense_scale, dense_scale.numel() // split_count) + qkv_scale = torch.split(qkv_scale, qkv_scale.numel() // split_count) + m4hh_scale = torch.split(m4hh_scale, m4hh_scale.numel() // split_count) + mh4h_scale = torch.split(mh4h_scale, mh4h_scale.numel() // split_count) + for s in range(split_count): + all_scales[s].append( + torch.cat([ + torch.cat((qkv_scale[s], torch.zeros_like(qkv_scale[s])), dim=1), + torch.cat((dense_scale[s], torch.zeros_like(dense_scale[s])), dim=1), mh4h_scale[s], + m4hh_scale[s] + ]).unsqueeze(0)) + for scales_a in all_scales: + torch.cat(scales_a) + return all_scales + + def sd_quantize_megatron(self, sd, quantize_bits, groups): + keys = sd.keys() + for key in keys: + value_list = [sd[key]] + if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key or \ + "mlp.dense_h_to_4h.weight" in key or "attention.query_key_value.weight" in key: + value_list = self.Quantize(value_list, quantize_bits, groups, key=key) + sd[key] = value_list[0] + + all_scales = self.merge_scales() + return sd, all_scales + + def model_quantize(self, model, quantize_policy, quantize_bits, groups): + all_scales = [] + + def quantize_fn(layer, policy_cls): + policy = policy_cls(layer) + + _, qkvw, _, dense_w, _, _ = policy.attention() + _, _h4h_w, _, _4hh_w, _ = policy.mlp() + keys = [qkvw, dense_w, _h4h_w, _4hh_w] + layer_scales = [] + + for key in range(len(keys)): + if self.mlp_extra_grouping and self.is_mlp(keys[key]): + data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups * 2) + elif policy_cls is HFBertLayerPolicy and self.is_qkv(keys[key]): + data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups * 3) + else: + data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups) + keys[key].copy_(data_quantized) + layer_scales.append((1 / data_scale.to(get_accelerator().current_device_name()).view(-1).unsqueeze(0))) + all_scales.append(self.merge_layer_scales(layer_scales)) + return layer + + def _quantize_module(model, policies): + for name, child in model.named_children(): + if child.__class__ in policies: + quantize_fn, replace_policy = policies[child.__class__] + setattr(model, name, quantize_fn(child, replace_policy)) + else: + _quantize_module(child, policies) + + return model + + policy = {} + if quantize_policy is not None: + for layer_name, replace_policy in quantize_policy.items(): + policy.update({layer_name: (quantize_fn, replace_policy)}) + else: + for plcy in replace_policies: + policy.update({plcy._orig_layer_class: (quantize_fn, plcy)}) + + quantized_module = _quantize_module(model, policy) + + return quantized_module, torch.cat(all_scales) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b045dbc8f974c8f6403e2edb006b7c14c3d9cc1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .partition_parameters import ZeroParamType +from .partition_parameters import ZeroParamStatus +from .partition_parameters import Init +from .partition_parameters import GatheredParameters +from .partition_parameters import register_external_parameter +from .parameter_offload import DeepSpeedZeRoOffload +from .partition_parameters import DeepSpeedTensorOverride + +from .tiling import TiledLinear +from .tiling import TiledLinearReturnBias + +from .mics import MiCS_Init + +from .stage3 import unwrap_model_for_generation diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..651f08dd511a96e4a21939cde2e2819d0c74c743 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4765943f3629e264afcb9c3f3927ee9de1adc7db Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40967e9016e649a721d53cfdb8c90b26ecfe375f Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bac3dece5ff224eddb484be529608474d4ed8209 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..695917fa3480b80a2db8a838ad855c3f8767ea29 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac86028dad41449f134347f2b7db95afa6be6a05 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b11737c1691b4e101a442be47e1846db48155fbf Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3983e201168ef565e94bbfb72e9df96faf8f6e4e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1c0491dfe2c429e74621e90ccfbc8b23d439200 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81d620b1a031b9ab5f98580c9a2f9135844dad2a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbaf4a12736aa4ed1db876e1975e49afba95b2dd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78d67e0ff01caed58f8ce07de90d03041b7bb242 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f299ff6c7d120b2d80ddaa3a4ff5c17f840db3a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..894765eeada933e775a88e612ee202faf43d216d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0fcdf6fd73ab7b7c8b5232a16c7424c5cc00ce6c Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0cea5e0070e1901cdca9f59b441459facb747fec Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..786bdc412518f1f51790b12ca9aadfda38de7d64 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/config.py new file mode 100644 index 0000000000000000000000000000000000000000..19ee9b51702e786773c3fd1781e3342031e805cc --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/config.py @@ -0,0 +1,361 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +from typing import Optional, Dict, Any +from enum import Enum +from pydantic import Field, model_validator +from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel +from deepspeed.utils import logger +from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum + +# ZeRO optimization. By default, this optimization is not enabled. +# Users have to configure the desired optimization (0 means disabled) in params.json as below example: +ZERO_FORMAT = """ +ZeRO optimization should be enabled as: +"session_params": { + "zero_optimization": { + "stage": [0|1|2], + "stage3_max_live_parameters" : 1000000000, + "stage3_max_reuse_distance" : 1000000000, + "stage3_use_all_reduce_for_fetch_params": [true|false], + "stage3_module_granularity_threshold": 0, + "allgather_partitions": [true|false], + "use_multi_rank_bucket_allreduce": [true|false], + "allgather_bucket_size": 500000000, + "reduce_scatter": [true|false], + "contiguous_gradients" : [true|false] + "overlap_comm": [true|false], + "reduce_bucket_size": 500000000, + "load_from_fp32_weights": [true|false], + "cpu_offload": [true|false] (deprecated), + "cpu_offload_param" : [true|false] (deprecated), + "cpu_offload_use_pin_memory": [true|false] (deprecated), + "sub_group_size" : 1000000000000, + "offload_param": {...}, + "offload_optimizer": {...}, + "ignore_unused_parameters": [true|false], + "round_robin_gradients": [true|false], + "zero_hpz_partition_size": 1, + "zero_quantized_weights": [true|false], + "zero_quantized_nontrainable_weights": [true|false], + "zero_quantized_gradients": [true|false], + "memory_efficient_linear": [true|false], + "override_module_apply": [true|false], + "zeropp_loco_param": {...}, + "log_trace_cache_warnings" : [true|false], + } +} +""" + +ZERO_OPTIMIZATION = "zero_optimization" + + +def read_zero_config_deprecated(param_dict): + zero_config_dict = {} + zero_config_dict["stage"] = 1 if param_dict[ZERO_OPTIMIZATION] else 0 + if zero_config_dict["stage"] > 0: + zero_config_dict["allgather_bucket_size"] = get_scalar_param(param_dict, "allgather_size", 5e8) + logger.warning( + "DeepSpeedConfig: this format of ZeRO optimization setup is deprecated. Please use the following format: {}". + format(ZERO_FORMAT)) + return zero_config_dict + + +def get_zero_config(param_dict): + if ZERO_OPTIMIZATION in param_dict: + zero_config_dict = param_dict[ZERO_OPTIMIZATION] + if isinstance(zero_config_dict, bool): + zero_config_dict = read_zero_config_deprecated(param_dict) + else: + zero_config_dict = {} + return DeepSpeedZeroConfig(**zero_config_dict) + + +class ZeroStageEnum(int, Enum): + """ Enum class for possible zero stages """ + disabled = 0 + optimizer_states = 1 + gradients = 2 + weights = 3 + max_stage = 3 + + +class DeepSpeedZeroConfig(DeepSpeedConfigModel): + """ + Sets parameters for ZeRO optimizations. + """ + + stage: ZeroStageEnum = 0 + """ + Chooses different stages of ZeRO Optimizer. Stage 0, 1, 2, and 3 refer + to disabled, optimizer state partitioning, and optimizer+gradient state + partitioning, and optimizer+gradient+parameter partitioning, respectively. + """ + + contiguous_gradients: bool = True + """ + Copies the gradients to a contiguous buffer as they are produced. Avoids + memory fragmentation during backward pass. + """ + + reduce_scatter: bool = True + """ + Uses reduce or reduce scatter instead of allreduce to average gradients + """ + + reduce_bucket_size: int = Field(pp_int(5e8), ge=0) + """ + Number of elements reduced/allreduced at a time. Limits the memory required + for the allgather for large model sizes + """ + + use_multi_rank_bucket_allreduce: bool = True + """ + Combine the reduce buckets of the different ranks and do an All-Reduce instead of multiple Reduce ops. + This feature is useful when the model is small and we want to scale it on too many GPUs which therefore + reduces the message sizes of each packet. + """ + + allgather_partitions: bool = True + """ + Chooses between allgather collective or a series of broadcast collectives + to gather updated parameters from all the GPUs at the end of each step + """ + + allgather_bucket_size: int = Field(pp_int(5e8), ge=0) + """ + Number of elements allgathered at a time. Limits the memory required for + the allgather for large model sizes + """ + + overlap_comm: Optional[bool] = None # None for dynamic default value (see validator `overlap_comm_valid` below) + """ + Attempts to overlap the reduction of the gradients with backward computation + """ + + load_from_fp32_weights: bool = True + """ + Boolean indicating whether to initialize fp32 master weights from fp32 + copies in checkpoint (no precision loss) or from model's fp16 copies (with + precision loss). This can be used to initialize optimizer state even when + checkpoint is missing optimizer state. + """ + + elastic_checkpoint: bool = False + """ + Enable loading checkpoint that was saved by job with different GPU count. + No longer supported. + """ + + offload_param: Optional[DeepSpeedZeroOffloadParamConfig] = None + """ + Enable offloading of model parameters to CPU or NVMe. This frees up GPU + memory for larger models or batch sizes. Valid only with stage 3. Expects a + dictionary containing values for :any:`DeepSpeedZeroOffloadParamConfig`. + """ + + offload_optimizer: Optional[DeepSpeedZeroOffloadOptimizerConfig] = None + """ + Enable offloading of optimizer state to CPU or NVMe, and optimizer + computation to CPU. This frees up GPU memory for larger models or batch + sizes. Valid for ZeRO stage 1, 2, 3. Expects a dictionary containing values + for :any:`DeepSpeedZeroOffloadOptimizerConfig`. + """ + + sub_group_size: int = Field(pp_int(1e9), ge=0) + """ + Tile size for parameter processing to fit massive models (with trillions of + parameters). Used by ZeRO3-Offload and ZeRO-Infinity + """ + + cpu_offload_param: Optional[bool] = Field( + None, + json_schema_extra={ + "deprecated": True, + "new_param": "offload_param", + "new_param_fn": (lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu) + if val else None) + }, + ) + """ Deprecated, please use ``offload_param`` """ + + cpu_offload_use_pin_memory: Optional[bool] = Field( + None, + json_schema_extra={ + "deprecated": True, + "new_param": "offload_param or offload_optimizer", + "set_new_param": False + }, + ) + """ Deprecated, please use ``offload_param`` or ``offload_optimizer`` """ + + cpu_offload: Optional[bool] = Field( + None, + json_schema_extra={ + "deprecated": + True, + "new_param": + "offload_optimizer", + "new_param_fn": (lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu) + if val else None) + }, + ) + """ Deprecated, please use ``offload_optimizer`` """ + + prefetch_bucket_size: int = Field(pp_int(5e7), ge=0, alias="stage3_prefetch_bucket_size") + """ + Maximum number of parameter elements to fetch ahead of use. Used by ZeRO3, + ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference. + """ + + param_persistence_threshold: int = Field(pp_int(1e5), ge=0, alias="stage3_param_persistence_threshold") + """ + Do not partition parameters smaller than this threshold. Smaller values use + less memory, but can greatly increase communication (especially + latency-bound messages). + """ + + model_persistence_threshold: int = Field(pp_int(sys.maxsize, "sys.maxsize"), + ge=0, + alias="stage3_model_persistence_threshold") + """ + Maximum number of parameter elements that can be persisted in GPU and not + partitioned. This imposes an upper bound on the number of unpartitioned + parameters resulting from param_persistence_threshold setting. Used by + ZeRO3-Offload, ZeRO-Infinity and ZeRO-Inference. + """ + + max_live_parameters: int = Field(pp_int(1e9), ge=0, alias="stage3_max_live_parameters") + """ + The maximum number of parameters resident per GPU before releasing. Smaller + values use less memory, but perform more communication. + """ + + max_reuse_distance: int = Field(pp_int(1e9), ge=0, alias="stage3_max_reuse_distance") + """ + Do not release a parameter if it will be reused within this threshold of + parameters. Smaller values use less memory, but perform more communication. + """ + + gather_16bit_weights_on_model_save: bool = Field(False, alias="stage3_gather_16bit_weights_on_model_save") + """ + Consolidate the weights before saving the model by ``save_16bit_model()``. + Since the weights are partitioned across GPUs, they aren’t part of + ``state_dict``, so this function automatically gathers the weights when + this option is enabled and then saves the fp16 model weights. + """ + + module_granularity_threshold: int = Field(pp_int(0), alias="stage3_module_granularity_threshold") + """ + The granularity of a module is determined by the ratio of "parameter_count / (1 + descendant count)". + ZeRO3 classifies modules with a granularity below the threshold as fine-grained, + which are treated as integral units during parameter fetching. This reduces host overhead + and the separate allgather overhead introduced by hooks for fine-grained layers when fetching parameters. + """ + + use_all_reduce_for_fetch_params: bool = Field(False, alias="stage3_use_all_reduce_for_fetch_params") + """ + Use all_reduce op when fetching module parameters at stage3. This improves performance by reducing + the overhead of concatenation and slicing on the host. + """ + + stage3_gather_fp16_weights_on_model_save: bool = Field(False, + json_schema_extra={ + "deprecated": True, + "new_param": "gather_16bit_weights_on_model_save" + }) + """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """ + + ignore_unused_parameters: bool = True + """ + Unused parameters in modules may be unexpected in static networks, but + could be normal in dynamic networks. This controls whether or not training + should terminate with an error message when unused parameters are detected. + This is set to ``True`` by default, which means unused parameters are + ignored and training continues. Now is just used in stage 2. + """ + + legacy_stage1: bool = False + """ + For backward-compatibility enable old ZeRO stage 1 implementation. Use at + your own risk, will be deprecated soon. + """ + + round_robin_gradients: bool = False + """ + Stage 1 and 2 optimization for CPU offloading that parallelizes gradient + copying to CPU memory among ranks by fine-grained gradient partitioning. + Performance benefit grows with gradient accumulation steps (more copying + between optimizer steps) or GPU count (increased parallelism). + """ + zero_hpz_partition_size: int = Field(1, ge=0) + """ + Number of ranks in zero parameters partitioning secondary group + """ + zero_quantized_weights: bool = False + """ + Boolean indicating whether to quantize zero parameters (weights) + for efficient all_gather comm + """ + zero_quantized_nontrainable_weights: bool = False + """ + Boolean indicating whether to quantize non-trainable zero parameters (weights) + for efficient memory usage and communication. Different from zero_quantized_weights + that stores the weights in original precision and only perform quantization during communication, + this flag will store the weights in quantized precision. This is useful for LoRA training. + """ + zero_quantized_gradients: bool = False + """ + Boolean indicating whether to use quantized zero gradients + for efficient all_2_all_reduce comm + """ + zeropp_loco_param: Optional[Dict[str, Any]] = None + """ + This dictionary contains parameters for using LoCo-Zero++, with two key parameters: + - `err_beta`: A coefficient for the moving average of quantization errors before and after gradient computation. + It ranges between 0 and 1, with a default value of 0.8. + - `reset_T`: The number of steps after which the moving-average error buffer is cleared. The default value is 1024. + These parameters can be adjusted based on performance needs. Example configuration in ds config: + "zeropp_loco_param": { "err_beta": 0.8, "reset_T": 1024 }. + See LoCo paper for more details: (https://arxiv.org/abs/2407.04480). + """ + + mics_shard_size: int = Field(-1, json_schema_extra={"new_param": "mics_shard_size"}) + + mics_hierarchical_params_gather: bool = False + + memory_efficient_linear: bool = True + """ + Use memory efficient linear implementation, for Stage 3. + """ + """ + Whether force load checkpoint in pipeline mode, current only for Stage 3. + """ + pipeline_loading_checkpoint: bool = False + + override_module_apply: bool = True + """ + Override nn.Module apply function, for Stage 3. + """ + + log_trace_cache_warnings: bool = False + """ + Whether to log warnings from trace cache, such as invalidation events. + """ + + # Validators + @model_validator(mode="after") + def overlap_comm_valid(self): + if self.overlap_comm is None: + self.overlap_comm = self.stage == ZeroStageEnum.weights + return self + + @model_validator(mode="after") + def offload_ratio_check(self): + offload_config = self.offload_optimizer + if offload_config and offload_config.ratio < 1.0: + assert self.stage == ZeroStageEnum.weights, "Partial offloading only supported for ZeRO Stage 3." + return self diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py new file mode 100644 index 0000000000000000000000000000000000000000..35b3d5c7dd5d333e621cc554a73bb3c90e2a450e --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py @@ -0,0 +1,287 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from deepspeed import comm as dist + + +def print_rank_0(message): + if dist.get_rank() == 0: + print(message) + + +class ContiguousMemoryAllocator(object): + + def __init__(self, size, dtype, device): + self.buffer = torch.zeros(size, dtype=dtype, device=device) + + #address to contiguous size available + self.contiguous_sizes = {} + + self.contiguous_sizes[0] = size + + #tensor id to its address + self.tensor_addresses = {} + + #tensor address to its size + self.tensor_sizes = {} + + #tensor address to ids + self.tensor_ids = {} + + #id to tensors + self.tensor_map = {} + + #id to params. Maps each tensor buffer to list of parameters that uses it + self.id_to_params = {} + + self.total_size = size + self.total_free = size + self.largest_contiguous = size + self.max_allocated = 0 + + self.count = 0 + + #create a tensor of size from the pre-allocated buffer + #if not enough free space will fail + #if not enough contiguous space, will defragment and allocate + def allocate_tensor(self, size): + free_before = self.total_free + + assert size <= self.total_free, "Not enough memory in buffer. Allocation failed" + if self.largest_contiguous < size: + print_rank_0("Needs defragmentation to allocate. Before Defragmentation:") + self.print_allocation(resolution=100) + self._defragment_memory() + #set the param data to the new tensor buffer locations + self._reset_param_data() + print_rank_0("After defragmentation:") + self.print_allocation(resolution=100) + + self.total_free = self.total_free - size + + allocated = self.total_size - self.total_free + if allocated > self.max_allocated: + self.max_allocated = allocated + + tensor_address = self._get_new_tensor_address(size) + + ret_tensor = self._get_new_tensor(tensor_address, size) + print_rank_0( + f"Free before allocation {free_before}. Allocating {size}. Free after allocation {self.total_free}. Max allocated {self.max_allocated}" + ) + assert self.total_free + size == free_before, "Allocation bookkeeping error" + + return ret_tensor + + #assigns the tensor data to the param data and keeps track of the assignment + #any change the underlying buffer from defragmentation will cause a + #reassignment of the param data + def assign_to_param(self, tensor, param, numel, shape): + tensor_id = id(tensor) + + assert tensor_id in self.tensor_map.keys(), "No such tensor allocated by the allocator." + assert tensor.numel() >= numel, "Assert tensor buffer does is not large enough" + assert not tensor_id in self.id_to_params.keys(), "This tensor has already been assigned to a param" + + self.id_to_params[tensor_id] = [param] + + replicated_tensor = tensor.narrow(0, 0, numel).view(shape) + param.data = replicated_tensor.data + param.contiguous_tensor_id = tensor_id + + #deletes the tensor and frees up the underlying buffer + def release_tensor(self, tensor): + free_before = self.total_free + tensor_id = id(tensor) + tensor_size = tensor.numel() + self._release_tensor(tensor_id) + self._unassign_params(tensor_id) + self.total_free += tensor_size + print_rank_0( + f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.") + assert self.total_free - tensor_size == free_before, "Release bookkeeping error" + + def release_tensor_with_id(self, tensor_id): + free_before = self.total_free + assert tensor_id in self.tensor_map.keys(), "Invalid tensor id" + tensor = self.tensor_map[tensor_id] + tensor_size = tensor.numel() + self._release_tensor(tensor_id) + self._unassign_params(tensor_id) + self.total_free += tensor_size + print_rank_0( + f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.") + assert self.total_free - tensor_size == free_before, "Release bookkeeping error" + + #shows the current memory allocation at specified resolution + def print_allocation(self, resolution=200): + total_size = self.buffer.numel() * 1.0 + empty = [] + for addr, size in self.contiguous_sizes.items(): + start = int(addr * resolution / total_size) + end = int((addr + size) * resolution / total_size) + empty.extend(range(start, end)) + s = '' + for i in range(resolution): + s += '.' if i in empty else '|' + print_rank_0(s) + + def max_allocated(self): + return self.max_allocated + + #to be called after defragmentation that moves the tensor buffers + #this call reassigns the data of all the parameters using the tensor buffers + def _reset_param_data(self): + for id, tensor in self.tensor_map.items(): + for param in self.id_to_params[id]: + param.data = tensor.narrow(0, 0, param.numel()).view(param.data.shape).data + + def _unassign_params(self, tensor_id): + if tensor_id in self.id_to_params.keys(): + del self.id_to_params[tensor_id] + + def _release_tensor(self, tensor_id): + assert tensor_id in self.tensor_addresses, f"Tensor id {tensor_id} not found" + + address = self.tensor_addresses[tensor_id] + contiguous_size = self.tensor_map[tensor_id].numel() + + del self.tensor_addresses[tensor_id] + del self.tensor_ids[address] + del self.tensor_map[tensor_id] + del self.tensor_sizes[address] + + self._consolidate_address(address, contiguous_size) + self.largest_contiguous = self._largest_contiguous() + + def _consolidate_address(self, address, contiguous_size): + + #consolidate next buffer + end_address = address + contiguous_size + if end_address in self.contiguous_sizes: + contiguous_size += self.contiguous_sizes[end_address] + del self.contiguous_sizes[end_address] + + #consolidate previous buffer + for addr, size in self.contiguous_sizes.items(): + if addr + size == address: + del self.contiguous_sizes[addr] + contiguous_size += size + address = addr + break + + self.contiguous_sizes[address] = contiguous_size + + def _defragment_memory(self): + empty_addresses = sorted(self.contiguous_sizes.keys()) + tensor_addresses = sorted(self.tensor_addresses.values()) + + tensor_index = 0 + + while tensor_index < len(tensor_addresses): + + empty_addr = empty_addresses[0] + empty_size = self.contiguous_sizes[empty_addr] + + tensor_addr = tensor_addresses[tensor_index] + tensor_size = self.tensor_sizes[tensor_addr] + tensor_id = self.tensor_ids[tensor_addr] + tensor = self.tensor_map[self.tensor_ids[tensor_addr]] + + assert tensor_size == tensor.numel(), \ + f"Size mismatch. {tensor_size} is allocated at addr {tensor_addr} but tensor size is {tensor.numel()} " + + assert empty_addr != tensor_addr, \ + f"Cannot have same empty address {empty_addr} and tensor address {tensor_addr}" + + if empty_addr < tensor_addr: + + if empty_size >= tensor_size: + dest_buffer = self.buffer.narrow(0, empty_addr, tensor_size) + src_buffer = self.buffer.narrow(0, tensor_addr, tensor_size) + dest_buffer.data.copy_(src_buffer.data) + else: + + #print_rank_0(f'empty addr : {empty_addr}, empty size {empty_size} tensor addr {tensor_addr} tensor size {tensor_size}') + src_addr = tensor_addr + dest_addr = empty_addr + while src_addr < (tensor_addr + tensor_size): + copy_size = min(empty_size, tensor_addr + tensor_size - src_addr) + + dest_buffer = self.buffer.narrow(0, dest_addr, copy_size) + src_buffer = self.buffer.narrow(0, src_addr, copy_size) + + dest_buffer.data.copy_(src_buffer.data) + + src_addr += copy_size + dest_addr += copy_size + + self._replace_old_address_with_new(tensor_id, empty_addr) + + tensor_index += 1 + + else: + tensor_index += 1 + + empty_addresses = sorted(self.contiguous_sizes.keys()) + + def _replace_old_address_with_new(self, tensor_id, new_address): + + tensor = self.tensor_map[tensor_id] + tensor_size = tensor.numel() + tensor.data = self.buffer.narrow(0, new_address, tensor_size).data + + self._release_tensor(tensor_id) + self._mark_as_occupied(new_address, tensor_size) + + self.tensor_ids[new_address] = tensor_id + self.tensor_map[tensor_id] = tensor + self.tensor_addresses[tensor_id] = new_address + self.tensor_sizes[new_address] = tensor_size + + def _get_new_tensor_address(self, size): + tensor_address = None + for address, contiguous_size in self.contiguous_sizes.items(): + if contiguous_size >= size and \ + (tensor_address is None or \ + contiguous_size < self.contiguous_sizes[tensor_address]): + tensor_address = address + assert tensor_address is not None, "address cannot be None" + return tensor_address + + def _get_new_tensor(self, address, size): + available_contiguous_size = self.contiguous_sizes[address] + + assert size <= available_contiguous_size, \ + f"Tensor numel {size} is large than available contiguous size {available_contiguous_size}" + self.count += 1 + new_tensor = self.buffer.narrow(0, address, size) + tensor_id = id(new_tensor) + self.tensor_addresses[tensor_id] = address + self.tensor_sizes[address] = size + + self.tensor_ids[address] = tensor_id + self.tensor_map[tensor_id] = new_tensor + + self._mark_as_occupied(address, size) + + return new_tensor + + def _largest_contiguous(self): + if len(self.contiguous_sizes) > 0: + return max([size for _, size in self.contiguous_sizes.items()]) + else: + return 0 + + def _mark_as_occupied(self, address, size): + available_contiguous_size = self.contiguous_sizes[address] + del self.contiguous_sizes[address] + + if available_contiguous_size != size: + self.contiguous_sizes[address + size] = available_contiguous_size - size + + self.largest_contiguous = self._largest_contiguous() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/linear.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..8c8db60768ebca9f7f972c13e0d3d20176158c96 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/linear.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +#Linear Module to use with ZeRO Stage 3 to allow for parameter memory release +#after the module execution during forward +#Instead of saving variables using save_for_backward, we save variable ids +#Allowing us to retrieve the variable without creating pointer to it +#Which allows for underlying tensor to be garbage collected +#When partitioned as needed by the Zero Stage 3 optimizer +#TODO instead of patching Linear module, we could patch the ctx.save_for_backward +#ctx.saved_tensors so that this approach works for all nn modules that are built upon +#torch.nn.function. However the issue is that many modules uses C++ implementations +#which does not have pytorch implementation. Eg torch.addmm which acts as a functional +#when implemented outside of torch.autograd.Function + +import math +import functools + +import torch +from torch import Tensor +from torch.nn.parameter import Parameter +from torch.nn import init +from torch.nn.modules.module import Module +from deepspeed.runtime.utils import noop_decorator +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator + + +def print_rank_0(message, debug=False, force=False): + if dist.get_rank() == 0 and (debug or force): + print(message) + + +try: + # Fix `torch.[device].amp.custom_fwd/bwd` FutureWarning in torch 2.4 + if hasattr(torch, 'amp') and hasattr(torch.amp, 'custom_fwd') and hasattr(torch.amp, 'custom_bwd'): + autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=get_accelerator().device_name()) + autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=get_accelerator().device_name()) + else: + # original implementation + autocast_custom_fwd = get_accelerator().amp().custom_fwd + autocast_custom_bwd = get_accelerator().amp().custom_bwd +except (ImportError, AttributeError) as exp: + autocast_custom_fwd = noop_decorator + autocast_custom_bwd = noop_decorator + + +class LinearFunctionForZeroStage3(torch.autograd.Function): + + # Note that both forward and backward are @staticmethods + @staticmethod + @autocast_custom_fwd + # bias is an optional argument + def forward(ctx, input, weight, bias=None): + + ctx.save_for_backward(input, weight, bias) + + if input.dim() == 2 and bias is not None: + # fused op is marginally faster + ret = torch.addmm(bias, input, weight.t()) + else: + output = input.matmul(weight.t()) + if bias is not None: + output += bias + ret = output + + return ret + + # This function has only a single output, so it gets only one gradient + @staticmethod + @autocast_custom_bwd + def backward(ctx, grad_output): + # This is a pattern that is very convenient - at the top of backward + # unpack saved_tensors and initialize all gradients w.r.t. inputs to + # None. Thanks to the fact that additional trailing Nones are + # ignored, the return statement is simple even when the function has + # optional inputs. + input, weight, bias = ctx.saved_tensors + + grad_input = grad_weight = grad_bias = None + + #print(f"backward shaped grad_output {grad_output.shape}, input {input.shape}, weight {weight.shape} and bias {bias.shape if bias is not None else None}") + # These needs_input_grad checks are optional and there only to + # improve efficiency. If you want to make your code simpler, you can + # skip them. Returning gradients for inputs that don't require it is + # not an error. + if ctx.needs_input_grad[0]: + #print(f"Computing grad input weight {weight.shape} grad_output {grad_output.shape}") + grad_input = grad_output.matmul(weight) + #print(f"Computed grad input {grad_input.shape}") + if ctx.needs_input_grad[1]: + #print("Computing grad weight") + dim = grad_output.dim() + if dim > 2: + grad_weight = grad_output.reshape(-1, + grad_output.shape[-1]).t().matmul(input.reshape(-1, input.shape[-1])) + else: + grad_weight = grad_output.t().matmul(input) + #print(f"Computed grad weight grad_weight {grad_weight.shape}") + if bias is not None and ctx.needs_input_grad[2]: + #print("Computing grad bias") + if dim > 2: + grad_bias = grad_output.sum([i for i in range(dim - 1)]) + else: + grad_bias = grad_output.sum(0) + #print("Done computing grad bias") + #print("needs bias") + #print(f"backward shaped grad_input {grad_input.shape}, grad_weight {grad_weight.shape}, grad_bias {grad_bias.shape if grad_bias is not None else None}") + return grad_input, grad_weight, grad_bias + + +def zero3_linear_wrap(input, weight, bias=None): + if bias is None: + return LinearFunctionForZeroStage3.apply(input, weight) + else: + return LinearFunctionForZeroStage3.apply(input, weight, bias) + + +class LinearModuleForZeroStage3(Module): + r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`. + The weights are pre-transposed and stored as A^T instead of transposing during each + forward. Memory savings proportional to the parameter size. + + Args: + in_features: size of each input sample + out_features: size of each output sample + bias: If set to ``False``, the layer will not learn an additive bias. + Default: ``True`` + + Shape: + - Input: :math:`(N, *, H_{in})` where :math:`*` means any number of + additional dimensions and :math:`H_{in} = \text{in\_features}` + - Output: :math:`(N, *, H_{out})` where all but the last dimension + are the same shape as the input and :math:`H_{out} = \text{out\_features}`. + + Attributes: + weight: the learnable weights of the module of shape + :math:`(\text{out\_features}, \text{in\_features})`. The values are + initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where + :math:`k = \frac{1}{\text{in\_features}}` + bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. + If :attr:`bias` is ``True``, the values are initialized from + :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where + :math:`k = \frac{1}{\text{in\_features}}` + + Examples:: + + >>> m = nn.Linear(20, 30) + >>> input = torch.randn(128, 20) + >>> output = m(input) + >>> print(output.size()) + torch.Size([128, 30]) + """ + __constants__ = ['in_features', 'out_features'] + in_features: int + out_features: int + weight: Tensor + + def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: + super(LinearModuleForZeroStage3, self).__init__() + print("Building ZeRO module") + self.in_features = in_features + self.out_features = out_features + self.weight = Parameter(torch.Tensor(out_features, in_features)) + if bias: + self.bias = Parameter(torch.Tensor(out_features)) + else: + self.register_parameter('bias', None) + self.reset_parameters() + + def reset_parameters(self) -> None: + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) + bound = 1 / math.sqrt(fan_in) + init.uniform_(self.bias, -bound, bound) + + def forward(self, input: Tensor) -> Tensor: + return LinearFunctionForZeroStage3.apply(input, self.weight, self.bias) + + def extra_repr(self) -> str: + return 'in_features={}, out_features={}, bias={}'.format(self.in_features, self.out_features, self.bias + is not None) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics.py new file mode 100644 index 0000000000000000000000000000000000000000..628bf86a61da23182b2f33cc7b6edab2c1aee41f --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics.py @@ -0,0 +1,472 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sys +from typing import List + +import deepspeed +import torch +from deepspeed import comm as dist +from deepspeed.runtime.zero.utils import is_zero_param +from deepspeed.runtime.zero.mics_utils import (MiCS_CommGroups, create_mics_comm_groups, scale_tensors) +from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload +from deepspeed.runtime.zero.partition_parameters import Init, AllGatherCoalescedHandle, ZeroParamStatus +from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 +from deepspeed.utils import instrument_w_nvtx, log_dist, logger +from deepspeed.accelerator import get_accelerator +from torch import Tensor +from torch.nn import Parameter + + +def has_hierarchical_all_gather_groups(comm_groups: MiCS_CommGroups): + result = False + if comm_groups.param_intra_node_group is not None and comm_groups.param_inter_node_shard_group is not None: + result = True + return result + + +class MiCS_AllGatherCoalescedHandle(AllGatherCoalescedHandle): + """ This handle assumes that no need to + copy data out from a contiguous tensor + """ + + def __init__(self, allgather_handle, params: List[Parameter], partitions: List[Tensor], world_size: int) -> None: + super().__init__(allgather_handle, params, partitions, world_size) + + def wait(self, **kwargs) -> None: + """ + """ + # let the current stream to op + try: + # print("HANDLE", self.allgather_handle) + instrument_w_nvtx(self.allgather_handle.wait)() + except (ValueError, RuntimeError) as e: + log_dist( + f"WARNING: Runtime Error while waiting the collective all-gather, possibly due to the _IllegalWork", + ranks=[0]) + log_dist(f"Error message: {e}", ranks=[0]) + + if self.complete: + return + + for _, param in enumerate(self.params): + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + param.ds_status = ZeroParamStatus.AVAILABLE + + self.complete = True + + +class MiCS_Init(Init): + + def __init__(self, + module=None, + data_parallel_group=None, + sequence_data_parallel_group=None, + mem_efficient_linear=True, + remote_device=None, + pin_memory=False, + config_dict_or_path=None, + config=None, + enabled=True, + dtype=None, + mpu=None): + """A context manager to partition the model parameters during the model + construction with MiCS partition strategy. Model states are partitioned + to the number of devices specified via ``mics_shard_size`` field in the + deepspeed config json file. The context manager also introduces + hierarchical communication method to reduce the cost of inter-node + communications, which can be enabled with + ``mics_hierarchical_params_gather`` field in deepspeed config. + + Args: + module (``torch.nn.Module``, optional): If provided, partition the model as + if it was constructed in the context. + data_parallel_group (``deepspeed.comm`` process group, optional): + The group of processes to partition among. Defaults to all processes. + Synonymous with sequence data parallel group for param partitioning + across both sequence and data parallel groups. + mem_efficient_linear (bool, optional): Replace + torch.nn.functional.linear with an implementation that allows + DeepSpeed to partition parameters. Defaults to ``True``. + remote_device (string, optional): The initial device to store model + weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU + memory. The model may still be moved to GPU based on the + offload settings for training. Defaults to param offload device if a config is + defined, otherwise GPU. + pin_memory (bool, optional): Potentially increase performance by + using pinned memory for model weights. ``remote_device`` must be + ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``. + config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration + for swapping fp16 params to NVMe. + config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. + enabled (bool, optional): If ``False``, this context has no + effect. Defaults to ``True``. + dtype (``dtype``, optional): Can be used to change the data type of the parameters. + Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` + mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}. + + This context follows the same logic as ``deepspeed.zero.Init()``, but + with the modification for partition size of each parameter. + + Examples + -------- + + #. Allocate a model and partition it among all processes: + + .. code-block:: python + # the config_dict_or_path is required to let the context manager know + # how partition the parameters. + # The configuration has to include the field ``mics_shard_size`` + with deepspeed.zero.MiCS_Init(config_dict_or_path=ds_config): + model = MyLargeModel() + + + #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes: + + .. code-block:: python + + with deepspeed.zero.MiCS_Init(data_parallel_group=mpu.get_data_parallel_group(), + remote_device="cpu", + pin_memory=True + config_dict_or_path=ds_config): + model = MyLargeModel() + + + #. Partition an already-allocated model in CPU memory: + + .. code-block:: python + + model = deepspeed.zero.MiCS_Init(module=model, + config_dict_or_path=ds_config) + """ + + assert config_dict_or_path is not None, "Must provide configuration for MiCS Initialization" + _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, mpu) + if not dist.is_initialized(): + dist.init_distributed() + assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm" + + if data_parallel_group is None: + ds_process_group = dist.get_world_group() + else: + ds_process_group = data_parallel_group + + if sequence_data_parallel_group is not None: + logger.warning( + f"sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.") + if data_parallel_group is not None: + raise ValueError( + "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments." + ) + self.ds_process_group = sequence_data_parallel_group + + self.mics_comm_groups = create_mics_comm_groups( + _ds_config.mics_shard_size, + ds_process_group, + hierarchical_allgather=_ds_config.mics_hierarchial_params_gather, + mpu=mpu) + + super().__init__(module, data_parallel_group, mem_efficient_linear, remote_device, pin_memory, + config_dict_or_path, config, enabled, dtype, mpu) + + def _convert_to_deepspeed_param(self, param): + super()._convert_to_deepspeed_param(param) + # attach communication groups to every param + param.comm = self.mics_comm_groups + + # record existing all_gather_coalesced implementation + # so that we can fallback later + old_all_gather_coalesced = param.all_gather_coalesced + + def _param_all_gather_coalesced(params, param_buffers=None, **kwargs): + """""" + mics_comm_groups: MiCS_CommGroups = params[0].comm + hierarchical_all_gather = has_hierarchical_all_gather_groups(mics_comm_groups) + if dist.has_coalescing_manager() and hierarchical_all_gather: + return self._hierarchical_all_gather_params(params, param_buffers) + elif dist.has_coalescing_manager(): + return self._flat_all_gather_with_coalescing_manager(params, param_buffers) + else: + return old_all_gather_coalesced(params, **kwargs) + + # change the all_gather_coalesced method + param.all_gather_coalesced = _param_all_gather_coalesced + + def _pre_all_gather(self, params, params_buffers=None): + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(params) + + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(param.ds_summary()) + param.ds_status = ZeroParamStatus.INFLIGHT + + # ensure that each rank has params in same order. the allgather + # is done by flattening the parameter list into a single tensor that + # can be allgathered in a single call - this means that if each rank + # gives a list of the same parameters in a different order we will + # silently get incorrect parameter values, and have very difficult + # to debug correctness issues. + params = sorted(params, key=lambda p: p.ds_id) + return params, params_buffers + + def _flat_all_gather_with_coalescing_manager(self, params, params_buffers=None): + """""" + # must have to change the status of the param + # and ensure they are on the device + params, params_buffers = self._pre_all_gather(params, params_buffers) + + mics_comm_groups: MiCS_CommGroups = params[0].comm + param_shard_size = mics_comm_groups.param_shard_size + + output_tensors = [] + input_tensors = [] + for i, p in enumerate(params): + t_size = p.ds_tensor.ds_numel * param_shard_size + if params_buffers is not None and params_buffers[i] is not None: + assert params_buffers[i].numel( + ) == t_size, f'params_to_gather_buffers[{i}] size {params_buffers[i].numel()} does not match with t_size {t_size}' + flat_out = params_buffers[i] + else: + flat_out = torch.empty(t_size, dtype=p.dtype, device=self.local_device, requires_grad=False).view(-1) + output_tensors.append(flat_out) + _flat_input = p.ds_tensor.data.view(-1) + input_tensors.append(_flat_input) + + all_gather_handle = dist.all_gather_coalesced(output_tensors, + input_tensors, + group=mics_comm_groups.param_shard_group, + async_op=True) + + for idx, param in enumerate(params): + param.data = output_tensors[idx].narrow(0, 0, param.ds_numel).view(param.ds_shape).data + + return MiCS_AllGatherCoalescedHandle(allgather_handle=all_gather_handle, + params=params, + partitions=[], + world_size=param_shard_size) + + def _hierarchical_all_gather_params(self, params, params_buffers=None): + """""" + params, params_buffers = self._pre_all_gather(params, params_buffers) + + mics_comm_groups: MiCS_CommGroups = params[0].comm + local_rank = dist.get_rank(group=mics_comm_groups.param_intra_node_group) + inter_node_comm_group = mics_comm_groups.param_inter_node_shard_group + intra_node_comm_group = mics_comm_groups.param_intra_node_group + param_shard_size = mics_comm_groups.param_shard_size + + inter_node_size = dist.get_world_size(group=inter_node_comm_group) + intra_node_size = dist.get_world_size(group=intra_node_comm_group) + param_tensors = [] + for i, p in enumerate(params): + param_size = p.ds_tensor.ds_numel * param_shard_size + if params_buffers is not None and params_buffers[i] is not None: + assert params_buffers[i].numel( + ) == param_size, f'param_buffers[{i}] size {params_buffers[i].numel()} does not match with param_size {param_size}' + param_tensor = params_buffers[i] + else: + param_tensor = torch.empty(param_size, dtype=p.dtype, device=self.local_device, + requires_grad=False).view(-1) + param_tensors.append(param_tensor) + + # inter node all-gather + inter_outputs = [] + inter_inputs = [] + for i, p in enumerate(params): + inter_size = p.ds_tensor.ds_numel * inter_node_size + _out = param_tensors[i].narrow(0, local_rank * inter_size, inter_size) + inter_outputs.append(_out) + inter_inputs.append(p.ds_tensor.data.view(-1).to(self.local_device)) + # sync enqueue + dist.all_gather_coalesced(inter_outputs, inter_inputs, group=inter_node_comm_group, async_op=False) + + # intra node all-gather + intra_outputs = [] + intra_inputs = [] + for i, p in enumerate(params): + # partition param into multiple chunks for allgather + # because inter-node all-gather outputs are in a continues memory + # while in param memory, those inter-node data are placed in different + # location. + # each chunk is an intra-node output + param_chunk = param_tensors[i].view( + (inter_node_size, intra_node_size, p.ds_tensor.ds_numel)).narrow(1, local_rank, 1) + param_chunk.copy_(inter_outputs[i].detach().clone().view(param_chunk.size())) + output_chunks = torch.chunk(param_tensors[i], inter_node_size) + for j, _out in enumerate(output_chunks): + intra_chunk_size = intra_node_size * p.ds_tensor.ds_numel + local_offset = local_rank * p.ds_tensor.ds_numel + _in = param_tensors[i].narrow(0, j * intra_chunk_size + local_offset, p.ds_tensor.ds_numel) + intra_outputs.append(_out) + intra_inputs.append(_in) + + all_gather_handle = dist.all_gather_coalesced(intra_outputs, + intra_inputs, + group=intra_node_comm_group, + async_op=True) + for i, param in enumerate(params): + param.data = param_tensors[i].narrow(0, 0, param.ds_numel).view(param.ds_shape).data + + return MiCS_AllGatherCoalescedHandle( + allgather_handle=all_gather_handle, + params=params, + partitions=[], + world_size=param_shard_size, + ) + + def get_partition_dp_group(self, param): + return param.comm.param_shard_group + + def get_partition_rank(self): + return self.mics_comm_groups.param_shard_rank + + @property + def num_partitions(self): + return self.mics_comm_groups.param_shard_size + + +class MiCS_Offload(DeepSpeedZeRoOffload): + """ Wrapper to change the behavior for parameter sharding + """ + + def _convert_to_zero_parameters(self, ds_config, module, mpu): + """ overload the parent class function for convert the parameters + + """ + log_dist(f'Convert to zero parameters from MiCS Offload manager', ranks=[0]) + non_zero_params = [p for p in module.parameters() if not is_zero_param(p)] + if non_zero_params: + zero_params = [p for p in module.parameters() if is_zero_param(p)] + if zero_params: + zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) + else: + group = None + if mpu: + group = mpu.get_data_parallel_group() + + MiCS_Init(module=module, + data_parallel_group=group, + dtype=self.dtype, + config_dict_or_path=ds_config, + remote_device=self.offload_device, + pin_memory=self.offload_param_pin_memory, + mpu=mpu) + + +class MiCS_Optimizer(DeepSpeedZeroOptimizer_Stage3): + """ + MiCS Optimizer + """ + + def __init__(self, + module, + init_optimizer, + timers, + ds_config, + static_loss_scale=1, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + prefetch_bucket_size=50000000, + max_reuse_distance=1000000000, + max_live_parameters=1000000000, + param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, + dp_process_group=None, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + offload_param_config=None, + sub_group_size=1000000000000, + offload_ratio=0.0, + mpu=None, + clip_grad=0, + gradient_accumulation_dtype=torch.float16, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1, + gradient_accumulation_steps=1, + elastic_checkpoint=False, + aio_config=None): + + log_dist("Init MiCS optimizer", ranks=[0]) + super().__init__(module, init_optimizer, timers, ds_config, static_loss_scale, dynamic_loss_scale, + dynamic_loss_args, verbose, contiguous_gradients, reduce_bucket_size, prefetch_bucket_size, + max_reuse_distance, max_live_parameters, param_persistence_threshold, + model_persistence_threshold, dp_process_group, reduce_scatter, overlap_comm, + offload_optimizer_config, offload_param_config, sub_group_size, offload_ratio, mpu, clip_grad, + gradient_accumulation_dtype, communication_data_type, postscale_gradients, + gradient_predivide_factor, gradient_accumulation_steps, elastic_checkpoint, aio_config) + first_param = next(module.parameters()) + # overload the dp_process_group and partition_count + assert hasattr(first_param, "comm"), " ".join([ + "Sharded parameters don't have the MiCS_CommGroups attached.", + "Might due to the use of deepspeed.zero.Init context for initializing the weights.", + "To use MiCS sharding, please use deepspeed.zero.MiCS_Init instead for initializing parameter." + ]) + self.dp_process_group = first_param.comm.param_shard_group + self.partition_count = first_param.comm.param_shard_size + + def initialize_ds_offload( + self, + *args, + **kwargs, + ): + return MiCS_Offload(*args, **kwargs) + + def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: + grad_buffers = super().partition_grads(params_to_release, grad_partitions) + # perform all-reduce among replication groups + # the function will perform accumulation boundary check + self.allreduce_mics_shard_grads(params_to_release, grad_buffers) + + @instrument_w_nvtx + def allreduce_mics_shard_grads(self, params, partitioned_grads_buffers: List[Tensor]): + """ + """ + # TODO: improve the condition check + if not self.is_gradient_accumulation_boundary or \ + len(partitioned_grads_buffers) == 0: + return + + mics_comm_groups: MiCS_CommGroups = params[0].comm + param_repli_group = mics_comm_groups.param_repli_group + param_repli_size = mics_comm_groups.param_repli_size + + if param_repli_size is None or param_repli_size <= 1: + return + if not get_accelerator().on_accelerator(partitioned_grads_buffers[0]): + raise RuntimeError("Local sharding has no support for CPU offloading") + + if dist.has_all_reduce_coalesced(): + scale_tensors(partitioned_grads_buffers, param_repli_size) + dist.all_reduce_coalesced(tensors=partitioned_grads_buffers, group=param_repli_group) + else: + # manually coalescing all-reduce + aggregated_buffer: Tensor = torch.cat(partitioned_grads_buffers) + aggregated_buffer.div_(param_repli_size) + dist.all_reduce(aggregated_buffer, group=param_repli_group) + offset = 0 + for grad_buff in partitioned_grads_buffers: + grad_buff.view(-1).copy_(aggregated_buffer.narrow(0, offset, grad_buff.numel())) + offset += grad_buff.numel() + + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None): + r""" Loading the ZeRO-3/MiCS partitioned checkpoints + Because the self.dp_process_group is replaced with the communicator for + partition group we can call the load_state_dict logic from ZeRO-3. + """ + super().load_state_dict(state_dict_list, load_optimizer_states, load_from_fp32_weights, checkpoint_folder) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics_utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..06b83160bd6c4cd93f5250a12c67f5abda01d920 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics_utils.py @@ -0,0 +1,203 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from dataclasses import dataclass +from typing import List + +import numpy as np +import torch +from torch import Tensor + +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.utils import logger + + +def _log_rank0(msg): + if dist.get_rank() == 0: + logger.info(msg) + + +@torch.jit.script +def scale_tensors(tensors: List[Tensor], scale: int): + for t in tensors: + t.div_(scale) + + +@dataclass +class MiCS_CommGroups: + """""" + param_shard_group = None + param_shard_size = -1 + param_shard_rank = -1 + + param_repli_group = None + param_repli_size = -1 + param_repli_rank = -1 + + param_intra_node_group = None + param_inter_node_shard_group = None + + +def create_mics_comm_groups( + shard_size, + dp_group, + hierarchical_allgather=False, + mpu=None, +): + """ + create shard-group, replicate-group from config_file + TODO: consider broadcast the config from rank0 + + Returns: + MiCS_CommGroups + """ + # env var for debugging purpose + ndevices_per_node = int(os.environ.get("NDEV_PER_NODE", get_accelerator().device_count())) + _log_rank0(f'creating MiCS communication groups with per node device size {ndevices_per_node}') + groups = MiCS_CommGroups() + + if mpu is not None: + assert dp_group == mpu.get_data_parallel_group() + + # full size of the world + world_size = dist.get_world_size() + # global rank + global_rank = dist.get_rank() + + config = _generate_mics_config(world_size, ndevices_per_node, shard_size, 1) + ranks_of_shard_group = config['shard_groups'] + ranks_of_repli_group = config['replicate_groups'] + if len(ranks_of_repli_group) == 0: + assert len(ranks_of_shard_group) == 1, "replicate groups are empty only for single shard group" + for r in ranks_of_shard_group[0]: + ranks_of_repli_group.append([r]) + + # for simplicity + assert _sizes_all_same(ranks_of_repli_group), "replicate groups must have the same size" + assert _sizes_all_same(ranks_of_shard_group), "shard groups must have the same size" + + assert sum([len(g) for g in ranks_of_shard_group]) == dist.get_world_size(), "all sharded ranks " + if len(ranks_of_shard_group) > 1: # if only shard on one group then no need for replicate groups + assert len(ranks_of_shard_group) == len( + ranks_of_repli_group[0]), "number of shard groups must equal to the size of each replicate group" + + global_rank = dist.get_rank() + # create shard groups + for shard_ranks in ranks_of_shard_group: + _group = dist.new_group(shard_ranks) + if global_rank in shard_ranks: + groups.param_shard_group = _group + groups.param_shard_size = len(shard_ranks) + groups.param_shard_rank = dist.get_rank(_group) + logger.info(f'rank {global_rank}, shard group' + f' {groups.param_shard_rank}/{dist.get_world_size(group=_group)}') + + # create replicate groups + for repli_ranks in ranks_of_repli_group: + if len(repli_ranks) > 1: + _group = dist.new_group(repli_ranks) + if global_rank in repli_ranks: + groups.param_repli_group = _group + groups.param_repli_size = len(repli_ranks) + groups.param_repli_rank = dist.get_rank(group=_group) + logger.info(f'rank {global_rank} ' + f'replicate group {groups.param_repli_rank}/{dist.get_world_size(group=_group)}') + else: + groups.param_repli_group = None + groups.param_repli_size = 1 + groups.param_repli_rank = 0 + logger.info(f'rank {global_rank} replicate group 0/1') + + # assign shard group size as world size + assert groups.param_shard_size == len(ranks_of_shard_group[0]) + + if hierarchical_allgather: + # create hierarchy inter-node, intra-node groups + # n_span_nodes = config['shard_span'] + n_span_nodes = config['span_nodes'] + assert n_span_nodes > 1, "sharding spans on single node, no need for hierarchy allgather" + assert len(ranks_of_shard_group[0]) % n_span_nodes == 0 + + n_gpu_per_node = len(ranks_of_shard_group[0]) // n_span_nodes + intra_node_ranks_group = [] + inter_node_ranks_group = [] + for shard_group in ranks_of_shard_group: + _intra_node_ranks = [] + for i in range(0, len(shard_group), n_gpu_per_node): + _intra_node_ranks.append(shard_group[i:i + n_gpu_per_node]) + _inter_node_ranks = [] + for i in range(n_gpu_per_node): + _ranks = [_g[i] for _g in _intra_node_ranks] + _inter_node_ranks.append(_ranks) + + intra_node_ranks_group.append(_intra_node_ranks) + inter_node_ranks_group.append(_inter_node_ranks) + + _log_rank0(f"create for hierarchy all-gather groups: intra nodes {intra_node_ranks_group}") + _log_rank0(f"create for hierarchy all-gather groups: inter nodes {inter_node_ranks_group}") + + # create communicators + for shard_group in intra_node_ranks_group: + for intra_node_ranks in shard_group: + _group = dist.new_group(intra_node_ranks) + if global_rank in intra_node_ranks: + groups.param_intra_node_group = _group + _log_rank0(f'create group for intra node ranks {intra_node_ranks}') + + for shard_group in inter_node_ranks_group: + for inter_node_ranks in shard_group: + _group = dist.new_group(inter_node_ranks) + if global_rank in inter_node_ranks: + groups.param_inter_node_shard_group = _group + _log_rank0(f'create group for inter node ranks {inter_node_ranks}') + return groups + + +def _generate_mics_config(world_size, ndev_per_node, shard_size, pp_size=1): + """Generating the configuration for sharding This shard config generation assume + that the pipeline stages are partitioned in order, i.e., first ranks + hold the stage0, etc. + + Args: + + shard_size (int): zero3 data-parallel shard size, FIXME: + change the name later + + pp_size (int): pipeline parallel size, currently, only work with + pipeline parallelism + zero + + """ + assert world_size % pp_size == 0 + assert (world_size // pp_size) % shard_size == 0, \ + f"dp group size is not dividable by dp_shard_size, "\ + f" (world_size {world_size}, pp_size {pp_size}, dp_shard_size {shard_size})" + + config = {} + shard_groups = np.arange(world_size).reshape(-1, shard_size) + replicate_groups = [] + for i in range(shard_size): + same_shard_ranks = shard_groups[:, i].tolist() + n_ranks = len(same_shard_ranks) + replicate_size = n_ranks // pp_size + replicate_groups.extend([same_shard_ranks[j:j + replicate_size] for j in range(0, n_ranks, replicate_size)]) + + config['replicate_groups'] = replicate_groups + config['shard_groups'] = shard_groups.tolist() + config["span_nodes"] = len(shard_groups[0]) // ndev_per_node + return config + + +def _sizes_all_same(groups): + """all groups have same length""" + all_same = True + for g in groups: + if len(g) != len(groups[0]): + return False + return all_same diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ca35d7a7d1691c8ec4e6a7016cdb22b31ef459ac --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_config.py @@ -0,0 +1,109 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import Enum +from pathlib import Path +from pydantic import Field, model_validator +from typing import Optional + +from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int + + +class OffloadDeviceEnum(str, Enum): + """ Enum for valid offload devices """ + none = "none" + cpu = "cpu" + nvme = "nvme" + + +class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel): + """ Set options for parameter offload. Valid only with stage 3. """ + + device: OffloadDeviceEnum = "none" + """ + Device memory to offload model parameters. Supported options are `cpu` and + `nvme`. + """ + + nvme_path: Optional[Path] = None + """ Filesystem path for NVMe device for parameter offloading. """ + + buffer_count: int = Field(5, ge=0) + """ Number of buffers in buffer pool for parameter offloading to NVMe. """ + + buffer_size: int = Field(pp_int(1e8), ge=0) + """ Size of buffers in buffer pool for parameter offloading to NVMe. """ + + max_in_cpu: int = Field(pp_int(1e9), ge=0) + """ + Number of parameter elements to maintain in CPU memory when offloading to + NVMe is enabled. + """ + + pin_memory: bool = False + """ + Offload to page-locked CPU memory. This could boost throughput at the cost + of extra memory overhead. + """ + + +class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel): + """ Set options for optimizer offload. Valid with stage 1, 2, and 3. """ + + device: OffloadDeviceEnum = "none" + """ + Device memory to offload optimizer state. Supported options are `cpu` and + `nvme`. Optimizer computation is offload to CPU regardless of device option. + """ + + nvme_path: Optional[Path] = None + """ Filesystem path for NVMe device for optimizer state offloading. """ + + buffer_count: int = Field(4, ge=0) + """ + Number of buffers in buffer pool for optimizer state offloading to NVMe. + This should be at least the number of states maintained per parameter by + the optimizer. For example, Adam optimizer has 4 states (parameter, + gradient, momentum, and variance). + """ + + pin_memory: bool = False + """ + Offload to page-locked CPU memory. This could boost throughput at the cost + of extra memory overhead. + """ + + pipeline_read: bool = False + """ + 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" diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_states.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_states.py new file mode 100644 index 0000000000000000000000000000000000000000..d2c9b87a2bcaf302e09055277a5e3e311ca6ae60 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_states.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Set +import torch + +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum + + +def _make_offload_state_key(key): + return f"{key}_offload_buffer" + + +def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_key(state, key): + offload_buf_key = _make_offload_state_key(key) + if offload_buf_key not in state: + state[offload_buf_key] = torch.empty_like(state[key], device=device) + if pin_memory: + state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) + state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) + state[key].data = state[offload_buf_key] + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_key(state, "exp_avg_sq") + + +def reload_adam_states(optimizer, device, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_back_key(state, key): + state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking) + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_back_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_back_key(state, "exp_avg_sq") + + +def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]: + """Retrieve the devices of the specified state of the model. + + Args: + model (DeepSpeedEngine): The model whose device allocations are to be checked. + state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved. + + Returns: + Set[torch.device]: A set of devices of the specified state. + + """ + if state == OffloadStateTypeEnum.hp_params: + return set(model.optimizer.get_hp_param_device(p) for p in model.parameters()) + elif state == OffloadStateTypeEnum.lp_params: + return set(p.ds_tensor.device for p in model.parameters()) + elif state == OffloadStateTypeEnum.lp_grads: + return {model.optimizer.grad_partitions_flat_buffer.device} + elif state == OffloadStateTypeEnum.optim_states: + return set(model.optimizer.get_hp_param_device(p, "exp_avg") for p in model.parameters()) | \ + set(model.optimizer.get_hp_param_device(p, "exp_avg_sq") for p in model.parameters()) + elif state == OffloadStateTypeEnum.contiguous_grad_buffer: + return set(bucket.buffer.device for bucket in model.optimizer.ipg_buckets.values() + if bucket.buffer is not None) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py new file mode 100644 index 0000000000000000000000000000000000000000..000d0ebde1c7fbfe9e225ebd999e8c1d1a5e185c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py @@ -0,0 +1,589 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import torch +from collections import OrderedDict +from deepspeed.utils import z3_leaf_module, set_z3_leaf_module +from deepspeed.runtime.utils import see_memory_usage +from deepspeed.runtime.zero.utils import apply_to_tensors_only, is_zero_param +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.zero.partition_parameters import _init_external_params +from deepspeed.runtime.zero.partition_parameters import * +from deepspeed.runtime.zero.partitioned_param_coordinator import PartitionedParameterCoordinator, InflightParamRegistry, iter_params +from deepspeed.accelerator import get_accelerator +from deepspeed import utils + +FWD_MODULE_STACK = list() + + +#for each tensor in outputs run the forward_function and register backward_function as hook +def _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, outputs): + if type(outputs) is tuple: + touched_outputs = [] + for output in outputs: + touched_output = _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, + output) + touched_outputs.append(touched_output) + return tuple(touched_outputs) + elif type(outputs) is torch.Tensor: + forward_function(outputs) + if outputs.requires_grad: + outputs.register_hook(backward_function) + return outputs + else: + return outputs + + +class ZeROOrderedDict(OrderedDict): + + def __init__(self, parent_module, *args, **kwargs): + """A replacement for ``collections.OrderedDict`` to detect external ZeRO params. + + Args: + parent_module (``collections.OrderedDict``): the collection to replace + """ + + super().__init__(*args, **kwargs) + self._parent_module = parent_module + self._in_forward = False + + def __reduce__(self): + r0, _, *r2 = super().__reduce__() + return (r0, (self._parent_module, )) + tuple(r2) + + def __getitem__(self, key): + param = super().__getitem__(key) + + # Params can be registered as None (e.g., bias) + if param is None: + return param + + # TODO: only weaken this check during compilation + if hasattr(param, "ds_status") and param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if self._parent_module._parameters._in_forward: + register_external_parameter(FWD_MODULE_STACK[-1], param) + param.all_gather() + print_rank_0(f'Registering external parameter from getter {key} ds_id = {param.ds_id}', force=False) + + return param + + +def _inject_parameters(module, cls): + for module in module.modules(): + module._original_parameters = module._parameters + + if cls == ZeROOrderedDict: + new_param = cls(parent_module=module) + else: + new_param = cls() + + for key, param in module._parameters.items(): + new_param[key] = param + + module._parameters = new_param + + +class DeepSpeedZeRoOffload(object): + + def __init__( + self, + module, + timers, + ds_config, + overlap_comm=True, + prefetch_bucket_size=50000000, + max_reuse_distance=1000000000, + max_live_parameters=1000000000, + param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, + dp_process_group=None, + offload_param_config=None, + mpu=None, + zero_param_parallel_group=None, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + zero_module_granularity_threshold=0, + log_trace_cache_warnings=False, + ): + + see_memory_usage("DeepSpeedZeRoOffload initialize [begin]", force=True) + + print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) + + self.module = module + self.timers = timers + self.dtype = list(module.parameters())[0].dtype + self.dp_process_group = dp_process_group + self.offload_device = None + self.offload_param_pin_memory = False + self.zero_param_parallel_group = zero_param_parallel_group + self.zero_quantized_weights = zero_quantized_weights + self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights + self.log_trace_cache_warnings = log_trace_cache_warnings + + if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none: + self.offload_device = offload_param_config.device + self.offload_param_pin_memory = offload_param_config.pin_memory + + self._convert_to_zero_parameters(ds_config, module, mpu) + + for m in module.modules(): + _init_external_params(m) + + _inject_parameters(module, ZeROOrderedDict) + + self.param_numel_persistence_threshold = int(param_persistence_threshold) + self.model_persistence_threshold = int(model_persistence_threshold) + self.persistent_parameters = self.mark_persistent_parameters(self.param_numel_persistence_threshold, + self.model_persistence_threshold) + + self._prefetch_bucket_sz = int(prefetch_bucket_size) + self._max_reuse_distance_in_numel = int(max_reuse_distance) + self._max_available_parameters_in_numel = int(max_live_parameters) + self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream( + ) if overlap_comm else get_accelerator().default_stream() + + if not hasattr(module, "ds_inflight_param_registry"): + module.ds_inflight_param_registry = InflightParamRegistry() + self.__inflight_param_registry = module.ds_inflight_param_registry + + self.fast_sharding_for_leaf_module = False + + if zero_module_granularity_threshold > 0: + self.min_granularity_value = sys.maxsize + self.min_granularity_layer = None + self.granularity_info = set() + self.z3_leaf_layers = [] + self._set_z3_leaf_modules_by_threshold(module, zero_module_granularity_threshold) + self.fast_sharding_for_leaf_module = True + + self.param_coordinator = PartitionedParameterCoordinator( + prefetch_bucket_sz=self._prefetch_bucket_sz, + max_reuse_distance_in_numel=self._max_reuse_distance_in_numel, + max_available_parameters_in_numel=self._max_available_parameters_in_numel, + allgather_stream=self.__allgather_stream, + inflight_param_registry=self.__inflight_param_registry, + prefetch_nvme=self.offload_device == OffloadDeviceEnum.nvme, + timers=self.timers, + zero_quantized_weights=self.zero_quantized_weights, + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights, + fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module, + log_trace_cache_warnings=self.log_trace_cache_warnings, + ) + + self.forward_hooks = [] + self.backward_hooks = [] + + self.setup_zero_stage3_hooks() + print_rank_0( + f'Created module hooks: forward = {len(self.forward_hooks)}, backward = {len(self.backward_hooks)}', + force=False) + + see_memory_usage("DeepSpeedZeRoOffload initialize [end]", force=True) + + @instrument_w_nvtx + def partition_all_parameters(self): + """Partitioning Parameters that were not partitioned usually if parameters + of modules whose input parameters do not require grad computation do not + trigger post call and will therefore will remain unpartitioned""" + self.get_param_coordinator().release_and_reset_all(self.module) + for param in iter_params(self.module, recurse=True): + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f"{param.ds_summary()} expected to be released") + + def get_param_coordinator(self): + return self.param_coordinator + + def empty_partition_cache(self): + self.partition_all_parameters() + + def _convert_to_zero_parameters(self, ds_config, module, mpu): + non_zero_params = [p for p in module.parameters() if not is_zero_param(p)] + if non_zero_params: + zero_params = [p for p in module.parameters() if is_zero_param(p)] + if zero_params: + zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) + else: + group = None + # parallel_state_sp doesn't have get_data_parallel_group + if mpu and hasattr(mpu, "get_data_parallel_group"): + group = mpu.get_data_parallel_group() + + Init(module=module, + data_parallel_group=group, + dtype=self.dtype, + config_dict_or_path=ds_config, + remote_device=self.offload_device, + pin_memory=self.offload_param_pin_memory, + mpu=mpu, + zero_param_parallel_group=self.zero_param_parallel_group, + zero_quantized_weights=self.zero_quantized_weights, + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights) + + def destroy(self): + self._remove_module_hooks() + + def _remove_module_hooks(self): + num_forward_hooks = len(self.forward_hooks) + num_backward_hooks = len(self.backward_hooks) + + for hook in self.forward_hooks: + hook.remove() + + for hook in self.backward_hooks: + hook.remove() + + self.fwd_pre_hook.remove() + + print_rank_0(f'Deleted module hooks: forward = {num_forward_hooks}, backward = {num_backward_hooks}', + force=False) + + def setup_zero_stage3_hooks(self): + self.hierarchy = 0 + + #reset step if in inference mode + @instrument_w_nvtx + def _start_of_forward_hook(module, *args): + + self.get_param_coordinator().reset_step() + + self.fwd_pre_hook = self.module.register_forward_pre_hook(_start_of_forward_hook) + + #likely one of them should be enough but just to be safe + self._register_deepspeed_module(self.module) + + # Add top module to stack trace + global FWD_MODULE_STACK + FWD_MODULE_STACK.append(self.module) + + def mark_persistent_parameters(self, param_threshold, model_threshold): + persistent_params = [] + total_persistent_parameters = 0 + params_count = 0 + for name, param in self.module.named_parameters(recurse=True): + if param.ds_numel + total_persistent_parameters > model_threshold: + continue + + if param.ds_numel <= param_threshold: + params_count += 1 + param.ds_persist = True + persistent_params.append(param) + total_persistent_parameters += param.ds_numel + + print_rank_0( + f"Parameter Offload - Persistent parameters statistics: param_count = {params_count}, numel = {total_persistent_parameters}", + force=True) + + return persistent_params + + def _register_deepspeed_module(self, module, count=[0]): + my_count = count[0] + module.ds_id = my_count + + #print(f"{module.__class__} : {module.ds_id}") + + if z3_leaf_module(module): + for param in module.parameters(): + param.ds_z3_leaf_module = module + else: + for child in module.children(): + count[0] = count[0] + 1 + self._register_deepspeed_module(child, count=count) + + @torch.compiler.disable + def _pre_forward_module_hook(module, *args): + self.pre_sub_module_forward_function(module) + + @instrument_w_nvtx + def _post_forward_module_hook(module, input, output): + + global FWD_MODULE_STACK + FWD_MODULE_STACK.pop() + if output is None: + output = [] + elif not isinstance(output, (list, tuple)): + if torch.is_tensor(output): + output = [output] + else: + #print(f'got UNKNOWN type {type(output)}') + outputs = [] + output = output if isinstance(output, dict) else vars(output) + for name, val in output.items(): + if not name.startswith('__') and torch.is_tensor(val): + outputs.append(val) + output = outputs + + for item in filter(lambda item: is_zero_param(item) or hasattr(item, 'ds_param_alias'), output): + key = id(item) if hasattr(item, 'ds_id') else id(item.ds_param_alias) + actual_external_param = item if hasattr(item, 'ds_id') else item.ds_param_alias + + if not any(key in m._external_params for m in FWD_MODULE_STACK): + actual_external_param.is_external_param = True + module_to_register = FWD_MODULE_STACK[-1] + register_external_parameter(module_to_register, actual_external_param) + print_rank_0( + f'Registering dangling parameter for module {module_to_register.__class__.__name__}, ds_id = {actual_external_param.ds_id}.', + force=False) + + # It's possible that the parameter was already external to the completed module. If so, remove it the + # registration as it will be covered by the outer module instead. + if key in module._external_params: + print_rank_0( + f' Unregistering nested dangling parameter from module {module.__class__.__name__}, ds_id = {actual_external_param.ds_id}', + force=False) + unregister_external_parameter(module, actual_external_param) + + actual_external_param.all_gather() + + self.post_sub_module_forward_function(module) + + def _bwd_hook_unexpected_inputs_msg(value): + return f"A module has unknown inputs or outputs type ({type(value)}) and the tensors embedded in it cannot be detected. " \ + "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and " \ + "output tensors and therefore may not get triggered properly." + + def _pre_backward_module_hook(module, inputs, output): + + return apply_to_tensors_only(module.pre_bwd_fn.apply, + output, + warning_msg_fn=_bwd_hook_unexpected_inputs_msg) + + #This is an alternate to doing _post_backward_module_hook + #it uses tensor.register_hook instead of using torch.autograd.Function + def _alternate_post_backward_module_hook(module, inputs): + module.ds_grads_remaining = 0 + + #print(f"Before Forward {module.__class__.__name__}") + + def _run_after_backward_hook(*unused): + module.ds_grads_remaining = module.ds_grads_remaining - 1 + if module.ds_grads_remaining == 0: + #print(f"After backward {module.__class__.__name__}") + self.post_sub_module_backward_function(module) + + def _run_before_forward_function(input): + if input.requires_grad: + module.ds_grads_remaining += 1 + + return _apply_forward_and_backward_to_tensors_only(module, _run_before_forward_function, + _run_after_backward_hook, inputs) + + @torch.compiler.disable + def _post_backward_module_hook(module, inputs): + module.ds_grads_remaining = 0 + + return apply_to_tensors_only(module.post_bwd_fn.apply, + inputs, + warning_msg_fn=_bwd_hook_unexpected_inputs_msg) + + # Pre forward hook + self.forward_hooks.append(module.register_forward_pre_hook(_pre_forward_module_hook)) + + # Post forward hook + self.forward_hooks.append(module.register_forward_hook(_post_forward_module_hook)) + + # Pre backward hook + if not hasattr(module, "pre_bwd_fn"): + + @instrument_w_nvtx + def _run_before_backward_function(sub_module): + # some models (e.g. Albert) may run multiple forwards on the same layer in a loop + # before doing backwards, so each backward will need a pre-fetch - using reference + # counting to support this scenario + #print(f"COUNTER before: {sub_module.applied_pre_backward_ref_cnt}") + if sub_module.applied_pre_backward_ref_cnt > 0: + self.pre_sub_module_backward_function(sub_module) + sub_module.applied_pre_backward_ref_cnt -= 1 + #print(f"COUNTER after: {sub_module.applied_pre_backward_ref_cnt}") + + class PreBackwardFunctionForModule(torch.autograd.Function): + + @staticmethod + def forward(ctx, outputs): + # Capture `module` and _run_before_backward_function + ctx.module = module + ctx.pre_backward_function = _run_before_backward_function + if not hasattr(ctx.module, "applied_pre_backward_ref_cnt"): + ctx.module.applied_pre_backward_ref_cnt = 0 + ctx.module.applied_pre_backward_ref_cnt += 1 + outputs = outputs.detach() + return outputs + + @staticmethod + def backward(ctx, *args): + ctx.pre_backward_function(ctx.module) + return args + + module.pre_bwd_fn = PreBackwardFunctionForModule + + self.backward_hooks.append(module.register_forward_hook(_pre_backward_module_hook)) + + # post backward hook + if not hasattr(module, "post_bwd_fn"): + + @instrument_w_nvtx + def _run_after_backward_function(sub_module): + if sub_module.ds_grads_remaining == 0: + self.post_sub_module_backward_function(sub_module) + + class PostBackwardFunctionModule(torch.autograd.Function): + + @staticmethod + def forward(ctx, output): + ctx.module = module + if output.requires_grad: + #TODO SOME TIMES post backward does not seem to be triggered debug in detail + #Should only cause increase in memory not correctness issue + #if output.grad_fn.__class__.__name__ == 'ViewBackward': + # ctx.view=True + # print(f"Warning view tensor for input to module : {module.__class__.__name__}. Backward hooks may not trigger properly") + #assert len(module.parameters(recurse=False)), "The input tensor to the module is a view, and autograd Function or register_hook is not triggered with view tensors." + #if module.ds_grads_remaining == 0: + # print(f"Before Forward: {ctx.module.__class__.__name__}") + module.ds_grads_remaining += 1 + ctx.post_backward_function = _run_after_backward_function + output = output.detach() + return output + + @staticmethod + def backward(ctx, *args): + ctx.module.ds_grads_remaining = ctx.module.ds_grads_remaining - 1 + if ctx.module.ds_grads_remaining == 0: + ctx.post_backward_function(ctx.module) + return args + + module.post_bwd_fn = PostBackwardFunctionModule + + self.backward_hooks.append(module.register_forward_pre_hook(_post_backward_module_hook)) + + @torch.no_grad() + def pre_sub_module_forward_function(self, sub_module): + see_memory_usage(f"Before sub module function {sub_module.__class__.__name__}", force=False) + + global FWD_MODULE_STACK + FWD_MODULE_STACK.append(sub_module) + + param_coordinator = self.get_param_coordinator() + param_coordinator.trace_prologue(sub_module) + if param_coordinator.is_record_trace(): + param_coordinator.record_module(sub_module) + param_coordinator.fetch_sub_module(sub_module, forward=True) + + see_memory_usage(f"Before sub module function {sub_module.__class__.__name__} after fetch", force=False) + + @torch.no_grad() + def post_sub_module_forward_function(self, sub_module): + see_memory_usage( + f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} before release", + force=False) + + param_coordinator = self.get_param_coordinator() + param_coordinator.release_sub_module(sub_module, forward=True) + + see_memory_usage( + f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} after release", + force=False) + + @torch.no_grad() + def pre_sub_module_backward_function(self, sub_module): + # assert sub_module.training, "backward pass is invalid for module in evaluation mode" + param_coordinator = self.get_param_coordinator() + param_coordinator.trace_prologue(sub_module) + if param_coordinator.is_record_trace(): + param_coordinator.record_module(sub_module) + param_coordinator.fetch_sub_module(sub_module, forward=False) + + @torch.no_grad() + def post_sub_module_backward_function(self, sub_module): + # assert sub_module.training, "backward pass is invalid for module in evaluation mode" + see_memory_usage( + f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} before release", + force=False) + + self.get_param_coordinator().release_sub_module(sub_module, forward=False) + + see_memory_usage( + f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} after release", + force=False) + + def _set_z3_leaf_modules_by_threshold(self, module, zero_module_granularity_threshold): + + self._get_granularity_recursively(module) + print_rank_0(f"{'MODULE NAME'.ljust(30)}|{'GRANULARITY VALUE'.rjust(20)}", force=True) + for granularity in self.granularity_info: + print_rank_0(granularity, force=True) + + if self.min_granularity_value <= zero_module_granularity_threshold: + self._set_leaf_by_threshold_preorder(module, zero_module_granularity_threshold) + utils.logger.info( + f"z3_leaf_module was set by stage3_module_granularity_threshold:{zero_module_granularity_threshold}") + for layer in self.z3_leaf_layers: + print_rank_0(f"{layer.__class__.__name__}:{layer.ds_model_granularity}", force=True) + else: + utils.logger.warning( + f"The smallest module granularity is [{self.min_granularity_layer}:{self.min_granularity_value}]. "\ + f"To make stage3_module_granularity_threshold effective, you need to set stage3_module_granularity_threshold >= {self.min_granularity_value}. "\ + f"Current Value:{zero_module_granularity_threshold}" + ) + + def _get_granularity_recursively(self, module): + """This function is used to recursively obtain the granularity of each module.""" + + # avoid setting as leaf for particularly large models, even if the granularity is very small + # an oversized leaf module increases the number of live parameters, introducing memory overhead + Z3_MAX_LEAF_SIZE = 1e9 + + if not list(module.parameters()): + # skip Modules without parameters, such as GELU, etc. + module.ds_model_granularity = sys.maxsize + return 0, 0 + + num_layers = 0 + num_params = 0 + num_params += sum(p.ds_numel for p in module.parameters(recurse=False)) + if not any(module.children()): + # torch leaf module + module.ds_model_granularity = sys.maxsize + return 1, num_params + + for child in module.children(): + layers_in_child, params_in_child = self._get_granularity_recursively(child) + num_layers += layers_in_child + num_params += params_in_child + + if module.__class__.__name__ in torch.nn.modules.container.__all__: + # Do not set container modules like ModuleList as leaf modules + # as this will prevent hooks from being set on their children + # and they may do not invoke the forward method + module.ds_model_granularity = sys.maxsize + return num_layers, num_params + + num_layers += 1 + ds_model_granularity = (num_params // num_layers) if num_params <= Z3_MAX_LEAF_SIZE else sys.maxsize + module.ds_model_granularity = ds_model_granularity + # module.ds_model_num_layers = num_layers + # module.ds_model_num_params = num_params + if self.min_granularity_value > ds_model_granularity: + self.min_granularity_value = ds_model_granularity + self.min_granularity_layer = module.__class__.__name__ + self.granularity_info.add(f"{module.__class__.__name__.ljust(30)}|{str(ds_model_granularity).rjust(20)}") + + return num_layers, num_params + + def _set_leaf_by_threshold_preorder(self, module, granularity_treshhold): + '''Set modules as leaf modules based on the threshold, prioritizing parent nodes.''' + + num_params = sum(p.ds_numel for p in module.parameters()) + if num_params == 0: + # skip Modules without parameters, such as GELU, etc. + return + if module.ds_model_granularity <= granularity_treshhold: + set_z3_leaf_module(module, True) + self.z3_leaf_layers.append(module) + return + + for sub_module in module.children(): + self._set_leaf_by_threshold_preorder(sub_module, granularity_treshhold) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py new file mode 100644 index 0000000000000000000000000000000000000000..3482d30e35b9aa54c89102b5ab1a6b2439ba464b --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py @@ -0,0 +1,2329 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +import os +import types +from typing import Callable, Iterable +from enum import Enum +import functools +import itertools +from typing import List +from collections import defaultdict +import logging +import torch +from torch import Tensor +from deepspeed import comm as dist +from torch.nn import Module +from torch.nn import Parameter + +from .linear import zero3_linear_wrap + +from deepspeed.utils import groups +import deepspeed +from ..utils import see_memory_usage, get_only_unique_item +from deepspeed.runtime.zero.config import DeepSpeedZeroConfig +from deepspeed.runtime.zero.utils import assert_ints_same_as_other_ranks, is_zero_param +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.config_utils import get_config_default +from deepspeed.utils import instrument_w_nvtx, logger +from deepspeed.comm.comm import init_distributed +from deepspeed.utils.debug import (debug_param2name_id_shape, debug_param2name_id_shape_device, debug_module2name, + debug_param2name_id, debug_param2name_id_shape_status) +from deepspeed.accelerator import get_accelerator +from ..swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper, PartitionedParamStatus +from deepspeed.inference.quantization.utils import _quantize_param, WEIGHT_QUANTIZATION_LAYERS, wrap_quantized_functional, wrap_load_from_state_dict +from deepspeed.runtime.torch_autocast import sort_dtypes, get_autocast_dtype, has_autocast_dtype + +partitioned_param_data_shape = [0] +zero_init_context = 0 +top_level_context = None + + +class DeepSpeedTensorOverride(Enum): + dtype = 1 + device = 2 + + +DEFAULT_TENSOR_OVERRIDES = [DeepSpeedTensorOverride.dtype, DeepSpeedTensorOverride.device] + + +def get_allgather_dtype(param, param_ds_tensor): + autocast = has_autocast_dtype(param) + return get_autocast_dtype(param) if autocast else param_ds_tensor.dtype + + +class NoGatherHandle: + + def __init__(self, param: Parameter) -> None: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to be available") + + if hasattr(param.ds_tensor, "ds_quant_scale"): + param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to( + device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape) + else: + param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(), + non_blocking=True).view(param.ds_shape) + self.__param = param + + def wait(self, **kwargs) -> None: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().synchronize() + self.__param.ds_status = ZeroParamStatus.AVAILABLE + + +class NoGatherCoalescedHandle: + + def __init__(self, params: List[Parameter]) -> None: + self.__params = params + self.__complete = False + + for param in self.__params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + if hasattr(param.ds_tensor, "ds_quant_scale"): + param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to( + device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape) + else: + param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(), + non_blocking=True).view(param.ds_shape) + + @instrument_w_nvtx + def wait(self, **kwargs) -> None: + if self.__complete: + return + + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().synchronize() + for param in self.__params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + param.ds_status = ZeroParamStatus.AVAILABLE + + self.__complete = True + + +def _dist_allgather_fn(input_tensor: Tensor, output_tensor: Tensor, group=None): + return instrument_w_nvtx(dist.allgather_fn)(output_tensor, input_tensor, group=group, async_op=True) + + +def print_rank_0(message, debug=False, force=False): + rank = dist.get_rank() + if rank == 0 and (debug or force): + print(message) + # other variations + # - print for all ranks w/o interleaving + # printflock(f"[{rank}] {message}") + # - print to log file per rank + # log_rank_file(rank, message) + + +def debug_rank0(msg: str) -> None: + if dist.get_rank() == 0: + logger.debug(msg) + + +def _init_external_params(module): + if not hasattr(module, '_external_params'): + module._external_params = {} + + def external_parameters(self): + return self._external_params.items() + + def all_parameters(self): + return itertools.chain(self.named_parameters(self, recurse=False), external_parameters(self)) + + module.ds_external_parameters = types.MethodType(external_parameters, module) + module.all_parameters = types.MethodType(all_parameters, module) + + +def register_external_parameter(module, parameter): + """Instruct DeepSpeed to coordinate ``parameter``'s collection and partitioning in + the forward and backward passes of ``module``. + + This is used when a parameter is accessed outside of its owning module's + ``forward()``. DeepSpeed must know to collect it from its partitioned + state and when to release the memory. + + .. note:: + This is only applicable to training with ZeRO stage 3. + + Args: + module (``torch.nn.Module``): The module that requires ``parameter`` in its forward pass. + parameter (``torch.nn.Parameter``): The parameter to register. + + Raises: + RuntimeError: If ``parameter`` is not of type ``torch.nn.Parameter``. + + + Examples + ======== + + #. Register a weight that is used in another module's forward pass (line 6). + Parameter ``layer1.weight`` is used by ``layer2`` (line 11). + + .. code-block:: python + :linenos: + :emphasize-lines: 6,11 + + class ModuleZ3(torch.nn.Module): + def __init__(self, *args): + super().__init__(self, *args) + self.layer1 = SomeLayer() + self.layer2 = OtherLayer() + deepspeed.zero.register_external_parameter(self, self.layer1.weight) + + def forward(self, input): + x = self.layer1(input) + # self.layer1.weight is required by self.layer2.forward + y = self.layer2(x, self.layer1.weight) + return y + """ + if not isinstance(parameter, torch.nn.Parameter): + raise RuntimeError('Parameter is not a torch.nn.Parameter') + + if not hasattr(module, '_external_params'): + _init_external_params(module) + + key = id(parameter) + module._external_params[key] = parameter + + +def unregister_external_parameter(module, parameter): + """Reverses the effects of :meth:`register_external_parameter`. + + Args: + module (``torch.nn.Module``): The module to affect. + parameter (``torch.nn.Parameter``): The parameter to unregister. + + Raises: + RuntimeError: If ``parameter`` is not of type ``torch.nn.Parameter``. + RuntimeError: If ``parameter`` is not a registered external parameter of ``module``. + """ + if not isinstance(parameter, torch.nn.Parameter): + raise RuntimeError('Parameter is not a torch.nn.Parameter') + + if not hasattr(module, '_external_params') or id(parameter) not in module._external_params: + raise RuntimeError('Parameter is not a registered external parameter of module.') + + key = id(parameter) + del module._external_params[key] + + +class ZeroParamType(Enum): + + # same as regular pytorch parameters + NORMAL = 1 + + # parameters are partitioned across data parallel process + PARTITIONED = 2 + + # the parameter is held with a unique process rank + # and is not available on all other process + REMOTE = 3 + + +class ZeroParamStatus(Enum): + # parameters are fully present and ready for use on all processes + AVAILABLE = 1 + + # parameters are either partitioned or remote in some or all process + NOT_AVAILABLE = 2 + + # parameters are being gathered. + INFLIGHT = 3 + + +_orig_torch_tensor = torch.tensor +_orig_torch_empty = torch.empty +_orig_torch_zeros = torch.zeros +_orig_torch_ones = torch.ones +_orig_torch_full = torch.full +_orig_torch_arange = torch.arange +_orig_torch_eye = torch.eye +_orig_torch_randn = torch.randn + + +def zero_wrapper_for_fp_tensor_constructor(fn: Callable, target_fp_dtype: torch.dtype, + target_device: torch.device) -> Callable: + + def wrapped_fn(*args, **kwargs) -> Tensor: + if kwargs.get("device", None) is None and target_device is not None: + kwargs['device'] = target_device + tensor: Tensor = fn(*args, **kwargs) + if target_fp_dtype is not None and tensor.is_floating_point(): + tensor.data = tensor.data.to(target_fp_dtype) + + return tensor + + return wrapped_fn + + +def get_new_tensor_fn_for_dtype(target_fp_dtype: torch.dtype, target_device: torch.device) -> Callable: + + def new_tensor(cls, *args, **kwargs) -> Tensor: + if not args: + args = (0, ) + if target_device is None: + tensor = _orig_torch_empty(0).new_empty(*args, **kwargs) + else: + tensor = _orig_torch_empty(0, device=target_device).new_empty(*args, **kwargs) + + if tensor.is_floating_point() and target_fp_dtype is not None: + tensor = tensor.to(target_fp_dtype) + + return tensor + + return new_tensor + + +# https://stackoverflow.com/a/63851681/9201239 +def get_all_subclasses(cls, include_root=True): + subclass_list = [] + + def recurse(cl): + for subclass in cl.__subclasses__(): + subclass_list.append(subclass) + recurse(subclass) + + recurse(cls) + + ret = set(subclass_list) + if include_root: + ret.add(cls) + return ret + + +@instrument_w_nvtx +def free_param(param: Parameter) -> None: + """Free underlying storage of a parameter.""" + assert not param.ds_active_sub_modules, param.ds_summary() + if get_accelerator().on_accelerator(param.data): + # need to make sure that we don't free the parameter while it is still + # being used for computation + if not get_accelerator().is_synchronized_device(): + param.data.record_stream(get_accelerator().current_stream()) + # param.data doesn't store anything meaningful in partitioned state + param.data = torch.empty(0, dtype=param.dtype, device=param.device) + param.ds_status = ZeroParamStatus.NOT_AVAILABLE + + +reuse_buffers = False +temp_contiguous_tensor = None +empty_buffers = {} + + +# Inserts _post_init_method at the end of init method +# for all sub classes of torch.nn.Module +class InsertPostInitMethodToModuleSubClasses(object): + num_module_parameters = 0 + num_module_elements = 0 + + def __init__(self, enabled=True, mem_efficient_linear=True, ds_config=None, dtype=None): + self.mem_efficient_linear = mem_efficient_linear + self.enabled = enabled + self._set_dtype(ds_config, dtype) + assert self.dtype in [ + torch.half, torch.bfloat16, torch.float + ], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.bfloat16, torch.float]" + self.wrapped_cls = set() + self.skip_init_depth = 0 + + self.quantized_initialization = None + if ds_config is not None and ds_config.weight_quantization_config and ds_config.weight_quantization_config.quantized_initialization: + self.quantized_initialization = ds_config.weight_quantization_config.quantized_initialization + + def __enter__(self): + if not self.enabled: + return + + global zero_init_context + if zero_init_context == 0: + self.patch_init_and_builtins() + global top_level_context + top_level_context = self + + zero_init_context += 1 + + def __exit__(self, exc_type, exc_value, traceback): + if not self.enabled: + return + + global zero_init_context + zero_init_context -= 1 + + # Exiting the top level context + if zero_init_context == 0: + self.unpatch_init_and_builtins() + global top_level_context + top_level_context = None + + if dist.get_rank() == 0: + billion_elems = InsertPostInitMethodToModuleSubClasses.num_module_elements / 1e9 + num_params = InsertPostInitMethodToModuleSubClasses.num_module_parameters + logger.info( + f"finished initializing model - num_params = {num_params}, num_elems = {billion_elems:.2f}B") + + # Now that we cleaned up the metaclass injection, raise the exception. + if exc_type is not None: + return False + + # To be implemented by inheriting classes + def _post_init_method(self, module): + pass + + def _set_dtype(self, ds_config, dtype): + if ds_config is not None and dtype is None: + if ds_config.bfloat16_config.enabled and ds_config.float16_config.enabled: + raise RuntimeError("bfloat16 and fp16 cannot be enabled at once") + + if ds_config.bfloat16_config.enabled: + self.dtype = torch.bfloat16 + elif ds_config.float16_config.enabled: + self.dtype = torch.half + else: + self.dtype = torch.float + else: + self.dtype = dtype or torch.float16 if get_accelerator().is_fp16_supported( + ) else torch.bfloat16 if get_accelerator().is_bf16_supported else torch.float32 + + def patch_init_and_builtins(self): + + def apply_with_gather(orig_module_apply_fn: Callable) -> Callable: + """many models make use of child modules like Linear or Embedding which + perform their own weight initialization in their __init__ methods, + but will then have more weight initialization in a parent module's __init__ + method that modifies weights of child modules, which is typically done + using the Module.apply method. + + since the Init context manager partitions child modules immediately after + they are initialized, without modifying apply we would entirely skip + any initialization done by parent modules. + + to get around this issue, we wrap the function passed to Module.apply + so that the applied function is applied to child modules correctly. + """ + + def get_wrapped_fn_to_apply(fn_to_apply: Callable) -> Callable: + if hasattr(fn_to_apply, "wrapped"): + return fn_to_apply + + @functools.wraps(fn_to_apply) + def wrapped_fn_to_apply(module_to_apply_fn_to: Module) -> None: + """gathers parameters before calling apply function. afterwards + parameters are broadcasted to ensure consistency across all ranks + then re-partitioned. + + takes the following steps: + 1. allgathers parameters for the current module being worked on + 2. calls the original function + 3. broadcasts root rank's parameters to the other ranks + 4. re-partitions the parameters + """ + + # TODO Delay error checking for dangling partitioned parameters to post module init + # raise RuntimeError(f"not all parameters for {module_to_apply_fn_to.__class__.__name__}, " + # f"were zero params, is it possible that the parameters were " + # f"overwritten after they were initialized? " + # f"params: {[p for p in module_to_apply_fn_to.parameters(recurse=False)]} ") + + params_to_apply_fn_to: Iterable[Parameter] = list( + sorted([p for p in module_to_apply_fn_to.parameters(recurse=False) if is_zero_param(p)], + key=lambda p: p.ds_id)) + + for param in params_to_apply_fn_to: + param.all_gather() + + fn_to_apply(module_to_apply_fn_to) + + for param in params_to_apply_fn_to: + dist.broadcast(param.data, 0, group=param.ds_process_group) + + for param in params_to_apply_fn_to: + param.partition(has_been_updated=True) + + wrapped_fn_to_apply.wrapped = True + + return wrapped_fn_to_apply + + @functools.wraps(orig_module_apply_fn) + def wrapped_apply(module: Module, fn_to_apply: Callable) -> None: + orig_module_apply_fn(module, get_wrapped_fn_to_apply(fn_to_apply)) + + return wrapped_apply + + def hook_for_skip_init(module): + # this function is intended for handling the logic of torch.nn.utils.skip_init + # skip_init:module_cls(*args, **kwargs).to_empty(device=final_device), where kwargs['device']='meta' + # the function call occurs between module_cls(*args, **kwargs) and to_empty(device=final_device). + def partition_after_empty_init(f): + + @functools.wraps(f) + def wrapper(module, *args, **kwargs): + _module = f(module, *args, **kwargs) + # here is the post-hook for module.apply(empty_like...) + # after module.apply(empty_like...), the module has completed its empty init on real device + # since skip_init won't involve any computations or weight adjustments, we can directly utilize post_init + self._post_init_method(_module) + return _module + + return wrapper + + def post_wrapper_to_empty(f): + # append some wrapper restoration after to_empty() call + @functools.wraps(f) + def wrapper(*args, **kwargs): + res = f(*args, **kwargs) + # restore _apply hook + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _disable_class_apply(subclass) + # self restore + module.to_empty = f + return res + + return wrapper + + def _enable_class_apply(cls): + if '_apply' in cls.__dict__: + cls._old_apply_of_skip_init_hook = cls._apply + cls._apply = partition_after_empty_init(cls._apply) + + def _disable_class_apply(cls): + if hasattr(cls, '_old_apply_of_skip_init_hook'): + cls._apply = cls._old_apply_of_skip_init_hook + + # add hooks for to_empty: apply_(empty_like) + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _enable_class_apply(subclass) + + # add a restore hook when exiting skip_init + module.to_empty = post_wrapper_to_empty(module.to_empty) + + def partition_after(f): + + @functools.wraps(f) + def wrapper(module, *args, **kwargs): + + # important logic: We want to run post_init only after child's __init__ is + # completed, and do nothing after __init__ of any of its parents and grandparents in + # the inheritance ancestry. This way the partitioning will need to happen only once + # when the whole object is ready to be partitioned and not before. This is because + # often the child module will need to tweak the weights - for example running a + # custom weights init function. So if a parent created the weights param, the child + # won't need to gather it in order to tweak it + + print_rank_0(f'Before initializing {module.__class__.__name__}', force=False) + + is_child_module = False + if not hasattr(module, "_ds_child_entered"): + # child's __init__ was called, since parents all see the same object they can now skip post_init + is_child_module = True + setattr(module, "_ds_child_entered", True) + + init_on_meta = 'device' in kwargs and kwargs['device'] == 'meta' + if init_on_meta: + self.skip_init_depth += 1 + + f(module, *args, **kwargs) + if init_on_meta and self.skip_init_depth == 1: + # check and handle the logic of empty_init + hook_for_skip_init(module) + if is_child_module: + # child's __init__ is done, now we can run a single post_init on the child object + delattr(module, "_ds_child_entered") + + print_rank_0(f'Running post_init for {module.__class__.__name__}', force=False) + if self.skip_init_depth == 0: + self._post_init_method(module) + + print_rank_0(f'After initializing followed by post init for {module.__class__.__name__}', force=False) + if init_on_meta: + self.skip_init_depth -= 1 + + return wrapper + + def _enable_class(cls): + if '__init__' in cls.__dict__: + cls._old_init = cls.__init__ + cls.__init__ = partition_after(cls.__init__) + + def _init_subclass(cls, **kwargs): + if '__init__' in cls.__dict__: + cls._old_init = cls.__init__ + cls.__init__ = partition_after(cls.__init__) + + # Replace .__init__() for all existing subclasses of torch.nn.Module recursively + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _enable_class(subclass) + + # holding onto some methods so we can put them back the way they were in __exit__ + torch.nn.modules.module.Module._old_init_subclass = torch.nn.modules.module.Module.__init_subclass__ + torch.nn.modules.module.Module._old_apply = torch.nn.modules.module.Module.apply + torch.Tensor.__old_new__ = torch.Tensor.__new__ + + # Replace .__init__() for future subclasses of torch.nn.Module + torch.nn.modules.module.Module.__init_subclass__ = classmethod(_init_subclass) + if Init.override_module_apply: + torch.nn.modules.module.Module.apply = apply_with_gather(torch.nn.modules.module.Module._old_apply) + + if self.tensor_overrides: + self._add_tensor_creation_wrappers() + + if self.mem_efficient_linear: + print_rank_0( + "nn.functional.linear has been overridden with a more memory efficient version. This will persist unless manually reset.", + force=False) + if not hasattr(InsertPostInitMethodToModuleSubClasses, "linear_bk"): + InsertPostInitMethodToModuleSubClasses.linear_bk = torch.nn.functional.linear + torch.nn.functional.linear = zero3_linear_wrap + + if self.quantized_initialization: + print_rank_0("nn.functional.linear has been overridden with quantized linear version.", force=False) + torch.nn.functional.linear = wrap_quantized_functional(torch.nn.functional.linear) + torch.nn.functional.embedding = wrap_quantized_functional(torch.nn.functional.embedding) + for cls in WEIGHT_QUANTIZATION_LAYERS: + cls._load_from_state_dict = wrap_load_from_state_dict(cls._load_from_state_dict) + + logger.info("Enable Zero3 engine with INT4 quantization.") + + self.patched = True + + def unpatch_init_and_builtins(self): + if self.patched: + + def _disable_class(cls): + if hasattr(cls, '_old_init'): + cls.__init__ = cls._old_init + + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _disable_class(subclass) + + # putting methods back the way we found them + torch.nn.modules.module.Module.__init_subclass__ = torch.nn.modules.module.Module._old_init_subclass + if Init.override_module_apply: + torch.nn.modules.module.Module.apply = torch.nn.modules.module.Module._old_apply + + if self.tensor_overrides: + self._remove_tensor_creation_wrappers() + + self.patched = False + + def _add_tensor_creation_wrappers(self): + if DeepSpeedTensorOverride.dtype in self.tensor_overrides: + target_fp_dtype = self.dtype + else: + target_fp_dtype = None + if DeepSpeedTensorOverride.device in self.tensor_overrides: + target_device = self.local_device + else: + target_device = None + + torch.Tensor.__new__ = get_new_tensor_fn_for_dtype(target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.tensor = zero_wrapper_for_fp_tensor_constructor(_orig_torch_tensor, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.empty = zero_wrapper_for_fp_tensor_constructor(_orig_torch_empty, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.zeros = zero_wrapper_for_fp_tensor_constructor(_orig_torch_zeros, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.ones = zero_wrapper_for_fp_tensor_constructor(_orig_torch_ones, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.full = zero_wrapper_for_fp_tensor_constructor(_orig_torch_full, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.arange = zero_wrapper_for_fp_tensor_constructor(_orig_torch_arange, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.eye = zero_wrapper_for_fp_tensor_constructor(_orig_torch_eye, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + torch.randn = zero_wrapper_for_fp_tensor_constructor(_orig_torch_randn, + target_fp_dtype=target_fp_dtype, + target_device=target_device) + + def _remove_tensor_creation_wrappers(self): + torch.Tensor.__new__ = torch.Tensor.__old_new__ + torch.tensor = _orig_torch_tensor + torch.empty = _orig_torch_empty + torch.zeros = _orig_torch_zeros + torch.ones = _orig_torch_ones + torch.full = _orig_torch_full + torch.arange = _orig_torch_arange + torch.eye = _orig_torch_eye + torch.randn = _orig_torch_randn + + +def shutdown_init_context(): + """ + This function is used to initialize deepspeed engine inside the context of Init. + We need to remove the wrappers but keep the context. + """ + if top_level_context: + top_level_context.unpatch_init_and_builtins() + + +def restore_init_context(): + """ + This function is used to restore the wrappers after deepspeed engine is initialized. + """ + if top_level_context: + top_level_context.patch_init_and_builtins() + + +class AllGatherHandle: + + def __init__(self, handle, param: Parameter, quantization=None, original_dtype=None) -> None: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to be available") + + # Only one of original_dtype or quantization is provided + assert (original_dtype is None) != (quantization is None) + + self.__handle = handle + self.__param = param + self.__quantization = quantization + self.__original_dtype = original_dtype + + def wait(self, handle_dependency=True) -> None: + instrument_w_nvtx(self.__handle.wait)() + + if self.__original_dtype: + self.__param.data = self.__param.data.to(self.__original_dtype) + elif self.__quantization: + instrument_w_nvtx(self.__quantization.quant_handle.wait)() + self.__param.data = self.__quantization.backend.dequantize( + self.__quantization.quantized_param, self.__quantization.scale_buffer).to(self.__param.device) + self.__param.ds_status = ZeroParamStatus.AVAILABLE + + +class AllGatherCoalescedHandle: + + data_buffer = [] + + def __init__( + self, + allgather_handle, + params: List[Parameter], + partitions: List[Tensor], + world_size: int, + use_secondary_tensor=False, + quantization=None, + ) -> None: + self.allgather_handle = allgather_handle + self.params = params + self.partitions = partitions + self.world_size = world_size + self.use_secondary_tensor = use_secondary_tensor + self.complete = False + self.quantization = quantization + + for param in self.params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + + @instrument_w_nvtx + def wait(self, handle_dependency=True) -> None: + if self.complete: + return + + instrument_w_nvtx(self.allgather_handle.wait)() + + if self.quantization: + instrument_w_nvtx(self.quantization.quant_handle.wait)() + flat_tensor = self.quantization.backend.dequantize( + self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device) + + self.partitions: List[Parameter] = [] + for i in range(self.world_size): + self.partitions.append( + flat_tensor.narrow(0, self.quantization.partition_sz * i, self.quantization.partition_sz)) + + # split the single tensor out into individual tensors + param_offset = 0 + for param in self.params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + partitions: List[Tensor] = [] + ds_tensor_numel = param.ds_tensor.ds_numel + if self.use_secondary_tensor: + ds_tensor_numel *= param.ds_secondary_tensor_num_of_groups + for rank in range(self.world_size): + param_start = rank * ds_tensor_numel + if param_start < param.ds_numel: + part_to_copy = self.partitions[rank].narrow(0, param_offset, + min(param.ds_numel - param_start, ds_tensor_numel)) + partitions.append(part_to_copy) + # Note that dtypes of param and partitions can be different (currently for torch.autocast support) + param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape).to(param.ds_tensor.dtype) + param.ds_status = ZeroParamStatus.AVAILABLE + if not get_accelerator().is_synchronized_device() and handle_dependency: + for part_to_copy in partitions: + part_to_copy.record_stream(get_accelerator().current_stream()) + + param_offset += ds_tensor_numel + + self.complete = True + if not get_accelerator().is_synchronized_device() and not handle_dependency: + # if the device needs to handle dependencies and opts for explicit processing outside the function. + AllGatherCoalescedHandle.data_buffer.append(partitions) + + @staticmethod + def free_buffer(): + AllGatherCoalescedHandle.data_buffer = [] + + +class MultipleAllGatherHandles: + + def __init__(self, handles: List[AllGatherCoalescedHandle]): + self.handles = handles + + def wait(self, handle_dependency=True) -> None: + for handle in self.handles: + handle.wait(handle_dependency) + + +class AllReduceCoalescedHandle: + + def __init__(self, handle, params: List[Parameter]) -> None: + self.handle = handle + self.params = params + self.complete = False + + for param in self.params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + + @instrument_w_nvtx + def wait(self) -> None: + if self.complete: + return + + instrument_w_nvtx(self.handle.wait)() + + for param in self.params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + param.ds_status = ZeroParamStatus.AVAILABLE + + self.complete = True + + +class QuantizationInfo: + # a placeholder object to store all quant related vars used in handles + def __init__(self) -> None: + self.quantized_param = None + self.backend = None + self.quant_handle = None + self.scale_buffer = None + + +class CUDAQuantizer: + async_flag = True + target_group_size = 8000 # the optimal size is 4k, so we set the target to be below 8k + group_size_cache = dict() + quantizer_cuda_module = None + + def __init__(self) -> None: + if CUDAQuantizer.quantizer_cuda_module is None: + CUDAQuantizer.quantizer_cuda_module = deepspeed.ops.op_builder.QuantizerBuilder().load() + + def quantize(self, param, groups=None): + if groups is None: + try: + groups = self.group_size_cache[param.numel()] + except KeyError: + groups = math.ceil(param.numel() / self.target_group_size) + while groups < param.numel(): + if param.numel() % (8 * groups) == 0: + break + groups += 1 + while True: + if param.numel() % (8 * groups * 2) == 0 and param.numel( + ) / groups > self.target_group_size: #hard limit of 16k group_size + groups *= 2 + else: + break + assert ( + param.numel() % (8 * groups) == 0 + ), f"Qantized weight requires the number of weights be a multiple of 8. Yet {param.numel()} cannot be divided by 8*{groups}" + assert (param.numel() / groups < 16000), f"{param.numel()} / {groups} is larger than 16k" + assert param.numel( + ) > groups, f"Adaptive grouping algorithm cannot find a group size for input tensor of size {param.numel()}" + self.group_size_cache[param.numel()] = groups + return self.quantizer_cuda_module.quantize(param.to(get_accelerator().device_name()), groups, 8, + self.quantizer_cuda_module.Symmetric) + + def dequantize(self, quantized_param, scale): + return self.quantizer_cuda_module.dequantize(quantized_param, scale, scale.numel(), 8, + self.quantizer_cuda_module.Symmetric) + + +def _no_gather_coalesced(params: Iterable[Parameter]) -> AllGatherCoalescedHandle: + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f"expect param.ds_status == ZeroParamStatus.NOT_AVAILABLE, got{param.ds_summary()}") + param.ds_status = ZeroParamStatus.INFLIGHT + + params = sorted(params, key=lambda p: p.ds_id) + if len(params) == 1: + param, = params + return NoGatherHandle(param) + return NoGatherCoalescedHandle(params) + + +# Replaces all parameters in module with Scattered Parameters +class Init(InsertPostInitMethodToModuleSubClasses): + param_id = 0 + param_persistence_threshold = get_config_default(DeepSpeedZeroConfig, "param_persistence_threshold") + model_persistence_threshold = get_config_default(DeepSpeedZeroConfig, "model_persistence_threshold") + num_persisted_parameters = 0 + num_persisted_elements = 0 + apply_param_persistence = False + override_module_apply = get_config_default(DeepSpeedZeroConfig, "override_module_apply") + + def __init__(self, + module=None, + data_parallel_group=None, + mem_efficient_linear=True, + remote_device=None, + pin_memory=False, + config_dict_or_path=None, + config=None, + enabled=True, + dtype=None, + mpu=None, + zero_param_parallel_group=None, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + sequence_data_parallel_group=None, + param_swapper=None, + tensor_overrides=DEFAULT_TENSOR_OVERRIDES): + """A context to enable massive model construction for training with + ZeRO-3. Models are automatically partitioned (or, sharded) across the + system and converted to half precision. + + Args: + module (``torch.nn.Module``, optional): If provided, partition the model as + if it was constructed in the context. + data_parallel_group (``deepspeed.comm`` process group, optional): + The group of processes to partition among. Defaults to all processes. + Synonymous with sequence data parallel group for param partitioning + across both sequence and data parallel groups. + mem_efficient_linear (bool, optional): Replace + torch.nn.functional.linear with an implementation that allows + DeepSpeed to partition parameters. Defaults to ``True``. + remote_device (string, optional): The initial device to store model + weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU + memory. The model may still be moved to GPU based on the + offload settings for training. Defaults to param offload device if a config is + defined, otherwise GPU. + pin_memory (bool, optional): Potentially increase performance by + using pinned memory for model weights. ``remote_device`` must be + ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``. + config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration + for swapping fp16 params to NVMe and other things like ``dtype``. + config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. + enabled (bool, optional): If ``False``, this context has no + effect. Defaults to ``True``. + dtype (``dtype``, optional): Can be used to change the data type of the parameters. + Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` + mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}. + zero_param_parallel_group(``object``, optional): Parallel (comm) group for dual partitioning of ZeRO params. + zero_quantized_weights (bool, optional): If ``True``, turn on quantized weights in all gather weights. Default is ``False`` + zero_quantized_nontrainable_weights (bool, optional): If ``True``, nontrainable weights will be stored in quantized format. Default is ``False`` + param_swapper (``deepspeed.runtime.swap_tensor.partitioned_param_swapper.AsyncPartitionedParameterSwapper``, optional): [Experimental] Use existing parameter swapper. Defaults to ``None``. + This argument will be removed in the near future. + tensor_overrides ([`deepspeed.runtime.zero.DeepSpeedTensorOverride`], optional): Tensor attributes to override. Defaults to overriding dtype and device. + + This context accelerates model initialization and enables models that + are too large to allocate in their entirety in CPU memory. It has the + following effects: + + #. allocates tensors to either GPU or CPU memory or NVMe + #. converts floating point tensors to half precision + #. immediately partitions tensors among the group of data-parallel devices + #. (*optional*) replaces ``torch.nn.functional.linear`` with a more + memory-efficient implementation + + These modifications allow for models that exceed the size of local CPU/GPU + memory/NVMe, but fit within the total NVMe capacity (*i.e.*, aggregate CPU + or GPU memory or NVMe) across all nodes. Consider initializing a model with one + trillion parameters, whose weights occupy two terabytes (TB) in half + precision. The initial CPU allocation in full precision requires 4TB of + memory *per process*, and so a system with 8 GPUs per node would need 32TB of + CPU memory due to data-parallel redundancies. Instead, by immediately + partitioning tensors we remove the redundancies. The result is that + regardless of the number of GPUs, we still only require the original 4TB. This + allows for a linear increase in model size with the aggregate system memory. + For example, if a node has 1TB of memory and 8 GPUs, we could fit a trillion + parameter model with 4 nodes and 32 GPUs. + + Important: If the fp16 weights of the model can't fit onto a single GPU memory + this feature must be used. + + .. note:: + Initializes ``deepspeed.comm`` if it has not already been done so. + See :meth:`deepspeed.init_distributed` for more information. + + .. note:: + Only applicable to training with ZeRO-3. + + Examples + -------- + + #. Allocate a model and partition it among all processes: + + .. code-block:: python + + with deepspeed.zero.Init(): + model = MyLargeModel() + + + #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes: + + .. code-block:: python + + with deepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(), + remote_device="cpu", + pin_memory=True): + model = MyLargeModel() + + + #. Partition an already-allocated model in CPU memory: + + .. code-block:: python + + model = deepspeed.zero.Init(module=model) + """ + if config is not None: + config_dict_or_path = config + logger.warning('zero.Init: the `config` argument is deprecated. Please use `config_dict_or_path` instead.') + _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, + mpu) if config_dict_or_path is not None else None + if _ds_config is not None: + mem_efficient_linear = _ds_config.zero_config.memory_efficient_linear + + self.tensor_overrides = tensor_overrides + super().__init__(enabled=enabled, mem_efficient_linear=mem_efficient_linear, ds_config=_ds_config, dtype=dtype) + if not dist.is_initialized(): + init_distributed() + assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm" + + if data_parallel_group is None: + self.ds_process_group = dist.get_world_group() + else: + self.ds_process_group = data_parallel_group + + if sequence_data_parallel_group is not None: + logger.warning( + f"sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.") + if data_parallel_group is not None: + raise ValueError( + "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments." + ) + self.ds_process_group = sequence_data_parallel_group + + self.rank = dist.get_rank(group=self.ds_process_group) + self.dp_world_size = dist.get_world_size(group=self.ds_process_group) + + self.zero_param_process_group = zero_param_parallel_group + if _ds_config is not None and _ds_config.zero_config.zero_hpz_partition_size > 1 and self.zero_param_process_group is None: + groups._create_zero_param_parallel_group(_ds_config.zero_config.zero_hpz_partition_size) + self.zero_param_process_group = groups._get_zero_param_intra_parallel_group() + + self.num_ranks_in_param_group = self.dp_world_size + self.rank_in_group = self.rank + self.num_param_groups = 1 + + if self.zero_param_process_group is not None: + self.num_ranks_in_param_group = groups._get_zero_param_intra_parallel_group_world_size() + self.num_param_groups = int(self.dp_world_size / self.num_ranks_in_param_group) + self.rank_in_group = groups._get_zero_param_intra_parallel_rank_in_mygroup() + print_rank_0(f"hpZeRO group size: {self.num_ranks_in_param_group}", force=True) + + logger.debug( + "hpZeRO partition parameter my rank in world {} my rank in group {} ranks in my param partition group: {} " + .format(self.rank, self.rank_in_group, groups._get_zero_param_intra_parallel_group_ranks())) + + # Local device is the device where the parameters are consumed, must be default device. + # It is the device where parameters are fully instantiated using allgather + self.local_device = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])) + get_accelerator().set_device(self.local_device) + + self.quantized_weights = zero_quantized_weights + if _ds_config is not None and _ds_config.zero_config.zero_quantized_weights and not self.quantized_weights: + self.quantized_weights = _ds_config.zero_config.zero_quantized_weights + self.quantized_nontrainable_weights = zero_quantized_nontrainable_weights + if _ds_config is not None and _ds_config.zero_config.zero_quantized_nontrainable_weights and not self.quantized_nontrainable_weights: + self.quantized_nontrainable_weights = _ds_config.zero_config.zero_quantized_nontrainable_weights + + self.module = module + if (self.quantized_weights or self.quantized_nontrainable_weights): + self.quantizer_module = CUDAQuantizer() + print_rank_0(f'Using quantizer for weights: {self.quantizer_module.__class__.__name__}', force=True) + + if _ds_config is not None: + Init.override_module_apply = _ds_config.zero_config.override_module_apply + + if _ds_config.zero_config.offload_param is not None: + remote_device = _ds_config.zero_config.offload_param.device + pin_memory = _ds_config.zero_config.offload_param.pin_memory + + self._validate_remote_device(remote_device, _ds_config) + + # Remote device is the device where parameter partitions are stored + # It can be same as local_device or it could be CPU or NVMe. + self.remote_device = self.local_device if remote_device in [None, OffloadDeviceEnum.none] else remote_device + self.pin_memory = pin_memory if (self.remote_device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme + ]) else False + + # Enable fp16 param swapping to NVMe + if self.remote_device == OffloadDeviceEnum.nvme: + self.param_swapper = param_swapper or AsyncPartitionedParameterSwapper(_ds_config, self.dtype) + else: + self.param_swapper = None + + # If we are provided an already-allocated module to prepare. + if module is not None: + assert isinstance(module, torch.nn.Module) + self._convert_to_zero_parameters(module.parameters(recurse=True)) + + self.use_all_gather_into_tensor = dist.has_all_gather_into_tensor() + if not self.use_all_gather_into_tensor: + logger.info(f"all_gather_into_tensor API is not available in torch {torch.__version__}") + + self.use_all_reduce_for_fetch_params = get_config_default(DeepSpeedZeroConfig, + "use_all_reduce_for_fetch_params") + if _ds_config is not None: + self.use_all_reduce_for_fetch_params = _ds_config.zero_config.use_all_reduce_for_fetch_params + + def _update_persist_config(self, ds_config): + Init.apply_param_persistence = True + Init.param_persistence_threshold = ds_config.zero_config.param_persistence_threshold + Init.model_persistence_threshold = ds_config.zero_config.model_persistence_threshold // self.num_partitions + + def _zero_init_param(self, param): + self._convert_to_deepspeed_param(param) + if dist.get_world_group() == self.get_dp_process_group(): + dist.broadcast(param.data, 0, self.get_dp_process_group()) + else: + dist.broadcast(param.data, dist.get_global_rank(self.get_dp_process_group(), 0), + self.get_dp_process_group()) + param.partition() + + def _convert_to_zero_parameters(self, param_list): + for param in param_list: + if is_zero_param(param): + continue + + param.data = param.data.to(self.local_device) + self._zero_init_param(param) + + def _validate_remote_device(self, remote_device, ds_config): + if ds_config is not None: + if remote_device in [None, OffloadDeviceEnum.cpu]: + if ds_config.zero_config.offload_param is not None: + offload_param_device = ds_config.zero_config.offload_param.device + assert offload_param_device != OffloadDeviceEnum.nvme, \ + f"'device' in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}." + + if remote_device == OffloadDeviceEnum.nvme: + assert ds_config.zero_config.offload_param is not None, \ + f'"offload_param" must be defined in DeepSpeed Config if remote device is {OffloadDeviceEnum.nvme}.' + + assert ds_config.zero_config.offload_param.nvme_path is not None, \ + f'"nvme_path" in DeepSpeed Config cannot be None if remote device is {OffloadDeviceEnum.nvme}' + + def _post_init_method(self, module): + #see_memory_usage(f"Before converting params in {module.__class__.__name__}", force=False) + print_rank_0(f'Converting Params in {module.__class__.__name__}', force=False) + see_memory_usage(f"Before converting and partitioning params in {module.__class__.__name__}", force=False) + + for name, param in module.named_parameters(recurse=False): + print_rank_0(f'Analyzing param {name} in {module.__class__.__name__}', force=False) + InsertPostInitMethodToModuleSubClasses.num_module_parameters += 1 + InsertPostInitMethodToModuleSubClasses.num_module_elements += param.numel() + if not is_zero_param(param): + if not get_accelerator().on_accelerator(param): + param.data = param.data.to(self.local_device) + + if name == 'weight' and self.quantized_initialization and type(module) in WEIGHT_QUANTIZATION_LAYERS: + _quantize_param(param, self.quantized_initialization) + + self._zero_init_param(param) + print_rank_0( + f"Partitioning param {debug_param2name_id_shape(param)} module={debug_module2name(module)}") + + see_memory_usage( + f"Param count {InsertPostInitMethodToModuleSubClasses.num_module_elements}. After converting and partitioning params in {module.__class__.__name__}", + force=False) + + def _convert_to_deepspeed_param(self, param): + + # Partitioned, Normal, Remote + param.ds_param_type = ZeroParamType.PARTITIONED + + # Replicated vs Partitioned vs Inflight + param.ds_status = ZeroParamStatus.AVAILABLE + + # Stores the shape of the original tensor + param.ds_shape = param.shape + + # Stores the number of elements in the original parameter without padding + param.ds_numel = param.numel() + + # Stores the partitioned copy of the tensor + param.ds_tensor = None + + # Keeps track of how many active sub-modules need this param at any given point in time + param.ds_active_sub_modules = set() + + # If this flag is true, then the parameters are replicated throughput training + # And only partitioned before the step + if Init.apply_param_persistence and param.ds_numel <= Init.param_persistence_threshold and Init.num_persisted_elements + param.ds_numel <= Init.model_persistence_threshold: + param.ds_persist = True + Init.num_persisted_parameters += 1 + Init.num_persisted_elements += param.ds_numel + else: + param.ds_persist = False + + param.is_external_param = False + + # The group that the parameter is scattered across. + param.ds_process_group = self.ds_process_group + + # Stores the secondary partitioned copy of the tensor + param.ds_secondary_tensor = None + + #Process group for secondary partition all (group) gather + param.ds_zero_param_process_group = self.zero_param_process_group + param.ds_secondary_tensor_group_size = self.num_ranks_in_param_group + param.ds_secondary_tensor_num_of_groups = self.num_param_groups + + # This is set to the Async Param swapper if remote device is nvme + # else this is set to None + param.nvme_swapper = self.param_swapper + + # DeepSpeed Param ID + param.ds_id = Init.param_id + Init.param_id += 1 + + def all_gather(param_list=None, async_op=False, hierarchy=0): + cls = param + if param_list is None: + param_list = [cls] + return self._all_gather(param_list, async_op=async_op, hierarchy=hierarchy) + + def _all_gather_dtype(params, world_size, rank_in_group, ds_process_group, allgather_dtype): + # make sure all params have the same dtype + dtype = params[0].dtype # we assume len(params) > 0 + assert all(p.dtype == dtype for p in params), "all params must have the same dtype" + + partition_sz = sum(p.ds_tensor.ds_numel for p in params) + + use_secondary_tensor = params[0].ds_secondary_tensor is not None + + if use_secondary_tensor: + partition_sz = sum(p.ds_tensor.ds_numel * p.ds_secondary_tensor_num_of_groups for p in params) + + flat_tensor = torch.empty(partition_sz * world_size, + dtype=allgather_dtype, + device=get_accelerator().current_device_name(), + requires_grad=False) + + partitions: List[Parameter] = [] + for i in range(world_size): + partitions.append(flat_tensor.narrow(0, partition_sz * i, partition_sz)) + + if use_secondary_tensor: + instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype) + for p in params + ], + out=partitions[rank_in_group]) + else: + instrument_w_nvtx(torch.cat)( + [p.ds_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype) for p in params], + out=partitions[rank_in_group]) + handle = _dist_allgather_fn(partitions[rank_in_group], flat_tensor, ds_process_group) + #Fix get_partition_dp_group(params[0])) + + return AllGatherCoalescedHandle( + allgather_handle=handle, + params=params, + partitions=partitions, + world_size=world_size, + use_secondary_tensor=use_secondary_tensor, + ) + + @instrument_w_nvtx + def all_gather_coalesced(params: Iterable[Parameter], + safe_mode: bool = False, + quantize: bool = False) -> AllGatherCoalescedHandle: + + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(params) + + if self.num_partitions == 1: + return _no_gather_coalesced(params) + + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(param.ds_summary()) + param.ds_status = ZeroParamStatus.INFLIGHT + + #use appropriate all gather process group + ds_process_group = self.ds_process_group + rank_in_group = self.rank + world_size = self.dp_world_size + use_secondary_tensor = params[0].ds_secondary_tensor is not None + if self.zero_param_process_group and use_secondary_tensor: + ds_process_group = self.zero_param_process_group #intragroup + rank_in_group = self.rank_in_group + world_size = self.num_ranks_in_param_group + + #pprint(dir(ds_process_group)) + # ensure that each rank has params in same order. the allgather + # is done by flattening the parameter list into a single tensor that + # can be allgathered in a single call - this means that if each rank + # gives a list of the same parameters in a different order we will + # silently get incorrect parameter values, and have very difficult + # to debug correctness issues. + params = sorted(params, key=lambda p: p.ds_id) + + if logger.isEnabledFor(logging.DEBUG): + debug_rank0(f"-allgather_coalesced: {[p.ds_id for p in params]}") + + if safe_mode: + # ensure that same list (with same ordering) of parameters are + # being allgathered across all ranks, otherwise could mix + # data between tensors. + assert_ints_same_as_other_ranks([p.ds_id for p in params]) + # ensure that tensors from each rank agree on the same ds_numel + # otherwise could mix data between tensors. + assert_ints_same_as_other_ranks([p.ds_tensor.ds_numel for p in params]) + + if len(params) == 1: + # have an opportunity to avoid some intermediate memory allocations + param = params[0] + buffer_size = math.ceil(param.ds_numel / world_size) * world_size + if use_secondary_tensor: + buffer_size = param.ds_secondary_tensor.shape[0] * world_size #make sure out is appropriately sized + + param_ds_tensor = param.ds_secondary_tensor if use_secondary_tensor else param.ds_tensor + + original_dtype = param_ds_tensor.dtype + if quantize: + allgather_dtype = torch.int8 + else: + allgather_dtype = get_allgather_dtype(param, param_ds_tensor) + + param_buffer = torch.empty( + buffer_size, + dtype=allgather_dtype, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + if not quantize: + handles = _dist_allgather_fn( + param_ds_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype), + param_buffer, + ds_process_group, + ) + param.data = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device) + return AllGatherHandle(handles, param, original_dtype=original_dtype) + else: + if hasattr(param_ds_tensor, "ds_quant_scale"): + scales = param_ds_tensor.ds_quant_scale + quantized_param = param_ds_tensor.data + else: + quantized_param, scales = self.quantizer_module.quantize(param_ds_tensor) + handle = _dist_allgather_fn(quantized_param.to(get_accelerator().current_device_name()), + param_buffer, ds_process_group) + + quant_scale_buffer = torch.empty( + scales.numel() * world_size, + dtype=scales.dtype, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + quant_handle = _dist_allgather_fn(scales.to(get_accelerator().current_device_name()), + quant_scale_buffer, ds_process_group) + quant_info = QuantizationInfo() + quant_info.quantized_param = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to( + param.device) + quant_info.backend = self.quantizer_module + quant_info.quant_handle = quant_handle + quant_info.scale_buffer = quant_scale_buffer + return AllGatherHandle(handle, param, quantization=quant_info) + + else: + if self.use_all_reduce_for_fetch_params and not quantize and not use_secondary_tensor: + + # Use all_reduce instead of all_gather to fetch the module params + flat_buffer_size = sum(p.ds_numel_aligned for p in params) + flat_tensor = torch.zeros(flat_buffer_size, + dtype=get_only_unique_item(p.ds_tensor.dtype for p in params), + device=get_accelerator().current_device_name(), + requires_grad=False) + start_param = 0 + for param in params: + param.data = flat_tensor.narrow(0, start_param, param.ds_numel).view(param.ds_shape) + start = start_param + param.ds_tensor.ds_numel * self.get_partition_rank() + flat_tensor.narrow(0, start, param.ds_tensor.ds_numel).copy_(param.ds_tensor) + + start_param += param.ds_numel + + handle = dist.all_reduce(flat_tensor, group=ds_process_group, async_op=True) + + return AllReduceCoalescedHandle(handle=handle, params=params) + else: + if not quantize: + dtype_params = defaultdict(list) + for p in params: + allgather_dtype = get_allgather_dtype(p, p.ds_tensor) + dtype_params[allgather_dtype].append(p) + handles = [] + for dtype in sort_dtypes(dtype_params.keys()): + handles.append( + _all_gather_dtype(params, world_size, rank_in_group, ds_process_group, + allgather_dtype)) + + return MultipleAllGatherHandles(handles) + + else: + partition_sz = sum(p.ds_tensor.ds_numel for p in params) + + if use_secondary_tensor: + partition_sz = sum(p.ds_tensor.ds_numel * p.ds_secondary_tensor_num_of_groups + for p in params) + + flat_tensor = torch.empty(partition_sz * world_size, + dtype=torch.int8, + device=get_accelerator().current_device_name(), + requires_grad=False) + + if use_secondary_tensor: + if hasattr(params[0].ds_secondary_tensor, "ds_quant_scale"): + quantized_param = instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.data.to(get_accelerator().current_device_name()) + for p in params + ]) + scales = instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.ds_quant_scale.to(get_accelerator().current_device_name()) + for p in params + ]) + else: + quantized_param, scales = self.quantizer_module.quantize( + instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.to(get_accelerator().current_device_name()) + for p in params + ])) + else: + if hasattr(params[0].ds_tensor, "ds_quant_scale"): + quantized_param = instrument_w_nvtx(torch.cat)( + [p.ds_tensor.data.to(get_accelerator().current_device_name()) for p in params]) + scales = instrument_w_nvtx(torch.cat)([ + p.ds_tensor.ds_quant_scale.to(get_accelerator().current_device_name()) + for p in params + ]) + else: + quantized_param, scales = self.quantizer_module.quantize( + instrument_w_nvtx(torch.cat)( + [p.ds_tensor.to(get_accelerator().current_device_name()) for p in params])) + quant_scale_buffer = torch.empty( + scales.numel() * world_size, + dtype=torch.float32, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + handle = _dist_allgather_fn(quantized_param, flat_tensor, ds_process_group) + quant_handle = _dist_allgather_fn(scales, quant_scale_buffer, ds_process_group) + quant_info = QuantizationInfo() + quant_info.quantized_param = flat_tensor + quant_info.backend = self.quantizer_module + quant_info.quant_handle = quant_handle + quant_info.scale_buffer = quant_scale_buffer + quant_info.partition_sz = partition_sz + quant_info.world_size = world_size + return AllGatherCoalescedHandle( + allgather_handle=handle, + params=params, + partitions=None, + world_size=world_size, + use_secondary_tensor=use_secondary_tensor, + quantization=quant_info, + ) + + def partition(param_list=None, hierarchy=0, has_been_updated=False, free_data=True): + cls = param + print_rank_0(f"{'--'*hierarchy}----Partitioning param {debug_param2name_id_shape_device(cls)}", + force=False) + if param_list is None: + param_list = [cls] + self._partition(param_list, has_been_updated=has_been_updated, free_data=True) + + def reduce_gradients_at_owner(param_list=None, hierarchy=0): + cls = param + if param_list is None: + param_list = [cls] + print_rank_0( + f"{'--'*hierarchy}----Reducing Gradients for param with ids {[param.ds_id for param in param_list]} to owner" + ) + self._reduce_scatter_gradients(param_list) + + def partition_gradients(param_list=None, partition_buffers=None, hierarchy=0, accumulate=False): + cls = param + print_rank_0( + f"{'--'*hierarchy}----Partitioning param gradient with id {debug_param2name_id_shape_device(cls)}") + if param_list is None: + param_list = [cls] + if isinstance(partition_buffers, torch.Tensor): + partition_buffers = [partition_buffers] + + self._partition_gradients(param_list, partition_buffers=partition_buffers, accumulate=accumulate) + + def aligned_size(): + return self._aligned_size(param) + + def padding_size(): + return self._padding_size(param) + + def partition_numel(): + return self._partition_numel(param) + + def item_override(): + param.all_gather() + return param._orig_item() + + def ds_summary(slf: torch.Tensor, use_debug_name: bool = False) -> dict: + return { + "id": debug_param2name_id(slf) if use_debug_name else slf.ds_id, + "status": slf.ds_status.name, + "numel": slf.numel(), + "ds_numel": slf.ds_numel, + "shape": tuple(slf.shape), + "ds_shape": tuple(slf.ds_shape), + "requires_grad": slf.requires_grad, + "grad_shape": tuple(slf.grad.shape) if slf.grad is not None else None, + "persist": slf.ds_persist, + "active_sub_modules": slf.ds_active_sub_modules, + "ds_tensor.shape": slf.ds_tensor.shape if slf.ds_tensor is not None else None + } + + def convert_to_zero_parameters(param_list): + self._convert_to_zero_parameters(param_list) + + def allgather_before(func: Callable) -> Callable: + + def wrapped(*args, **kwargs): + param.all_gather() + return func(*args, **kwargs) + + return wrapped + + # Collectives for gathering and partitioning parameters + param.all_gather = all_gather + param.all_gather_coalesced = all_gather_coalesced + param.partition = partition + + # Collective for averaging gradients + param.reduce_gradients_at_owner = reduce_gradients_at_owner + param.partition_gradients = partition_gradients + + # Partitioning size utilities + param.aligned_size = aligned_size + param.padding_size = padding_size + param.partition_numel = partition_numel + param.ds_summary = types.MethodType(ds_summary, param) + + param.item = allgather_before(param.item) + + param.convert_to_zero_parameters = convert_to_zero_parameters + + def _aligned_size(self, param): + return param.ds_numel + self._padding_size(param) + + def _padding_size(self, param): + remainder = param.ds_numel % self.num_partitions + return (self.num_partitions - remainder) if remainder else 0 + + def _partition_numel(self, param): + return param.ds_tensor.ds_numel + + def _ensure_availability_of_partitioned_params(self, params): + swap_in_list = [] + swap_in_flight = [] + for param in params: + if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE: + assert param.ds_tensor.final_location == OffloadDeviceEnum.nvme and param.ds_status == ZeroParamStatus.NOT_AVAILABLE + swap_in_list.append(param) + if param.ds_tensor.status == PartitionedParamStatus.INFLIGHT: + assert param.ds_tensor.final_location == OffloadDeviceEnum.nvme and param.ds_status == ZeroParamStatus.NOT_AVAILABLE + swap_in_flight.append(param) + if len(swap_in_list) > 0: + swap_in_list[0].nvme_swapper.swap_in(swap_in_list, async_op=False) + elif len(swap_in_flight) > 0: + swap_in_flight[0].nvme_swapper.synchronize_reads() + + @instrument_w_nvtx + def _all_gather(self, param_list, async_op=False, hierarchy=None): + + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(param_list) + + handles = [] + all_gather_list = [] + for param in param_list: + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if async_op: + handle = self._allgather_param(param, async_op=async_op, hierarchy=hierarchy) + param.ds_status = ZeroParamStatus.INFLIGHT # if async_op else ZeroParamStatus.AVAILABLE + handles.append(handle) + else: + all_gather_list.append(param) + # note: param_list may contain params that are already in flight / aviailable. So we need to use all_gather_list + if not async_op: + if len(all_gather_list) == 1: + ret_value = self._allgather_params(all_gather_list, hierarchy=hierarchy) + else: + all_gather_quantize_list = [] + all_gather_nonquantize_list = [] + for param in all_gather_list: + if hasattr(param.ds_tensor, + "ds_quant_scale") or (hasattr(param, "ds_secondary_tensor") + and hasattr(param.ds_secondary_tensor, "ds_quant_scale")): + all_gather_quantize_list.append(param) + else: + all_gather_nonquantize_list.append(param) + # _allgather_params_coalesced always return None + self._allgather_params_coalesced(all_gather_nonquantize_list, hierarchy, quantize=False) + self._allgather_params_coalesced(all_gather_quantize_list, hierarchy, quantize=True) + for param in all_gather_list: + param.ds_status = ZeroParamStatus.AVAILABLE + return None + + return handles + + def _partition(self, param_list, force=False, has_been_updated=False, free_data=True): + for param in param_list: + print_rank_0(f"Before Partitioning Param {param.ds_id}", force=False) + if self.zero_param_process_group is not None: + self._partition_param_sec(param) + self._partition_param(param, has_been_updated=has_been_updated, free_data=True) + + param.ds_status = ZeroParamStatus.NOT_AVAILABLE + # if param.ds_tensor is not None: + # assert id(param.data) == id(param.ds_tensor.data), \ + # "After the parameters are initially partitioned, make sure we are not recreating the partition." + #print_rank_0(f"After Partitioning Param {param.ds_id} {param.ds_tensor.size()} {param.ds_tensor}",force=False) + @instrument_w_nvtx + def _partition_param(self, param, buffer=None, has_been_updated=False, free_data=True): + assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" + global reuse_buffers + print_rank_0(f"Param id {param.ds_id} status is {param.ds_status}", force=False) + if param.ds_status is ZeroParamStatus.AVAILABLE: + print_rank_0(f"Partitioning param id {param.ds_id} reuse buffers {reuse_buffers}", force=False) + # if reuse_buffers and False: + # numel = buffer.numel() + # buffer = param.data.view(-1) + # print_rank_0( + # "Returning buffer for param {param.ds_id} with numel {param.ds_numel} to empty buffers", + # force=False) + # if numel in empty_buffers: + # empty_buffers[numel].append(buffer) + + # if deepspeed.comm.get_rank(): + # print(f"Releasing {param.data.numel()}") + + if param.ds_tensor is not None and not has_been_updated: ##param already partitioned + + #print_rank_0(f"Param {param.ds_id} pri {param.ds_tensor.size()} loc? {param.ds_tensor.final_location}", force=True) + #param.data = param.ds_tensor.data + + see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) + # param.data does not store anything meaningful in partitioned state + if free_data: + free_param(param) + see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) + + if param.ds_tensor.final_location == OffloadDeviceEnum.nvme: + print_rank_0(f"Param {param.ds_id} partition released since it exists in nvme", force=False) + param.nvme_swapper.remove_partition_and_release_buffers([param]) + print_rank_0( + f"after swap Param {param.ds_id} {param.ds_tensor.shape} partition released since it exists in nvme", + force=False) + + return + + tensor_size = self._aligned_size(param) + partition_size = tensor_size // self.num_partitions + if param.ds_tensor is None: + final_location = None + if self.remote_device == OffloadDeviceEnum.nvme and self.param_swapper.swappable_tensor( + numel=partition_size): + final_location = OffloadDeviceEnum.nvme + buffer = self.param_swapper.get_buffer(param, partition_size) + partitioned_tensor = torch.empty(0, dtype=param.dtype, device=buffer.device) + partitioned_tensor.data = buffer.data + print_rank_0(f"ID {param.ds_id} Initializing partition for the first time for nvme offload.") + + else: + if param.ds_persist: + device = self.local_device + elif self.remote_device == OffloadDeviceEnum.nvme: + device = OffloadDeviceEnum.cpu + else: + device = self.remote_device + + partitioned_tensor = torch.empty(partition_size, dtype=param.dtype, device=device) + # quantize the tensor if it's not trainable + if not param.requires_grad and self.quantized_nontrainable_weights: + partitioned_tensor, partitioned_tensor.ds_quant_scale = self.quantizer_module.quantize( + partitioned_tensor) + + if device == OffloadDeviceEnum.cpu and self.pin_memory: + partitioned_tensor = get_accelerator().pin_memory(partitioned_tensor) + + partitioned_tensor.requires_grad = False + param.ds_tensor = partitioned_tensor + param.ds_tensor.ds_numel = partition_size + param.ds_tensor.status = PartitionedParamStatus.AVAILABLE + param.ds_tensor.final_location = final_location + param.ds_numel_aligned = tensor_size + + start = partition_size * self.get_partition_rank() + end = start + partition_size + + one_dim_param = param.contiguous().view(-1) + + if start < param.ds_numel and end <= param.ds_numel: + src_tensor = one_dim_param.narrow(0, start, partition_size) + + with torch.no_grad(): + # make sure param.ds_tensor requires_grad always be false, + # otherwise, torch tracer will complain. + param.ds_tensor.copy_(src_tensor) + + #partitioned_tensor = src_tensor.clone().detach().to(self.remote_device) + + else: + # partitioned_tensor = torch.zeros(partition_size, + # dtype=param.dtype, + # device=self.remote_device ) + + if start < param.ds_numel: + elems_to_copy = param.ds_numel - start + with torch.no_grad(): + # make sure param.ds_tensor requires_grad always be false, + # otherwise, torch tracer will complain. + param.ds_tensor.narrow(0, 0, + elems_to_copy).copy_(one_dim_param.narrow(0, start, elems_to_copy)) + + #print(f"Remote device {self.remote_device}") + + #param.ds_tensor = partitioned_tensor + + #param.data = param.ds_tensor.data + + # param.data does not store anything meaningful in partitioned state + + see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) + free_param(param) + see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) + + if param.ds_tensor.final_location == OffloadDeviceEnum.nvme: + self.param_swapper.swap_out_and_release([param]) + print_rank_0(f"ID {param.ds_id} Offloaded to nvme offload and buffers released.") + see_memory_usage(f"ID {param.ds_id} Offloaded to nvme offload and buffers released.", force=False) + + print_rank_0(f"ID {param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}") + + @instrument_w_nvtx + def _partition_param_sec(self, param, buffer=None, has_been_updated=False): + assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" + global reuse_buffers + ##support for NVME secondary param offload + #print_rank_0(f"SEC Param id {param.ds_id} status is {param.ds_status}", force=True) + if param.ds_status is ZeroParamStatus.AVAILABLE: + if param.ds_secondary_tensor is not None and not has_been_updated: ##param already partitioned + return + #check padding + tensor_size = self._aligned_size(param) + partition_size = tensor_size // self.dp_world_size + + secondary_partition_size = int(tensor_size // self.num_ranks_in_param_group) + if param.ds_secondary_tensor is None: + final_location = None + secondary_partitioned_tensor = torch.empty(secondary_partition_size, + dtype=param.dtype, + device=self.remote_device) + + if self.pin_memory: + secondary_partitioned_tensor = secondary_partitioned_tensor.pin_memory() + # quantize the tensor if it's not trainable + if not param.requires_grad and self.quantized_nontrainable_weights: + secondary_partitioned_tensor, secondary_partitioned_tensor.ds_quant_scale = self.quantizer_module.quantize( + secondary_partitioned_tensor) + secondary_partitioned_tensor.requires_grad = False + param.ds_secondary_tensor = secondary_partitioned_tensor + param.ds_secondary_tensor.ds_numel = secondary_partition_size + param.ds_secondary_tensor.status = PartitionedParamStatus.AVAILABLE + param.ds_secondary_tensor.final_location = final_location + + #use rank in group for secondary tensor + secondary_start = secondary_partition_size * self.rank_in_group + + secondary_end = secondary_start + secondary_partition_size + + one_dim_param = param.contiguous().view(-1) + + # ds_numel is unpadded, so the last chunk of the secondary tensor might not be secondary_partition_size + sec_numel = max(0, min(param.ds_numel - secondary_start, secondary_partition_size)) + + # copy from full tensor to secondary tensor + param.ds_secondary_tensor.narrow(0, 0, + sec_numel).copy_(one_dim_param.narrow(0, secondary_start, sec_numel)) + + # TODO: This is a temporary fix to avoid the issue that 2nd tensor all-gather happens before 2nd tensor partition is done + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().synchronize() + + print_rank_0(f"{param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}", + force=False) + + def _param_status(self, param): + if param.ds_tensor is not None: + print_rank_0( + f"Param id {param.ds_id}, param status: {param.ds_status}, param numel {param.ds_numel}, partitioned numel {param.ds_tensor.numel()}, data numel {param.data.numel()}" + ) + else: + print_rank_0( + f"Param id {param.ds_id}, param status: {param.ds_status}, param numel {param.ds_numel}, partitioned ds_tensor {param.ds_tensor}, data numel {param.data.numel()}" + ) + + def _allgather_param(self, param, async_op=False, hierarchy=0): + + partition_size = param.ds_tensor.ds_numel + + tensor_size = partition_size * self.num_partitions + aligned_param_size = self._aligned_size(param) + assert tensor_size == aligned_param_size, f'param id {param.ds_id} aligned size {aligned_param_size} does not match tensor size {tensor_size}' + + print_rank_0( + f"{'--'* hierarchy}---- Before allocating allgather param {debug_param2name_id_shape_status(param)} partition size={partition_size}" + ) + + see_memory_usage( + f'Before allocate allgather param {debug_param2name_id_shape_status(param)} partition_size={partition_size} ', + force=False) + flat_tensor = torch.zeros(aligned_param_size, dtype=param.dtype, device=param.device).view(-1) + see_memory_usage( + f'After allocate allgather param {debug_param2name_id_shape_status(param)} {aligned_param_size} {partition_size} ', + force=False) + + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + + print_rank_0( + f"{'--'* hierarchy}----allgather param with {debug_param2name_id_shape_status(param)} partition size={partition_size}" + ) + # if not flat_tensor.numel() > 100000: + # replicated_tensor = flat_tensor.narrow(0, + # 0, + # param.ds_numel).view(param.ds_shape) + # param.data = replicated_tensor.data + # return None + if self.use_all_gather_into_tensor: + handle = dist.all_gather_into_tensor(flat_tensor, + param.ds_tensor.to(get_accelerator().device_name()), + group=self.get_partition_dp_group(param), + async_op=async_op) + else: + partitions = [] + for i in range(self.num_partitions): + partitions.append(flat_tensor.narrow(0, partition_size * i, partition_size)) + + if i == dist.get_rank(group=self.get_partition_dp_group(param)): + partitions[i].data.copy_(param.ds_tensor.data, non_blocking=True) + + handle = dist.all_gather(partitions, + partitions[self.get_partition_rank()], + group=self.get_partition_dp_group(param), + async_op=async_op) + + replicated_tensor = flat_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape) + param.data = replicated_tensor.data + return handle + + def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False): + """ blocking call + avoid explicit memory copy in _allgather_params + """ + if len(param_list) == 0: + return + + if self.num_partitions == 1: + handle = _no_gather_coalesced(param_list) + handle.wait() + return None + + # collect local tensors and partition sizes + partition_sizes = [] + local_tensors = [] + if quantize: + quantize_scale_sizes = [] + quantize_scale_tensors = [] + for param in param_list: + partition_sizes.append(param.ds_tensor.ds_numel) + local_tensors.append(param.ds_tensor.to(get_accelerator().device_name())) + if quantize: + quantize_scale_sizes.append(param.ds_tensor.ds_quant_scale.numel()) + quantize_scale_tensors.append(param.ds_tensor.ds_quant_scale.to(get_accelerator().device_name())) + # allocate memory for allgather params + allgather_params = [] + if quantize: + allgather_quantize_scale = [] + for psize in partition_sizes: + tensor_size = psize * self.num_partitions + flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, + device=self.local_device).view(-1) + flat_tensor.requires_grad = False + allgather_params.append(flat_tensor) + if quantize: + for psize in quantize_scale_sizes: + tensor_size = psize * self.num_partitions + flat_tensor = torch.empty(tensor_size, + dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, + device=self.local_device).view(-1) + flat_tensor.requires_grad = False + allgather_quantize_scale.append(flat_tensor) + + # launch + launch_handles = [] + launch_quantize_handles = [] + for param_idx, param in enumerate(param_list): + input_tensor = local_tensors[param_idx].view(-1) + + if self.use_all_gather_into_tensor: + # try the _all_gather_base from Pytorch master + h = dist.all_gather_into_tensor(allgather_params[param_idx], + input_tensor, + group=self.get_partition_dp_group(param), + async_op=True) + if quantize: + quantize_handle = dist.all_gather_into_tensor(allgather_quantize_scale[param_idx], + quantize_scale_tensors[param_idx], + group=self.get_partition_dp_group(param), + async_op=True) + launch_quantize_handles.append(quantize_handle) + else: + output_list = [] + for i in range(self.num_partitions): + psize = partition_sizes[param_idx] + partition = allgather_params[param_idx].narrow(0, i * psize, psize) + output_list.append(partition) + if not get_accelerator().on_accelerator(partition): + logger.warning( + f'param {param_idx}, partition {i} is not on CUDA, partition shape {partition.size()}') + + # back to old all_gather function + h = dist.all_gather(output_list, input_tensor, group=self.get_partition_dp_group(param), async_op=True) + if quantize: + output_scale_list = [] + for i in range(self.num_partitions): + psize = quantize_scale_sizes[param_idx] + partition = allgather_quantize_scale[param_idx].narrow(0, i * psize, psize) + output_scale_list.append(partition) + quant_handle = dist.all_gather(output_scale_list, + quantize_scale_tensors[param_idx], + group=self.get_partition_dp_group(param), + async_op=True) + launch_quantize_handles.append(quant_handle) + launch_handles.append(h) + + # Wait ensures the operation is enqueued, but not necessarily complete. + launch_handles[-1].wait() + if quantize: + for quant_handle in launch_quantize_handles: + quant_handle.wait() + + # assign to param.data (not copy) + for i, param in enumerate(param_list): + gathered_tensor = allgather_params[i] + if quantize: + gathered_tensor = self.quantizer_module.dequantize(gathered_tensor, allgather_quantize_scale[i]) + param.data = gathered_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape).data + + # guarantee the communication to be completed + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + + return None + + @torch.no_grad() + def _allgather_params(self, param_list, hierarchy=0): + if len(param_list) == 0: + return + + partition_size = sum([param.ds_tensor.ds_numel for param in param_list]) + + tensor_size = partition_size * self.num_partitions + flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, device=self.local_device) + partitions = [] + for i in range(self.num_partitions): + start = partition_size * i + + partitions.append(flat_tensor.narrow(0, start, partition_size)) + + if i == self.get_partition_rank(): + offset = 0 + for param in param_list: + param_numel = param.ds_tensor.ds_numel + + partitions[i].narrow(0, offset, param_numel).copy_(param.ds_tensor.data) + + offset += param_numel + + if hasattr(param_list[0], 'ds_quant_scale'): + scale_size = sum([param.ds_tensor.ds_quant_scale.numel() for param in param_list]) + scale_tensor_size = scale_size * self.world_size + flat_scale_tensor = torch.empty(scale_tensor_size, + dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, + device=self.local_device) + scale_partitions = [] + for i in range(self.world_size): + start = scale_tensor_size * i + scale_partitions.append(flat_scale_tensor.narrow(0, start, scale_tensor_size)) + if i == self.rank: + offset = 0 + for param in param_list: + param_scale_numel = param.ds_tensor.ds_quant_scale.ds_numel + + scale_partitions[i].narrow(0, offset, + param_scale_numel).copy_(param.ds_tensor.ds_quant_scale.data) + + offset += param_scale_numel + + dist.all_gather_into_tensor(flat_tensor, + partitions[self.get_partition_rank()], + group=self.get_partition_dp_group(param), + async_op=False) + if hasattr(param_list[0], 'ds_quant_scale'): + dist.all_gather(flat_scale_tensor, + param_list[0].ds_quant_scale, + group=self.get_partition_dp_group(param), + async_op=False) + param_offset = 0 + + for param in param_list: + param_partition_size = param.ds_tensor.ds_numel + param_size = param.ds_numel + replicated_tensor = torch.empty(param.ds_shape, dtype=param.ds_tensor.dtype, device=self.local_device) + + for i in range(self.num_partitions): + + start = i * partition_size + + param_start = i * param_partition_size + + if param_start < param_size: + numel_to_copy = min(param_size - param_start, param_partition_size) + + part_to_copy = partitions[i].narrow(0, param_offset, numel_to_copy) + + replicated_tensor.view(-1).narrow(0, param_start, numel_to_copy).copy_(part_to_copy) + #param_offset += param.data.numel() + param_offset += param.ds_tensor.ds_numel + if hasattr(param_list[0], 'ds_quant_scale'): + replicated_tensor = self.quantizer_module.dequantize(replicated_tensor, flat_scale_tensor) + param.data = replicated_tensor.data + + return None + + def _reduce_scatter_gradients(self, param_list): + #print_rank_0([param.grad for param in param_list]) + #assert any([param.grad is None for param in param_list]), "None gradients cannot be reduce scattered" + + handles_and_reduced_partitions = [] + for param in param_list: + assert param.grad.numel( + ) == param.ds_numel, f"{param.grad.numel()} != {param.ds_numel} Cannot reduce scatter gradients whose size is not same as the params" + + handles_and_reduced_partitions.append(self._reduce_scatter_gradient(param)) + + for param, (handle, reduced_partition) in zip(param_list, handles_and_reduced_partitions): + if handle is not None: + handle.wait() + + # some ranks may have partitions that are padded to go beyond the grad size. + # For these ranks the output of reduce scatter is a separate buffer and needs + # to be copied in + partition_size = param.ds_tensor.ds_numel + start = self.get_partition_rank() * partition_size + end = start + partition_size + #print_rank_0("REduce scatter was executed for param {param.ds_id}") + if start < param.ds_numel < end: + elements = param.ds_numel - start + param.grad.view(-1).narrow(0, start, elements).copy_(reduced_partition.narrow(0, 0, elements)) + + def _reduce_scatter_gradient(self, param): + + partition_size = param.ds_tensor.ds_numel + #output = torch.empty(partition_size, dtype=param.dtype, device=param.device) + + total_size = partition_size * self.num_partitions + input_list = [] + + for i in range(self.num_partitions): + + start = i * partition_size + end = start + partition_size + + #print("before reduce scatter gradients") + if start < param.ds_numel and end <= param.ds_numel: + input = param.grad.view(-1).narrow(0, start, partition_size) + else: + input = torch.zeros(partition_size, dtype=param.dtype, device=param.device) + + if start < param.ds_numel: + elements = param.ds_numel - start + input.narrow(0, 0, elements).copy_(param.grad.view(-1).narrow(0, start, elements)) + #print("after reduce scatter gradients") + input_list.append(input) + + rank = dist.get_rank(group=self.get_partition_dp_group(param)) + handle = dist.reduce_scatter(input_list[rank], + input_list, + group=self.get_partition_dp_group(param), + async_op=True) + + return handle, input_list[rank] + + def _partition_gradients(self, param_list, partition_buffers=None, accumulate=False): + if partition_buffers is None: + partition_buffers = [None] * len(param_list) + + for param, partition_buffer in zip(param_list, partition_buffers): + self._partition_gradient(param, partition_buffer=partition_buffer, accumulate=accumulate) + + def _partition_gradient(self, param, partition_buffer=None, accumulate=False): + + #import pdb;pdb.set_trace() + # param.grad=None + # param.grad.test() + print_rank_0( + f"Partitioning param {param.ds_id} gradient of size {param.grad.numel()} type {param.grad.dtype} part_size {param.ds_tensor.ds_numel}" + ) + see_memory_usage("Before partitioning gradients", force=False) + partition_size = param.ds_tensor.ds_numel + + if partition_buffer is None: + assert not accumulate, "No buffer to accumulate to" + partition_buffer = torch.zeros(partition_size, dtype=param.dtype, device=param.device) + else: + assert partition_buffer.numel( + ) >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}" + + rank = dist.get_rank(group=self.get_partition_dp_group(param)) + start = partition_size * rank + end = start + partition_size + + dest_tensor_full_buffer = partition_buffer.view(-1).narrow(0, 0, partition_size) + + #print("before partition gradients") + if start < param.ds_numel: + elements = min(param.ds_numel - start, partition_size) + + dest_tensor = dest_tensor_full_buffer.narrow(0, 0, elements) + src_tensor = param.grad.view(-1).narrow(0, start, elements) + + # just copy the grad partition to the buffer + if not accumulate: + dest_tensor.copy_(src_tensor) + + # if source and destination are on same device, + # add to the provided buffer + elif src_tensor.device == dest_tensor.device: + dest_tensor.add_(src_tensor) + + # if source and destination are on different device, copy first to src + # then add and move back to the destination. This seems to run faster + # when src is gpu and dest is cpu + # adding directly to cpu is very slow + else: + acc_tensor = torch.empty(src_tensor.numel(), dtype=param.dtype, device=param.device) + + acc_tensor.copy_(dest_tensor) + acc_tensor.add_(src_tensor) + dest_tensor.copy_(acc_tensor) + + # partition_buffer.view(-1).narrow( + # 0, + # 0, + # elements).copy_(param.grad.view(-1).narrow(0, + # start, + # elements)) + + #print("after partition gradients") + param.grad.data = dest_tensor_full_buffer.data + see_memory_usage("After partitioning gradients", force=False) + + def get_partition_dp_group(self, param): + return param.ds_process_group + + def get_partition_rank(self): + """subclass can overload to specify different relative rank in + parameter partition group""" + return self.rank + + @property + def num_partitions(self): + return self.dp_world_size + + def get_dp_process_group(self): + """ Return the communication group with all data-parallel ranks """ + return self.ds_process_group + + +class GatheredParameters: + + def __init__(self, params, modifier_rank=None, fwd_module=None, enabled=True): + """A context that collects parameters that were partitioned via a + :class:`deepspeed.zero.Init` context. The parameters are partitioned + again upon exit. + + Args: + params (``torch.nn.Parameter``): A single parameter, or an iterable of parameters (list, tuple, generator) of parameters to collect. + It's assumed that all parameters are zero params. + modifier_rank (int, optional): If specified, this rank's parameter will be + broadcasted on exit from the context. This argument is required if ``params`` are + modified, so that all processes have a consistent view of the data. Defaults + to ``None``. + fwd_module (``torch.nn.Module``, optional): If specified, ``params`` will be + registered as external parameters of ``fwd_module``. See :meth:`deepspeed.zero.register_external_parameter`. + enabled (bool, optional): If ``False``, this context is a no-op. Defaults to ``True``. + + Important: Make sure to use ``modifier_rank`` that is not ``None`` (e.g., ``modifier_rank=0``) + if you need the GPU memory allocated by gather to be released upon exit from the context manager. + + Important: if ``params`` isn't an iterable of parameters or a single parameter it'll be silently ignored! + + Examples + ======== + + #. Allocate a partitioned module, initialize its weight on rank 0, and update all + processes. + + .. code-block:: python + + with deepspeed.zero.Init(): + linear = torch.nn.Linear(1000,1000) + + with deepspeed.zero.GatheredParameters(linear.weight, + modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + linear.weight.zero_() + + with deepspeed.zero.GatheredParameters(linear.weight, + modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + linear.weight.zero_() + + #. Collect a partitioned weight to pass to another module during + training. The parameter will be registered as an external parameter + and made available during the backward pass. + + .. code-block:: python + :emphasize-lines: 6 + + def forward(self, input): + x = self.layer1(input) + + # self.layer1.weight is required by self.layer2.forward + with deepspeed.zero.GatheredParameters(self.layer1.weight, + fwd_module=self): + y = self.layer2(x, self.layer1.weight) + return y + + + #. Pretrained model loading + + .. code-block:: python + + with deepspeed.zero.Init(): + model = MyModel() + + state_dict = torch.load(model_path, map_location="cpu") + + def load(module: nn.Module, prefix=""): + # because zero3 puts placeholders in model params, this context + # manager gathers (unpartitions) the params of the current layer, then loads from + # the state dict and then re-partitions them again + with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + module._load_from_state_dict(state_dict, prefix) + + for name, child in module._modules.items(): + if child is not None: + load(child, prefix + name + ".") + + load(model, prefix="") + + If this approach is not used, then the full model will first be copied to each GPU. For models + bigger than the memory of a single GPU, this method is required. + """ + + self.enabled = enabled + if not enabled: + return + + if isinstance(params, Iterable) and not isinstance(params, torch.Tensor): + # deal with generators like model.parameters() + # must convert to list to be able to iterate more than once if we get a generator + params = list(params) + else: + # single param + params = [params] + # enable if at least one is zero-param, otherwise a noop + if not any(is_zero_param(p) for p in params): + self.enabled = False + return + + self.params = [p for p in params if hasattr(p, "ds_id")] + self.params = sorted( + set(self.params), key=lambda x: x.ds_id + ) # remove the duplicates to prevent racing condition, we must also make sure the order is the same on all ranks otherwise we'll get deadlocks + self.src_rank = None + if modifier_rank is not None: + if self.params[0].ds_process_group == dist.get_world_group(): + self.src_rank = modifier_rank + else: + # A group was specified; convert DP rank to global rank + self.src_rank = dist.get_global_rank(self.params[0].ds_process_group, modifier_rank) + self.fwd_module = fwd_module + if self.fwd_module is not None: + # is a no-op if already registered + for p in self.params: + register_external_parameter(self.fwd_module, p) + + def __enter__(self): + if not self.enabled: + return + self.params[0].all_gather(param_list=self.params) + + def __exit__(self, *exc): + if not self.enabled: + return + if self.src_rank is None: + self.params[0].partition(param_list=self.params, has_been_updated=False) + return + + handles = [dist.broadcast(p.data, self.src_rank, group=p.ds_process_group, async_op=True) for p in self.params] + for h in handles: + h.wait() + self.params[0].partition(param_list=self.params, has_been_updated=True) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py new file mode 100644 index 0000000000000000000000000000000000000000..80b20cb1d3e883d91a5dce95a16a496c42b9faef --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py @@ -0,0 +1,585 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass +import collections +from collections import UserDict +from typing import Deque, Set + +from deepspeed import comm as dist +from deepspeed.utils import z3_leaf_module +from deepspeed.utils.logging import logger +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.zero.partition_parameters import * +from deepspeed.runtime.zero.partitioned_param_profiler import PartitionedParameterProfiler +from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus +from deepspeed.utils.debug import debug_param2name_id_shape +from deepspeed.accelerator import get_accelerator +import deepspeed.runtime.compiler as compiler +from deepspeed.runtime.compiler import is_compiling + +import logging + +ENABLE_PROFILER = False + + +def debug_rank0(message: str) -> None: + if dist.get_rank() == 0: + logger.debug(message) + + +@instrument_w_nvtx +def get_all_parameters(sub_module, recurse=False): + return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters()) + + +@compiler.disable +def iter_params(module: Module, recurse=False) -> Iterable[Parameter]: + return map(lambda pair: pair[1], get_all_parameters(module, recurse)) + + +class ZeRoTraceMode(Enum): + # Record trace of the network during a single forward+backward (for training) or forward (for inference) + RECORD = 1 + # Use recorded network trace to optimize current forward+backward or forward + COMPLETE = 2 + # Recorded trace does not match current forward+backward or forward pass. + INVALID = 3 + + +class InflightParamRegistry(UserDict): + """registry for parameters in flight""" + + def __setitem__(self, param: Parameter, handle: AllGatherCoalescedHandle) -> None: + if param in self.data: + raise RuntimeError(f"{param.ds_summary()} already in registry") + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"attempted to add non-inflight parameter to registry {param.ds_summary()}") + self.data[param] = handle + + +class PartitionedParameterCoordinator: + FORWARD_FETCH_SUBMIT = 'forward_fetch_submit' + FORWARD_FETCH_WAIT = 'forward_fetch_wait' + FORWARD_PREFETCH_SUBMIT = 'forward_prefetch_submit' + BACKWARD_FETCH_SUBMIT = 'backward_fetch_submit' + BACKWARD_FETCH_WAIT = 'backward_fetch_wait' + BACKWARD_PREFETCH_SUBMIT = 'backward_prefetch_submit' + FORWARD_ALL_GATHER = 'forward_all_gather' + BACKWARD_ALL_GATHER = 'backward_all_gather' + """Handles partitioning and gathering of parameters.""" + + @dataclass + class __ParamInTrace: + param: Parameter + step_id_last_used_at: int + + def __init__( + self, + prefetch_bucket_sz: int, + max_reuse_distance_in_numel: int, + max_available_parameters_in_numel: int, + allgather_stream: get_accelerator().Stream, + inflight_param_registry: InflightParamRegistry, + prefetch_nvme: bool = False, + timers=None, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + fast_sharding_for_leaf_module=False, + log_trace_cache_warnings=False, + ) -> None: + # mapping of param -> handle for each param that is currently in flight + self.__inflight_param_registry = inflight_param_registry + # keeps track of the number of submodules invoked so far. + self.__step_id: int = 0 + # network tracing mode + self.__trace_mode: ZeRoTraceMode = ZeRoTraceMode.INVALID + # sequence of submodules/parameters in forward pass + backward pass + self.__submodule_order: Iterable[Module] = [] + self.__param_order: Iterable[__class__.__ParamInTrace] = [] + self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) + self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque()) + # number of available params, and max number of available params + self.__n_available_params: int = 0 + self.__max_n_available_params: int = max_available_parameters_in_numel + # max distance between two use of the module beyond which module is released + self.__max_reuse_dist_in_numel: int = max_reuse_distance_in_numel + # queue for parameters to fetch. parameters will be popped off the left + # side of the dequeue as they are fetched + self.__param_queue: Deque[__class__.__ParamInTrace] = None + self.__prefetch_bucket_sz: int = prefetch_bucket_sz + self.__prefetch_nvme: bool = prefetch_nvme + self.hierarchy: int = 0 + self.zero_quantized_weights = zero_quantized_weights + self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights + + # stream that will be used for allgather operations + self.__allgather_stream: get_accelerator().Stream = allgather_stream + + # limit the number of fetch events that can be queued at once + # otherwise, what happens is memory is allocated by the host thread at the + # time of the call, but not used until later by the asynchronous cuda stream. + # allowing an infinite number of these to queue up causes a lot of memory + # pressure that then becomes detrimental to performance. + # this is a much less elegant way of fixing this vs something like using + # cudaMallocAsync/cudaFreeAsync. Choosing to not expose this to the user now + # because ideally in the future its replaced by an async allocation + # mechanism which doesn't require any configuration by the user. + self.__ongoing_fetch_events: Deque[get_accelerator().Event] = collections.deque() + # TODO. make this configurable via JSON + self.__max_ongoing_fetch_events: int = 2 + self.__profiler = PartitionedParameterProfiler(timers if ENABLE_PROFILER else None) + + # Whether to log trace cache warnings, e.g. invalidation events + self.__log_trace_cache_warnings = log_trace_cache_warnings + + # whether to enable fast fetch for the z3 leaf module. + # this will improve fetch speed but will not break down leaf module parameters to alleviate memory pressure. + self.fast_sharding_for_leaf_module = fast_sharding_for_leaf_module + + """Tracing and Tracking + TODO. consider performing trace before initializing PartitionedParameterCoordinator + and passing trace results into constructor. This way all the code in here can + just assume that the trace is complete and the results can be entirely + immutable. + + Bookkeeping operations used to track where we are in the forward/backward pass + """ + + def _clear_trace_structures(self) -> None: + self.__submodule_order = [] + self.__param_order = [] + self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) + self.__param_queue = None + + def is_complete_trace(self) -> bool: + return self.__trace_mode == ZeRoTraceMode.COMPLETE + + def is_invalid_trace(self) -> bool: + return self.__trace_mode == ZeRoTraceMode.INVALID + + def is_record_trace(self) -> bool: + return self.__trace_mode == ZeRoTraceMode.RECORD + + def _clean_inflight_param_registry(self) -> None: + for param, handle in self.__inflight_param_registry.items(): + handle.wait() + self.__release_param(param) + self.__inflight_param_registry.clear() + + def _invalidate_trace(self) -> None: + if self.is_invalid_trace(): + raise RuntimeError("attempted to invalidate already invalid trace") + self.__trace_mode = ZeRoTraceMode.INVALID + self._clear_trace_structures() + self._clean_inflight_param_registry() + + def trace_prologue(self, sub_module: Module) -> None: + if self.is_complete_trace(): + # sub_module must match expectation else invalidate trace cache + if len(self.__submodule_order) <= self.__step_id: + print_rank_0( + f"Invalidate trace cache @ step {self.__step_id} and module {sub_module.ds_id}: " + f"cache has only {len(self.__submodule_order)} modules", + force=self.__log_trace_cache_warnings) + self._invalidate_trace() + return + + if sub_module != self.__submodule_order[self.__step_id]: + expected_module_id = self.__submodule_order[self.__step_id].ds_id + print_rank_0( + f"Invalidate trace cache @ step {self.__step_id}: " + f"expected module {expected_module_id}, but got module {sub_module.ds_id}", + force=self.__log_trace_cache_warnings) + self._invalidate_trace() + + @compiler.disable + def record_module(self, sub_module: Module) -> None: + """adds sub module to trace""" + if is_compiling(): + return + + if not self.is_record_trace(): + raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}") + + self.__submodule_order.append(sub_module) + self.__step_id_module_fetched_for[sub_module.ds_id].append(self.__step_id) + + def record_parameters(self, sub_module: Module) -> None: + if is_compiling(): + return + """adds sub module to trace""" + if not self.is_record_trace(): + raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}") + + step_id = self.__step_id_module_fetched_for[sub_module.ds_id].popleft() + for param in sorted(set(iter_params(sub_module, recurse=z3_leaf_module(sub_module))), key=lambda p: p.ds_id): + self.__param_order.append(__class__.__ParamInTrace(param=param, step_id_last_used_at=step_id)) + + def construct_parameter_trace_from_module_trace(self): + """use module trace to construct parameter trace""" + self.__param_order = [] + for sub_module in self.__submodule_order: + self.record_parameters(sub_module) + + @compiler.disable + def reset_step(self) -> None: + """indicate that we have completed one fwd+bwd for the model""" + if is_compiling(): + return + + self._clean_inflight_param_registry() + + if not self.is_complete_trace(): # not self.trace_complete: + # Make sure that recorded submodule orders are identical across ranks + assert_ints_same_as_other_ranks([m.ds_id for m in self.__submodule_order]) + + if self.is_record_trace(): + # Successfully recorded a trace + self.construct_parameter_trace_from_module_trace() + # Make sure that recorded parameter orders are identical across ranks + assert_ints_same_as_other_ranks([p.param.ds_id for p in self.__param_order]) + assert_ints_same_as_other_ranks([p.step_id_last_used_at for p in self.__param_order]) + + self.__submodule_order = tuple(self.__submodule_order) # freeze + self.__param_order = tuple(self.__param_order) # freeze + self.__trace_mode = ZeRoTraceMode.COMPLETE + print_rank_0( + f"completed record trace of {len(self.__submodule_order)} sub modules: {[m.ds_id for m in self.__submodule_order]}", + force=False) + else: + # Enable trace recording for next forward/backward pass + self.__trace_mode = ZeRoTraceMode.RECORD + + else: + if self.__profiler is not None: + self.__profiler.log_events() + + self.__param_queue = collections.deque(self.__param_order) # reset fetch queue + self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) + self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque()) + self.__step_id = 0 + self.__n_available_params = 0 + self.__profiler.reset_events() + + def _dump_params(self, tag, sub_module, params, step_id=None): + if step_id is None: + step_id = self.__step_id + param_names = [debug_param2name_id_shape(p) for p in params] + print_rank_0(f'{tag} step = {step_id} p_names = {param_names}', force=False) + + def _dump_param_ids(self, tag, mod_id, p_ids, step_id=None): + if step_id is None: + step_id = self.__step_id + print_rank_0(f'{tag} mod = {mod_id}, step = {step_id}, p_ids = {p_ids}', force=False) + + """Fetch and Release + Fetching, prefetching, and releasing parameters + """ + + @compiler.disable + @instrument_w_nvtx + @torch.no_grad() + def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None: + """This method does the following (in order): + 1. kick off fetch for parameters in immediately required sub module + 2. kick off fetch for next few parameters we will need later (prefetch) + 3. block on parameters in immediately required sub module + """ + if logger.isEnabledFor(logging.DEBUG): + debug_rank0( + f"{self.__step_id}: M{current_submodule.ds_id}({type(current_submodule).__name__}) P{[p.ds_id for p in iter_params(current_submodule, recurse=z3_leaf_module(current_submodule))]} " + + str({ + "avail": f"{self.__n_available_params:.1e}", + "queue_sz": f"{len(self.__param_queue or [])}", + "inflight": [p.ds_id for p in self.__inflight_param_registry], + })) + + params_to_fetch = set(iter_params(current_submodule, recurse=z3_leaf_module(current_submodule))) + fetch_numel = sum( + [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + + if fetch_numel > 0: + event_name = __class__.FORWARD_FETCH_SUBMIT if forward else __class__.BACKWARD_FETCH_SUBMIT + self._dump_param_ids(event_name, current_submodule.ds_id, + [(p.ds_id, p.ds_shape) + for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + # self._dump_params(event_name, current_submodule, [p for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + + self.__profiler.start_event(event_name) + # kick off all gather for params in the immediately required submodule + #for param in params_to_fetch: + if logger.isEnabledFor(logging.DEBUG): + for param in params_to_fetch: + debug_rank0(f"-fetch: {param.ds_summary()}") + self.__all_gather_params(params_to_fetch, forward) + self.__profiler.stop_event(event_name, fetch_numel) + + wait_numel = 0 + wait_event_name = __class__.FORWARD_FETCH_WAIT if forward else __class__.BACKWARD_FETCH_WAIT + self.__profiler.start_event(wait_event_name) + fast_fetch = self.fast_sharding_for_leaf_module and z3_leaf_module(current_submodule) + # wait for parameters in the immediately needed submodule to become available + for param in params_to_fetch: + param.ds_active_sub_modules.add(current_submodule.ds_id) + if logger.isEnabledFor(logging.DEBUG): + debug_rank0(f"-wait: {param.ds_summary()}") + if param in self.__inflight_param_registry: + wait_numel += param.partition_numel() + with get_accelerator().stream(self.__allgather_stream): + while self.__ongoing_fetch_events and self.__ongoing_fetch_events[0].query(): + self.__ongoing_fetch_events.popleft() + if len(self.__ongoing_fetch_events) > self.__max_ongoing_fetch_events: + self.__ongoing_fetch_events.popleft().synchronize() + + self.__inflight_param_registry.pop(param).wait(handle_dependency=not fast_fetch) + + if not get_accelerator().handles_memory_backpressure() and not fast_fetch: + event = get_accelerator().Event() + event.record() + self.__ongoing_fetch_events.append(event) + + assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary() + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().wait_stream(self.__allgather_stream) + if fast_fetch: + AllGatherCoalescedHandle.free_buffer() + self.__profiler.stop_event(wait_event_name, wait_numel) + + # kick off parameter prefetches for upcoming modules + # don't prefetch if we dont have a completed model trace + if self.is_complete_trace(): + # go through the parameters we need for the current module and pop them + # off the fetch queue so that they aren't prefetched later. + # if params have already been popped off the fetch queue by earlier + # prefetches we won't look for them here + discarded_from_prefetch_queue = set() + params_not_already_fetched = set( + filter(lambda p: self.__most_recent_step_id_param_fetched_for[p] < self.__step_id, params_to_fetch)) + while self.__param_queue and len(discarded_from_prefetch_queue) < len(params_not_already_fetched): + param_in_trace = self.__param_queue.popleft() + self.__most_recent_step_id_param_fetched_for[ + param_in_trace.param] = param_in_trace.step_id_last_used_at + discarded_from_prefetch_queue.add(param_in_trace.param) + + if discarded_from_prefetch_queue != params_not_already_fetched: + raise RuntimeError( + f"tracing error at step {self.__step_id}: \n" + f"module id: {current_submodule.ds_id}, training: {current_submodule.training}\n" + f"expected the next {len(params_not_already_fetched)} parameters in the " + f"parameter fetch queue to be {tuple(p.ds_summary(use_debug_name=True) for p in params_not_already_fetched)} \n" + f"but got \n {tuple(p.ds_summary(use_debug_name=True) for p in discarded_from_prefetch_queue)}.") + + def _is_currently_on_nvme(param): + if param.nvme_swapper is None: + return False + + return param.ds_tensor.final_location == OffloadDeviceEnum.nvme \ + and param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE + + # kick off all gather for params in the next few submodules (prefetch) + if self.__prefetch_bucket_sz > 0: + max_params_to_prefetch = min(self.__max_n_available_params - self.__n_available_params, + self.__prefetch_bucket_sz) + params_to_prefetch = set() + numel_prefetching = 0 + while self.__param_queue and numel_prefetching < max_params_to_prefetch: + param_in_trace: __class__.__ParamInTrace = self.__param_queue.popleft() + + if _is_currently_on_nvme(param_in_trace.param): + # nvme prefetch is handled elsewhere. Need to break here to preserve fetch order + self.__param_queue.appendleft(param_in_trace) + break + + do_prefetch = param_in_trace.param.ds_status == ZeroParamStatus.NOT_AVAILABLE + if param_in_trace.param in params_to_prefetch: + # Avoid duplicates + do_prefetch = False + + self.__most_recent_step_id_param_fetched_for[param_in_trace.param] = \ + max(self.__most_recent_step_id_param_fetched_for[param_in_trace.param], + param_in_trace.step_id_last_used_at) + + if do_prefetch: + params_to_prefetch.add(param_in_trace.param) + numel_prefetching += param_in_trace.param.ds_numel + + if numel_prefetching > 0: + event_name = __class__.FORWARD_PREFETCH_SUBMIT if forward else __class__.BACKWARD_PREFETCH_SUBMIT + self.__profiler.start_event(event_name) + if logger.isEnabledFor(logging.DEBUG): + for param in params_to_prefetch: + debug_rank0(f"-prefetch: {param.ds_summary()}") + self.__all_gather_params(params_to_prefetch, forward) + self.__profiler.stop_event(event_name, numel_prefetching) + + if self.__prefetch_nvme: + self.__prefetch_nvme_param_partitions() + + self.__step_id += 1 + + @instrument_w_nvtx + @torch.no_grad() + def release_sub_module(self, submodule: Module, forward=False) -> None: + """release the parameters of a sub module, assuming they meet conditions to + be released.""" + #print_rank_0(f"release_sub_module {'fwd' if forward else 'bwd'}: {debug_module2name_id(submodule)}", force=False) + params_to_release = (self.__params_to_release(submodule, self.__step_id) if self.is_complete_trace() else set( + p.ds_id for p in iter_params(submodule, recurse=z3_leaf_module(submodule)))) + + free_data = not z3_leaf_module(submodule) or not self.fast_sharding_for_leaf_module + if not free_data: + # wait for the computation to finish and launch as early as possible. + empty_buffer = torch.empty(1, device=get_accelerator().current_device()) + + for param in iter_params(submodule, recurse=z3_leaf_module(submodule)): + param.ds_active_sub_modules.discard(submodule.ds_id) + if param.ds_id in params_to_release and not param.is_external_param: + self.__release_param(param, free_data) + if not free_data: + if param.ds_id in params_to_release and not param.is_external_param: + # empty buffer ensures that all computations are complete + param.data = empty_buffer + + @instrument_w_nvtx + @torch.no_grad() + def release_and_reset_all(self, module: Module) -> None: + """release all module parameters""" + for param in iter_params(module, recurse=True): + if param in self.__inflight_param_registry: + self.__inflight_param_registry.pop(param).wait() + + # TODO. make this throw if if there are still active submodules. currently + # there's a hook execution issue + param.ds_active_sub_modules.clear() + self.__release_param(param) + + for param in iter_params(module, recurse=True): + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f"{param.ds_summary()} expected to be released") + + @instrument_w_nvtx + def __all_gather_params(self, params: Set[Parameter], forward: bool) -> None: + quantized_params = [] + nonquantized_params = [] + for param in params: + if hasattr(param.ds_tensor, 'ds_quant_scale'): + quantized_params.append(param) + else: + nonquantized_params.append(param) + if quantized_params: + self.__all_gather_params_(quantized_params, forward, quantize=True) + if nonquantized_params: + self.__all_gather_params_(nonquantized_params, forward, quantize=self.zero_quantized_weights) + + def __all_gather_params_(self, params: Set[Parameter], forward: bool, quantize: bool = False) -> None: + """for each partitioned parameter, kick off an async allgather and store + the work handle for the in flight parameters.""" + partitioned_params = [] + all_gather_numel = 0 # numel = num of elements + for param in params: + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + partitioned_params.append(param) + all_gather_numel += param.ds_numel + + if partitioned_params: + self.__n_available_params += all_gather_numel + # here we need to handle a special case where some of the parameters have a valid hpz secondary tensor (e.g. they are not trainable so their secondary tensor never expire) but others do not. + partitioned_params_with_secondary_tensors = [ + p for p in partitioned_params if p.ds_secondary_tensor is not None + ] + partitioned_params_without_secondary_tensors = [ + p for p in partitioned_params if p.ds_secondary_tensor is None + ] + for param_group in [ + partitioned_params_with_secondary_tensors, partitioned_params_without_secondary_tensors + ]: + if not param_group: + continue + with get_accelerator().stream(self.__allgather_stream): + event_name = __class__.FORWARD_ALL_GATHER if forward else __class__.BACKWARD_ALL_GATHER + self.__profiler.start_event(event_name) + handle = param_group[0].all_gather_coalesced(param_group, quantize=quantize) + self.__profiler.stop_event(event_name, all_gather_numel) + for param in param_group: + assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary() + self.__inflight_param_registry[param] = handle + + # Release swap buffers for persisted params on nvme since they will never be partitioned or evicted from GPU + swap_persisted_params = [ + p for p in partitioned_params if p.ds_persist and p.ds_tensor.final_location == OffloadDeviceEnum.nvme + ] + if swap_persisted_params: + swap_persisted_params[0].nvme_swapper.remove_partition_and_release_buffers(swap_persisted_params) + + @compiler.disable + @instrument_w_nvtx + def __release_param(self, param: Parameter, free_data: bool = True) -> None: + if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules: + if logger.isEnabledFor(logging.DEBUG): + debug_rank0(f"-release: {param.ds_summary()}") + print_rank_0(f"release: {debug_param2name_id_shape(param)}", force=False) + param.partition(free_data=free_data) + self.__n_available_params -= param.ds_numel + + @instrument_w_nvtx + @functools.lru_cache(maxsize=None) + def __params_to_release(self, submodule_to_release: Module, step_id: int) -> Set[int]: + if not self.is_complete_trace(): + raise RuntimeError("expected trace to be complete") + + params_to_release = set( + p.ds_id for p in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)) + if not p.ds_persist) + + # Problem: When prefetcher scans the param trace, it skips AVAILABLE params. + # This creates issues if those params are released before the skipped uses: + # 1) It hurts performance as the skipped uses are never prefetched. + # 2) For nvme params, we run out of swap buffers because the prefetch order + # diverges from the trace. + # Solution: Don't release params whose reuse was skipped by prefetch. This is + # possible because we detect such skips during prefetch and mark those params. + for param in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)): + if self.__most_recent_step_id_param_fetched_for[param] > step_id: + params_to_release.discard(param.ds_id) + + # examine all modules within `max_reuse_dist_in_numel` of the current step, + # if we see any of the candidate parameters to be released reoccur while + # doing this, remove them from the set of parameters to release. + params_traversed = 0 + for module in self.__submodule_order[step_id:]: + if params_traversed >= self.__max_reuse_dist_in_numel: + break + for param in iter_params(module, recurse=z3_leaf_module(submodule_to_release)): + params_to_release.discard(param.ds_id) + params_traversed += param.ds_numel + + return params_to_release + + @instrument_w_nvtx + def __prefetch_nvme_param_partitions(self) -> None: + """swap in parameter partitions from nvme for those parameters that will be used + after the ones that are already being prefetched into full parameters + """ + if not self.is_complete_trace(): + return + + numel_in_flight = sum(param.ds_numel for param in self.__inflight_param_registry) + + numel_considered = 0 + swap_in_params = [] + for param_in_trace in self.__param_queue: + param = param_in_trace.param + if param.nvme_swapper is None: + continue + if (numel_considered > 2 * numel_in_flight + or len(swap_in_params) >= param.nvme_swapper.available_swap_in_buffers()): + break + if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE: + swap_in_params.append(param) + numel_considered += param.ds_numel + + if swap_in_params: + swap_in_params[0].nvme_swapper.swap_in(swap_in_params, async_op=True) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..b4ea11f3b83632a5f7f74343ef7092af6c5bf3a7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass +from deepspeed.utils import log_dist + + +class PartitionedParameterProfiler(object): + + @dataclass + class EventCounter: + name: str + count: int + num_elem: int + + def reset(self): + self.count = 0 + self.num_elem = 0 + + def increment(self, numel): + self.count += 1 + self.num_elem += numel + + def __init__(self, timers): + self.timers = timers + self.event_counters = {} + + def reset_events(self): + for event_ctr in self.event_counters.values(): + event_ctr.reset() + + def start_event(self, name): + if self.timers is None: + return + + if name not in self.event_counters: + self.event_counters[name] = __class__.EventCounter(name=name, count=0, num_elem=0) + self.timers(name).start() + + def stop_event(self, name, num_elem): + if self.timers is None: + return + assert name in self.event_counters, f'unknown event {name}' + self.event_counters[name].increment(num_elem) + self.timers(name).stop() + + def _log_timers(self): + if self.timers is None: + return + self.timers.log(names=list(self.event_counters.keys())) + + def _log_event_counters(self): + for event_ctr in self.event_counters.values(): + log_dist( + f'{event_ctr.name}: count = {event_ctr.count}, numel = {event_ctr.num_elem}', + #f'{event_ctr.name}: time = {self._log_timers()},count = {event_ctr.count}, numel = {event_ctr.num_elem}', + ranks=[0]) + + def log_events(self): + self._log_event_counters() + self._log_timers() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py new file mode 100644 index 0000000000000000000000000000000000000000..2937191a0ad56136ab2bb9219d487e42f251f3b2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py @@ -0,0 +1,3262 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import gc +import collections +import itertools +from typing import Deque, Dict, Set, List, Tuple, Container, Optional +from contextlib import contextmanager +from dataclasses import dataclass, field + +from deepspeed import comm as dist +from deepspeed.utils import groups, z3_leaf_parameter + +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed.runtime.base_optimizer import ZeROOptimizer +from deepspeed.utils import logger +from deepspeed.utils.torch import register_grad_hook +from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler +from deepspeed.runtime.torch_autocast import get_all_autocast_dtypes, is_autocast_initialized, sort_dtypes +from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced, all_to_all_quant_reduce, all_to_all_loco_quant_reduce +from deepspeed.runtime.utils import inf, is_model_parallel_parameter, get_only_unique_item, mask_nan_or_inf_with_val_inplace +from deepspeed.runtime.zero.partition_parameters import * +from deepspeed.runtime.zero.config import ZeroStageEnum +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum +from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload +from deepspeed.runtime.zero.utils import apply_to_tensors_only, get_mapping_to_flat_buffer +from deepspeed.runtime.zero.offload_states import offload_adam_states, reload_adam_states +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus +from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper +from deepspeed.runtime.swap_tensor.partitioned_optimizer_swapper import PartitionedOptimizerSwapper +from deepspeed.runtime.swap_tensor.pipelined_optimizer_swapper import PipelinedOptimizerSwapper +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FP32_FLAT_GROUPS, PARTITION_COUNT, ZERO_STAGE, LOSS_SCALER +from deepspeed.accelerator import get_accelerator + +# Toggle this to true to enable correctness test +# with gradient partitioning and without +pg_correctness_test = False + +OPTIMIZER_SWAP_IN_STATE_TIMER = 'optimizer_swap_in_state' +INIT_OPTIMIZER_TIMER = 'init_optimizer_state' +OPTIMIZER_SWAP_OUT_STATE_TIMER = 'optimizer_swap_out_state' +OPTIMIZER_STEP_TIMER = 'optimizer_step' + + +def print_rank_0(message, debug=False, force=False): + rank = dist.get_rank() + if rank == 0 and (debug or force): + logger.info(message) + # other variations + # - print for all ranks w/o interleaving + # printflock(f"[{rank}] {message}") + # - print to log file per rank + # log_rank_file(rank, message) + + +def input(msg): + return + + +def isclose(a, b, rtol=1e-09, atol=0.0): + return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) + + +def lcm(x, y): + from fractions import gcd # or can import gcd from `math` in Python 3 + return x * y // gcd(x, y) + + +def move_to_cpu(tensor_list): + for tensor in tensor_list: + tensor.data = tensor.data.cpu() + + +@contextmanager +def unwrap_model_for_generation(model): + """ + For ZeRO-3 models, we gather the weights once to speed up generation. + """ + with GatheredParameters(model.parameters()): + # Removes the optimizer hooks from a DeepSpeed ZeRO-3 model. + + # Remove hooks + if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): + optimizer_offload = model.optimizer.parameter_offload + elif model.optimizer is not None: + optimizer_offload = model.optimizer + + for hook in optimizer_offload.forward_hooks: + hook.remove() + for hook in optimizer_offload.backward_hooks: + hook.remove() + + optimizer_offload.forward_hooks = [] + optimizer_offload.backward_hooks = [] + + yield model + + # Adds the optimizer hooks from a DeepSpeed ZeRO-3 model. + if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): + optimizer_offload = model.optimizer.parameter_offload + elif model.optimizer is not None: + optimizer_offload = model.optimizer + optimizer_offload._register_deepspeed_module(optimizer_offload.module) + return + + +@dataclass +class IPGBucketZ3: + buffer: Optional[torch.Tensor] = None + buffer_meta: Optional[torch.Tensor] = None + params: List[torch.Tensor] = field(default_factory=list) + elements: int = 0 + + def clear(self): + self.buffer = None + self.buffer_meta = None + self.params.clear() + self.elements = 0 + + +INITIAL_MICRO_STEP_ID = -1 + + +class DeepSpeedZeroOptimizer_Stage3(ZeROOptimizer): + """ + DeepSpeedZeroOptimizer designed to reduce the memory footprint + required for training large deep learning models. + + For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models + https://arxiv.org/abs/1910.02054 + + For usage examples, refer to TODO: DeepSpeed Tutorial + + """ + + def __init__( + self, + module, + init_optimizer, + timers, + ds_config, + static_loss_scale=1.0, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + prefetch_bucket_size=50000000, + max_reuse_distance=1000000000, + max_live_parameters=1000000000, + param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, + dp_process_group=None, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + offload_param_config=None, + sub_group_size=1000000000000, + offload_ratio=0.0, + mpu=None, + clip_grad=0.0, + gradient_accumulation_dtype=torch.float32, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1.0, + gradient_accumulation_steps=1, + elastic_checkpoint=False, + aio_config=None, + all2all_process_group=None, + zero_hpz_partition_size=1, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + zero_module_granularity_threshold=0, + zeropp_loco_param=None, + log_trace_cache_warnings=False, + ): + see_memory_usage("Stage 3 initialize beginning", force=True) + + print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) + + if dist.get_rank() == 0: + logger.info(f"Reduce bucket size {reduce_bucket_size}") + logger.info(f"Prefetch bucket size {prefetch_bucket_size}") + # The fused optimizer does all the work. We need this layer for two reason: + # 1. maintain same user API from apex.fp16_utils + # 2. keep common stuff here in case we need to add ne552w fused optimizer later + + # differences from apex.fp16_utils: + # - assume all model params in fp16 + # - assume all params requires grad + # - flat by groups, not keeping state. TODO: remove state explicitly? + # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? + if not get_accelerator().is_available(): + raise SystemError("Cannot use fp16 without accelerator.") + + self.optimizer = init_optimizer + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + self.dtype = self.optimizer.param_groups[0]['params'][0].dtype + self.gradient_accumulation_dtype = gradient_accumulation_dtype + self._global_grad_norm = 0. + + self.custom_loss_scaler = False + self.external_loss_scale = None + + self.optimizer_swapper = None + self.swap_optimizer = False + + self.offload_optimizer = False + self.offload_optimizer_pin_memory = False + self.offload_optimizer_fast_init = False + self.offload_param = False + self.offload_param_pin_memory = False + self.params_in_nvme_and_cpu = False + self.max_params_in_cpu = 0 + self.partial_offload = offload_ratio + + #num of ranks in a ZeRO param partitioning group + self.zero_hpz_partition_size = zero_hpz_partition_size + + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + print_rank_0( + f"ZeRO Stage 3 param partitioning group {self.zero_hpz_partition_size} {zero_param_parallel_group}", + force=False) + if self.zero_hpz_partition_size > 1 and zero_param_parallel_group is None: + self._set_zero_group_parallelism() + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + + self.parameter_offload = self.initialize_ds_offload( + module=module, + timers=timers, + ds_config=ds_config, + overlap_comm=overlap_comm, + prefetch_bucket_size=prefetch_bucket_size, + max_reuse_distance=max_reuse_distance, + max_live_parameters=max_live_parameters, + param_persistence_threshold=param_persistence_threshold, + model_persistence_threshold=model_persistence_threshold, + dp_process_group=dp_process_group, + offload_param_config=offload_param_config, + mpu=mpu, + zero_param_parallel_group=zero_param_parallel_group, + zero_quantized_weights=zero_quantized_weights, + zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights, + zero_module_granularity_threshold=zero_module_granularity_threshold, + log_trace_cache_warnings=log_trace_cache_warnings, + ) + + self.persistent_parameters = self.parameter_offload.persistent_parameters + self._configure_offloading(offload_optimizer_config, offload_param_config) + + # backup fused_adam optimizer init + if self.offload_optimizer and self.partial_offload != 1.0: + backup_gpu_tensor = torch.randn(1, device=get_accelerator().device_name()).to(self.dtype) + backup_gpu_param = torch.nn.Parameter(backup_gpu_tensor) + assert type(init_optimizer) == DeepSpeedCPUAdam, 'Hybrid Optimizer Only Supports DeepSpeedCPUAdam' + self.backup_optimizer = torch.optim.AdamW([backup_gpu_param], + lr=self.optimizer.param_groups[0]["lr"], + betas=self.optimizer.param_groups[0]["betas"], + eps=self.optimizer.param_groups[0]["eps"], + weight_decay=self.optimizer.param_groups[0]["weight_decay"], + amsgrad=self.optimizer.param_groups[0]["amsgrad"]) + # Multiple param_groups configs for back-up optimizer + if len(self.optimizer.param_groups) > 1: + for i in range(1, len(self.optimizer.param_groups)): + self.backup_optimizer.add_param_group(self.optimizer.param_groups[i]) + + self.module = module + self.elastic_checkpoint = elastic_checkpoint + + self.device = get_accelerator().current_device_name() if not self.offload_optimizer else OffloadDeviceEnum.cpu + + self.inf_or_nan_tracker: Tensor = torch.zeros(1, dtype=torch.bool, device=self.device, requires_grad=False) + + self.deepspeed_adam_offload = (self.offload_optimizer and type(init_optimizer) == DeepSpeedCPUAdam) + + ### streams used for overlapping computation with communication + self.reduce_and_partition_stream = None if get_accelerator().is_synchronized_device() else get_accelerator( + ).Stream() if overlap_comm else get_accelerator().default_stream() + + ############################################################################ + + self.n_caching_allocator_flushes = 0 + + #-------------Stage 3 Setup-------------------# + + self.timers = timers + + self.all2all_process_group = all2all_process_group + + self.reduce_scatter = reduce_scatter + + self.dp_process_group = self.parameter_offload.dp_process_group + self.sequence_parallel_size = groups._get_sequence_parallel_world_size() + + self.all2all_process_group = all2all_process_group + + self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights + + self.partition_count = dist.get_world_size(group=self.dp_process_group) + + self.zeropp_loco_param = zeropp_loco_param + + if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): + self.model_parallel_group = None + self.model_parallel_rank = 0 + else: + self.model_parallel_group = mpu.get_model_parallel_group() + self.model_parallel_rank = mpu.get_model_parallel_rank() + + self.overflow = False + self.clip_grad = clip_grad + self.communication_data_type = communication_data_type + self.gradient_predivide_factor = gradient_predivide_factor + self.postscale_gradients = postscale_gradients + self.gradient_accumulation_steps = gradient_accumulation_steps + self.micro_step_id = 0 + self.reduce_bucket_size = int(reduce_bucket_size) + + if self.all2all_process_group is not None: + assert self.all2all_process_group is not None and self.reduce_scatter == True, "when enable all_to_all_reduce, reduce_scatter should also be enabled for data type checks." + + if self.reduce_scatter: + valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) + assert self.communication_data_type in valid_reduce_scatter_dtypes, f"ZeRO-3 supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" + assert self.gradient_predivide_factor == 1.0, "gradient_predivide_factor != 1.0 is not yet supported with ZeRO-3 with reduce scatter enabled" + assert self.postscale_gradients, "pre-scale gradients is not yet supported with ZeRO-3 with reduce scatter enabled" + + # Holds the mode parameter + # The param.data may not hold any meaningful data + # when param's status is NOT_AVAILABLE or IN_FLGHT + self.fp16_groups = [] + + # Hold partitioned parameters + self.fp16_partitioned_groups = [] + + # Holds a fused and flattened copy of the parameters + self.fp16_partitioned_groups_flat = [] + self.fp16_partitioned_groups_flat_numel = [] + self.fp16_partitioned_groups_flat_id = [] + + #defragmented pinned memory + self.param_groups_fp16_flat_cpu_memory = [] + + #a single 32-bit partition of the parallel partitioned parameters + #that this process will update + self.fp32_partitioned_groups_flat = [] + self.next_swappable_fp32_partitioned_groups = [] + + # number of elements per partition in each group + self.partition_size = [] + + self.all_reduce_print = False + + self.prefetch_elements = int(prefetch_bucket_size) + + self.contiguous_gradients = contiguous_gradients + + # padding on each partition for alignment purposes + self.groups_padding = [] + + self.sub_group_size = sub_group_size + + self.sub_group_to_group_id = {} + + # Trainable parameters + self.trainable_param_groups = self._get_trainable_parameter_groups() + + see_memory_usage("Before creating fp16 partitions", force=True) + self._create_fp16_partitions_with_defragmentation(self.trainable_param_groups) + num_fp16_subgroups = len(self.fp16_partitioned_groups_flat) + see_memory_usage(f"After creating fp16 partitions: {num_fp16_subgroups}", force=True) + + # Optimizer tensor swapping + if self.swap_optimizer: + self._configure_tensor_swapping(offload_optimizer_config, aio_config) + + self.is_gradient_accumulation_boundary: bool = True + + self.param_reduce_events: Deque[get_accelerator().Event] = collections.deque() + # TODO. make this configurable via JSON + self.max_param_reduce_events: int = 2 + + self.param_dict = {} + + # map between param_id and bool to specify if a param is in this partition + self.is_param_in_current_partition = {} + + if is_autocast_initialized(): + comm_dtypes = get_all_autocast_dtypes([p for params in self.fp16_groups for p in params]) + self.torch_autocast_gradscaler = torch.amp.GradScaler(device=get_accelerator().device_name()) + else: + comm_dtypes = {self.communication_data_type} + self.torch_autocast_gradscaler = None + + self.ipg_buckets: Dict[torch.dtype, IPGBucketZ3] = {dtype: IPGBucketZ3() for dtype in comm_dtypes} + + self.params_already_reduced = {} + self.previous_reduced_grads = None + + # model parameter traversal-based param id that's stable across runs + for params_group in self.fp16_groups: + for param in params_group: + param_id = self.get_param_id(param) + self.param_dict[param_id] = param + self.params_already_reduced[param_id] = False + + #Largest partitioned param + largest_partitioned_param_numel = 0 + for fp16_partitioned_group in self.fp16_partitioned_groups: + if len(fp16_partitioned_group) > 0: + largest_partitioned_param_numel = max( + largest_partitioned_param_numel, + max([max(tensor.numel(), tensor.ds_numel) for tensor in fp16_partitioned_group])) + + print_rank_0(f'Largest partitioned param numel = {largest_partitioned_param_numel}', force=False) + + self._setup_for_real_optimizer() + self.grad_position = {} + self.set_grad_positions() + + if self.offload_optimizer: + self.norm_for_param_grads = {} + + # stores if a partition has been reduced in this step + self.is_partition_reduced = {} + + # stores if a grad in a partition has been computed or not + self.is_grad_computed = {} + + # will store the averaged gradients required by this partition + self.averaged_gradients = {} + + #creates backward hooks for gradient partitioning + ###Calls all gather param + self._grad_acc_hooks = [] + self._leaf_module_hooks = [] + self.create_reduce_and_remove_grad_hooks() + + #exit(0) + + # we may have a way of fusing dynamic scale. Do not support for now + self.loss_scaler = CreateLossScaler(dtype=self.dtype, + static_loss_scale=static_loss_scale, + dynamic_scaling=dynamic_loss_scale, + dynamic_loss_args=dynamic_loss_args) + self.dynamic_loss_scale = self.loss_scaler.dynamic + + self.debug_fp16_grads = [{} for _ in self.fp16_groups] + + self._link_all_hp_params() + + self.offloaded_states: Set[OffloadDeviceEnum] = set() + + if dist.get_rank(group=self.dp_process_group) == 0: + see_memory_usage(f"After initializing ZeRO optimizer", force=True) + + def destroy(self): + self.parameter_offload.destroy() + for hook in self._grad_acc_hooks: + hook.remove() + for hook in self._leaf_module_hooks: + hook.remove() + print_rank_0("Removed grad acc hooks", force=False) + self.ipg_buckets.clear() + + def initialize_ds_offload( + self, + module, + timers, + ds_config, + overlap_comm, + prefetch_bucket_size, + max_reuse_distance, + max_live_parameters, + param_persistence_threshold, + model_persistence_threshold, + dp_process_group, + offload_param_config, + mpu, + zero_param_parallel_group, + zero_quantized_weights, + zero_quantized_nontrainable_weights, + zero_module_granularity_threshold, + log_trace_cache_warnings, + ): + return DeepSpeedZeRoOffload(module=module, + timers=timers, + ds_config=ds_config, + overlap_comm=overlap_comm, + prefetch_bucket_size=prefetch_bucket_size, + max_reuse_distance=max_reuse_distance, + max_live_parameters=max_live_parameters, + param_persistence_threshold=param_persistence_threshold, + model_persistence_threshold=model_persistence_threshold, + dp_process_group=dp_process_group, + offload_param_config=offload_param_config, + mpu=mpu, + zero_param_parallel_group=zero_param_parallel_group, + zero_quantized_weights=zero_quantized_weights, + zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights, + zero_module_granularity_threshold=zero_module_granularity_threshold, + log_trace_cache_warnings=log_trace_cache_warnings) + + def _get_trainable_parameter_groups(self): + param_groups = [] + PARAMS_KEY = "params" + for param_group in self.optimizer.param_groups: + trainable_params = [p for p in param_group[PARAMS_KEY] if p.requires_grad] + if len(trainable_params) == 0: + continue + + trainable_param_group = {} + for key in param_group.keys(): + if key == PARAMS_KEY: + trainable_param_group[PARAMS_KEY] = trainable_params + else: + trainable_param_group[key] = param_group[key] + param_groups.append(trainable_param_group) + + return param_groups + + def _set_zero_group_parallelism(self): + groups._create_zero_param_parallel_group(self.zero_hpz_partition_size) + + def invalidate_secondary_tensor(self): + for fpg in self.fp16_groups: + for param in fpg: + if param.ds_secondary_tensor is not None: + param.ds_secondary_tensor = None + + def _setup_for_real_optimizer(self): + see_memory_usage("Before creating fp32 partitions", force=True) + self._create_fp32_partitions() + see_memory_usage("After creating fp32 partitions", force=True) + dist.barrier() + + # To support pipelined optimizer swapping + self._create_next_swappable_fp32_groups() + + see_memory_usage("Before initializing optimizer states", force=True) + + self.initialize_optimizer_states() + see_memory_usage("After initializing optimizer states", force=True) + dist.barrier() + + if dist.get_rank() == 0: + logger.info(f"optimizer state initialized") + + # IPG + if self.contiguous_gradients: + for dtype, bucket in self.ipg_buckets.items(): + bucket.buffer = torch.empty(self.reduce_bucket_size, + dtype=dtype, + device=get_accelerator().current_device_name()) + + self.grad_partitions_flat_buffer = None + self.__param_id_to_grad_partition: Dict[int, Tensor] = {} + + all_params = list(itertools.chain.from_iterable(self.fp16_groups)) + + self.grad_partitions_flat_buffer: Tensor = torch.zeros(sum(p.partition_numel() for p in all_params), + dtype=self.gradient_accumulation_dtype, + device=self.device) + if self.offload_optimizer_pin_memory: + self.grad_partitions_flat_buffer = get_accelerator().pin_memory(self.grad_partitions_flat_buffer) + + offset = 0 + for param in all_params: + self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( + 0, offset, param.partition_numel()) + offset += param.partition_numel() + + def _link_all_hp_params(self): + for p in self.module.parameters(): + p._z3_optimizer = self + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + # TODO. factor out to a utility outside of stage3 + @staticmethod + def defragment(tensors: List[Tensor]) -> Tensor: + """move provided tensors into a contiguous flat buffer, with some additional + measures taken to reduce memory fragmentation""" + assert len(set(t.dtype for t in tensors)) == 1 + assert len(set(t.device for t in tensors)) == 1 + + cpu_buffer = torch.empty(sum(p.numel() for p in tensors), + dtype=get_only_unique_item(t.dtype for t in tensors), + device="cpu") + tensor_infos: List[Tuple[Tensor, int, int]] = get_mapping_to_flat_buffer(tensors) + orig_device = get_only_unique_item(t.device for t in tensors) + + offset = 0 + for tensor, offset, tensor_numel in tensor_infos: + # move the tensor from device memory to host memory + cpu_buffer.narrow(0, offset, tensor_numel).copy_(tensor) + tensor.data = torch.empty(0, dtype=tensor.dtype, device=tensor.device) + + gc.collect() + get_accelerator().empty_cache() + + # copy tensors (now flattened and contiguous) back to GPU + device_buffer = cpu_buffer.to(orig_device) + + # restore device tensors + for tensor, offset, tensor_numel in tensor_infos: + tensor.data = device_buffer.narrow(0, offset, tensor_numel) + + return device_buffer + + def _get_param_coordinator(self): + return self.parameter_offload.get_param_coordinator() + + def _configure_offloading(self, offload_optimizer_config, offload_param_config): + ###################### offload optimizer setup ################################## + if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: + self.offload_optimizer = True + self.offload_optimizer_pin_memory = offload_optimizer_config.pin_memory + self.swap_optimizer = offload_optimizer_config.device == OffloadDeviceEnum.nvme + self.offload_optimizer_fast_init = offload_optimizer_config.fast_init + + ###################### offload param setup ################################## + if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none: + self.offload_param = True + self.offload_param_pin_memory = offload_param_config.pin_memory + self.params_in_nvme_and_cpu = offload_param_config.device == OffloadDeviceEnum.nvme + self.max_params_in_cpu = offload_param_config.max_in_cpu + print_rank_0( + f"FP16 params swapping is {self.params_in_nvme_and_cpu}, Max params in CPU is {self.max_params_in_cpu}", + force=False) + + def _configure_tensor_swapping(self, offload_optimizer_config, aio_config): + nvme_swap_folder = os.path.join(offload_optimizer_config.nvme_path, 'zero_stage_3') + os.makedirs(nvme_swap_folder, exist_ok=True) + if dist.get_rank() == 0: + logger.info(f'Tensor Swapping: Adding optimizer tensors') + + swapper_type = PipelinedOptimizerSwapper if offload_optimizer_config.pipeline else PartitionedOptimizerSwapper + + self.optimizer_swapper = swapper_type(swap_config=offload_optimizer_config, + aio_config=aio_config, + base_folder=nvme_swap_folder, + optimizer=self.optimizer, + largest_numel=max(self.fp16_partitioned_groups_flat_numel), + device=self.device, + dtype=torch.float32, + timers=self.timers) + + def _move_to_flat_buffer(self, param_list, flat_buffer, avoid_copy=False): + '''If flat buffer is None then the parameters in the param_list are + not copied to the flat buffer. This is because they exceed the number of max_params_in_cpu + Some of these parameters may already be in CPU in unflattened buffers + or they maybe in GPU, or they maybe in NVME. If they are in NVME, then + they will be marked as NOT_AVAILABLE, and will be moved to CPU when they are + needed during training.''' + if flat_buffer is None: + # this dst buffer is on NVMe, so skip this + return + + start = 0 + for param in param_list: + src = param.ds_tensor + dest = flat_buffer.narrow(0, start, src.ds_numel) + start = start + src.ds_numel + '''if the parameter was initialized in nvme then bring it to the destination buffer directly''' + if src.status == PartitionedParamStatus.NOT_AVAILABLE: + print_rank_0( + f"Swapping in {param.ds_id} with partition size {param.partition_numel()} permanently to CPU") + param.nvme_swapper.swap_into_buffer(param, dest) + src.data = dest.data + src.status = PartitionedParamStatus.AVAILABLE + else: + assert src.status == PartitionedParamStatus.AVAILABLE, "Partitioned Param must be available here" + if not avoid_copy: + dest.data.copy_(src.data) + src.data = dest.data + + # Final location must be gpu/cpu in this case + param.ds_tensor.final_location = 'not-nvme' + + def _create_param_groups_fp16_flat_cpu_memory(self): + + aggregate_params_count = 0 + + for j, param_group in enumerate(self.trainable_param_groups): + params_in_group = sum([p.partition_numel() for p in param_group['params']]) + + flat_buffer_size = params_in_group + + if self.params_in_nvme_and_cpu and \ + aggregate_params_count + params_in_group > self.max_params_in_cpu: + + flat_buffer_size = max(0, self.max_params_in_cpu - aggregate_params_count) + + aggregate_params_count += params_in_group + + if flat_buffer_size > 0: + print_rank_0(f"group {j} flat buffer size {flat_buffer_size}", force=False) + self.param_groups_fp16_flat_cpu_memory.append(get_accelerator().pin_memory( + torch.empty(int(flat_buffer_size), dtype=self.dtype))) + else: + print_rank_0(f"No flat buffer size. Param group size was {params_in_group}", force=False) + + self.param_groups_fp16_flat_cpu_memory.append(torch.empty(1, dtype=self.dtype)) + + def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): + dist.barrier() + + param_groups: List[List[Parameter]] = tuple( + self._create_fp16_sub_groups(param_group["params"]) for param_group in fp16_param_groups) + + # bookkeeping related to param groups + for param_group_idx, param_group in enumerate(param_groups): + for sub_group in param_group: + sub_group_idx = len(self.fp16_groups) + + # record sub group and partitions + self.fp16_groups.append(sub_group) + self.fp16_partitioned_groups.append([param.ds_tensor for param in sub_group]) + + # record sub group -> group mapping + self.sub_group_to_group_id[sub_group_idx] = param_group_idx + + # record total elements of parameter partitions in sub group + self.fp16_partitioned_groups_flat_numel.append(sum(p.partition_numel() for p in sub_group)) + + # record ds_ids of parameter partitions in sub group + self.fp16_partitioned_groups_flat_id.append([p.ds_id for p in sub_group]) + + # record padding required to align group to world size (only applies to last rank) + rank_requires_padding = dist.get_rank( + self.dp_process_group) == dist.get_world_size(self.dp_process_group) - 1 + self.groups_padding.append([p.padding_size() if rank_requires_padding else 0 for p in sub_group]) + + # move parameters to flattened buffer + if not self.offload_param: # partitioned params remain in GPU during training + # move parameter partitions into a single contiguous flat buffer + parameter_partitions = self._get_parameter_partitions() + + # We need to keep the reference to this buffer to make sure you can free it in `offload_states` + self.lp_param_buffer = __class__.defragment(parameter_partitions) + self._set_fp16_partitioned_groups_flat() + + else: # partitioned params offloaded to CPU when not in use + # create a flat CPU memory allocation for each param group + self._create_param_groups_fp16_flat_cpu_memory() + for param_group_idx, param_group in enumerate(param_groups): + flat_offset = 0 + for i, sub_group in enumerate(param_group): + total_elements = sum(p.partition_numel() for p in sub_group) + print_rank_0(f"Params in nvme and cpu {self.params_in_nvme_and_cpu}") + #Flat buffer may not be available for parameters that reside in NVME + if not self.params_in_nvme_and_cpu or flat_offset + total_elements <= self.param_groups_fp16_flat_cpu_memory[ + param_group_idx].numel(): + fp16_partitioned_group_flat = self.param_groups_fp16_flat_cpu_memory[param_group_idx].narrow( + 0, flat_offset, total_elements) + print_rank_0( + f"Creating a flat buffer for subgroup {i} requiring {total_elements} elements, and cumulative CPU elements {flat_offset + total_elements}", + force=False) + + elif self.params_in_nvme_and_cpu: + fp16_partitioned_group_flat = None + print_rank_0(f"No flat buffer for sub group {i} of {total_elements} elements", force=False) + else: + assert False, "Either params are in nvme, or they are in CPU memory. This code path should not be triggered. Please see you max_params_in_cpu and params_in_nvme configs" + + self.fp16_partitioned_groups_flat.append(fp16_partitioned_group_flat) + flat_offset += total_elements + + self._move_to_flat_buffer(sub_group, + fp16_partitioned_group_flat, + avoid_copy=not self.offload_param) + + # if necessary, create a pinned memory buffer to be used for swapping out + # params to NVME after optimizer step + should_create_fp16_flat_reuse_buffer = any(flattened_partition_group is None + for flattened_partition_group in self.fp16_partitioned_groups_flat) + if should_create_fp16_flat_reuse_buffer: + max_partition_numel, largest_partition_numel = 0, None + for sub_group in self.fp16_groups: + total_elements = sum(t.partition_numel() for t in sub_group) + if total_elements > max_partition_numel: + largest_partition_numel = [t.ds_numel for t in sub_group] + max_partition_numel = total_elements + + assert len(largest_partition_numel) > 0, f'Unexpected that largest partition is empty' + self.fp16_groups[0][0].nvme_swapper.reserve_partitioned_swap_space(largest_partition_numel) + + def _get_parameter_partitions(self) -> List[Tensor]: + return [param.ds_tensor for sub_group in self.fp16_groups for param in sub_group] + + def _swap_in_sub_group_to_flat_buffer(self, flat_buffer, sub_group_id): + offset = 0 + elements_in_sub_group = sum([t.ds_numel for t in self.fp16_partitioned_groups[sub_group_id]]) + assert (flat_buffer.numel() == elements_in_sub_group) + for param, partitioned_param in zip(self.fp16_groups[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]): + dest = flat_buffer.narrow(0, offset, partitioned_param.ds_numel) + if partitioned_param.status == PartitionedParamStatus.NOT_AVAILABLE: + print_rank_0( + f"Swapping in {param.ds_id} with elements {param.ds_numel} and partition {param.partition_numel()}" + ) + param.nvme_swapper.swap_in([param], async_op=False) + dest.data.copy_(partitioned_param.data) + param.nvme_swapper.remove_partition_and_release_buffers([param]) + print_rank_0(f"Swapping in {param.ds_id} done") + else: + dest.data.copy_(partitioned_param.data) + offset += partitioned_param.ds_numel + + def _create_next_swappable_fp32_groups(self): + reverse_order_indices = [i for i in range(len(self.fp32_partitioned_groups_flat))] + reverse_order_indices.reverse() + + next_group = None + for i in reverse_order_indices: + self.next_swappable_fp32_partitioned_groups.append(next_group) + if self._swappable_optimizer_subgroup(i): + next_group = self.fp32_partitioned_groups_flat[i] + + self.next_swappable_fp32_partitioned_groups.reverse() + + def _get_sub_group_partitions(self, sub_group_id): + sub_group_partitions = [] + for param, partitioned_param in zip(self.fp16_groups[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]): + if partitioned_param.status == PartitionedParamStatus.NOT_AVAILABLE: + swap_path = param.nvme_swapper.get_path(param, True) + sub_group_partitions.append((partitioned_param, param.partition_numel(), swap_path)) + else: + sub_group_partitions.append((partitioned_param, partitioned_param.ds_numel, None)) + + return sub_group_partitions + + def _create_fp32_partitions(self): + cpu_memory_usage = 0 + cpu_memory_sub_groups = 0 + nvme_memory_usage = 0 + num_swappable_partitions = 0 + num_swap_from_nvme_partitions = 0 + num_swap_from_cpu_partitions = 0 + swap_from_nvme_memory_usage = 0 + swap_from_cpu_memory_usage = 0 + GIGA_BYTES = (1024**3) + + swappable_fp32_tensors = [] + swappable_fp16_src_tensors = [] + nvme_fp16_partitions_info = [] + nvme_fp16_num_elems = [] + nvme_fp32_dest_tensors = [] + fp32_element_size = torch.tensor([], dtype=torch.float32).element_size() + + # Assign portion of subgroup to cpu, the other to gpu. + if self.offload_optimizer: + self.subgroup_to_device = {} + sub_group_size = len(self.fp16_partitioned_groups_flat) + # print(f"Partial offload sub_group_size is {sub_group_size}, ratio is {self.partial_offload}\n") + for i in range(sub_group_size): + if i < int(self.partial_offload * sub_group_size): + self.subgroup_to_device[i] = 'cpu' + else: + self.subgroup_to_device[i] = get_accelerator()._name + + for i, tensor in enumerate(self.fp16_partitioned_groups_flat): + num_elements = self.fp16_partitioned_groups_flat_numel[i] + ds_id_begin = str(self.fp16_partitioned_groups_flat_id[i][0]) + ds_id_end = str(self.fp16_partitioned_groups_flat_id[i][-1]) + ds_id = ds_id_begin + '_' + ds_id_end + + # a partition of the fp32 master weights that will be updated by this process + if self._swappable_optimizer_subgroup(i): + self.fp32_partitioned_groups_flat.append(torch.Tensor()) + self.fp32_partitioned_groups_flat[i].ds_id = ds_id + nvme_memory_usage += (fp32_element_size * num_elements) + num_swappable_partitions += 1 + + if self.params_in_nvme_and_cpu and tensor is None: + num_swap_from_nvme_partitions += 1 + swap_from_nvme_memory_usage += (fp32_element_size * num_elements) + if self.offload_optimizer_fast_init: + sub_group_partitions = self._get_sub_group_partitions(i) + nvme_fp16_partitions_info.append(sub_group_partitions) + nvme_fp16_num_elems.append(num_elements) + nvme_fp32_dest_tensors.append(self.fp32_partitioned_groups_flat[i]) + else: + unpinned_fp32_buffer = torch.empty(num_elements, device=self.device, dtype=torch.float) + self._swap_in_sub_group_to_flat_buffer(unpinned_fp32_buffer, i) + self.optimizer_swapper.initialize_parameters(parameters=[self.fp32_partitioned_groups_flat[i]], + src_tensors=[unpinned_fp32_buffer]) + else: + num_swap_from_cpu_partitions += 1 + swap_from_cpu_memory_usage += (fp32_element_size * num_elements) + swappable_fp32_tensors.append(self.fp32_partitioned_groups_flat[i]) + swappable_fp16_src_tensors.append(self.fp16_partitioned_groups_flat[i]) + else: + cpu_memory_usage += (fp32_element_size * num_elements) + cpu_memory_sub_groups += 1 + + if self.params_in_nvme_and_cpu and tensor is None: + unpinned_fp32_buffer = torch.empty(num_elements, device=self.device, dtype=torch.float) + self._swap_in_sub_group_to_flat_buffer(unpinned_fp32_buffer, i) + self.fp32_partitioned_groups_flat.append(unpinned_fp32_buffer) + else: + if self.offload_optimizer: + self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i].to( + self.subgroup_to_device[i]).clone().float().detach()) + else: + self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i].to( + self.device).clone().float().detach()) + self.fp32_partitioned_groups_flat[i].ds_id = ds_id + + self.fp32_partitioned_groups_flat[i].requires_grad = True # keep this in case internal optimizer uses it + + if len(swappable_fp32_tensors) > 0: + self.optimizer_swapper.initialize_parameters(parameters=swappable_fp32_tensors, + src_tensors=swappable_fp16_src_tensors) + + if len(nvme_fp32_dest_tensors) > 0: + fp16_pinned_buffers = self.fp16_groups[0][0].nvme_swapper.reserve_available_buffers() + assert len(fp16_pinned_buffers) > 0 + self.optimizer_swapper.initialize_from_swapped_fp16_params(fp16_partitions_info=nvme_fp16_partitions_info, + fp16_num_elems=nvme_fp16_num_elems, + fp16_pinned_buffers=fp16_pinned_buffers, + fp32_parameters=nvme_fp32_dest_tensors) + self.fp16_groups[0][0].nvme_swapper.release_reserved_buffers() + + nvme_gigabytes = nvme_memory_usage / GIGA_BYTES + print_rank_0(f'Swappable FP32 Partitions: count={num_swappable_partitions} size={nvme_gigabytes:5.2f} GB', + force=False) + if self.params_in_nvme_and_cpu: + print_rank_0( + f'Swap from NVMe Partitions: count = {num_swap_from_nvme_partitions}, size = {swap_from_nvme_memory_usage/GIGA_BYTES:5.2f}GB', + force=False) + print_rank_0( + f'Swap from CPU Partitions: count = {num_swap_from_cpu_partitions}, size = {swap_from_cpu_memory_usage/GIGA_BYTES:5.2f}GB', + force=False) + + cpu_memory_gigabytes = cpu_memory_usage / GIGA_BYTES + print_rank_0(f'In-Memory FP32 Partitions: count={cpu_memory_sub_groups} size={cpu_memory_gigabytes:5.2f} GB', + force=False) + + # Clear for on-the-fly population before the optimizer step + for param_group in self.optimizer.param_groups: + param_group['params'] = [] + + def _create_fp16_sub_groups(self, params_group): + + params_group_numel = sum([param.partition_numel() for param in params_group]) + sub_group_size = self.sub_group_size + + if sub_group_size is None or sub_group_size >= params_group_numel: + return [params_group] + + sub_groups = [] + sub_group = [] + local_sub_group_size = 0 + for param in params_group: + + sub_group.append(param) + local_sub_group_size += param.partition_numel() + + if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): + + sub_groups.append(sub_group) + + sub_group = [] + local_sub_group_size = 0 + + return sub_groups + + def _optimizer_step(self, sub_group_id): + param_group_id = self.sub_group_to_group_id[sub_group_id] + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + if self.offload_optimizer: + cur_device = self.subgroup_to_device[sub_group_id] + if cur_device == 'cpu': + self.optimizer.param_groups[param_group_id]['params'] = [fp32_param] + cpu_loss = self.optimizer.step() + self.optimizer.param_groups[param_group_id]['params'] = [] + else: + self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] + gpu_loss = self.backup_optimizer.step() + self.backup_optimizer.param_groups[param_group_id]['params'] = [] + else: + self.optimizer.param_groups[param_group_id]['params'] = [fp32_param] + self.optimizer.step() + self.optimizer.param_groups[param_group_id]['params'] = [] + + def _swappable_optimizer_subgroup(self, sub_group_id): + if not self.swap_optimizer: + return False + + return self.optimizer_swapper.is_swappable_tensor(None, + numel=self.fp16_partitioned_groups_flat_numel[sub_group_id]) + + def _partitioned_params_swap_out(self, i): + offset = 0 + fp32_param = self.fp32_partitioned_groups_flat[i] + assert fp32_param is not None, \ + f'fp32 parameters of sub_group {i} is None' + + swap_fp16_params = [] + swap_fp32_params = [] + for param, partitioned_param in zip(self.fp16_groups[i], self.fp16_partitioned_groups[i]): + src = fp32_param.narrow(0, offset, partitioned_param.ds_numel) + if partitioned_param.status == PartitionedParamStatus.AVAILABLE: + partitioned_param.data.copy_(src.data) + else: + swap_fp32_params.append(src) + swap_fp16_params.append(param) + offset += partitioned_param.ds_numel + + if len(swap_fp16_params): + swap_fp16_params[0].nvme_swapper.swap_out_partitioned_params(dst_fp16_params=swap_fp16_params, + src_fp32_params=swap_fp32_params) + + def _set_fp16_partitioned_groups_flat(self): + # setup flat buffers per subgroup, these are each just sections of the + # contiguous flat buffer for all parameters that we created earlier + offset = 0 + for sub_group in self.fp16_groups: + sub_group_numel = sum(param.partition_numel() for param in sub_group) + self.fp16_partitioned_groups_flat.append(self.lp_param_buffer.narrow(0, offset, sub_group_numel)) + offset += sub_group_numel + + def initialize_optimizer_states(self): + num_subgroups = len(self.fp16_groups) + + largest_numel = max([sum([p.ds_numel for p in psg]) for psg in self.fp16_partitioned_groups]) + gradient_dtype = self.fp32_partitioned_groups_flat[0].dtype + gradient_buffer = torch.zeros(int(largest_numel), dtype=gradient_dtype, device=self.device) + + timer_names = set() + + # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers + # which do lazy initialization of the state at the first call to step. + is_adagrad = isinstance(self.optimizer, torch.optim.Adagrad) + + if self.swap_optimizer: + self.optimizer_swapper.init_timers() + + timer_names.add(INIT_OPTIMIZER_TIMER) + self.timers(INIT_OPTIMIZER_TIMER).start() + + for i, group in enumerate(self.fp16_groups): + swappable_optimizer_subgroup = self._swappable_optimizer_subgroup(i) + swappable_param_subgroup = self.fp16_partitioned_groups_flat[i] is None + + num_elements = int(self.fp16_partitioned_groups_flat_numel[i]) + + see_memory_usage( + f'[Begin] Initialize optimizer states {i} / {num_subgroups} subgroups, num_elems: {num_elements}, swappable opt/param:{swappable_optimizer_subgroup}/{swappable_param_subgroup}', + force=False) + + if swappable_optimizer_subgroup: + self._optimizer_states_and_gradient_swap_in(i, timer_names) + + if self.offload_optimizer and not swappable_optimizer_subgroup: + subgroup_gradient_buffer = torch.zeros(num_elements, dtype=gradient_dtype, device=self.device) + if self.offload_optimizer_pin_memory: + subgroup_gradient_buffer = get_accelerator().pin_memory(subgroup_gradient_buffer) + + self.fp32_partitioned_groups_flat[i].grad = subgroup_gradient_buffer.to(self.subgroup_to_device[i]) + else: + self.fp32_partitioned_groups_flat[i].grad = gradient_buffer.narrow(0, 0, num_elements) + + if swappable_param_subgroup: + self._partitioned_params_swap_out(i) + + if swappable_optimizer_subgroup: + self._optimizer_states_and_gradient_swap_out(i, timer_names) + + see_memory_usage( + f'[End] Initialize optimizer states {i} / {num_subgroups} subgroups, num_elems: {num_elements}, swappable opt/param:{swappable_optimizer_subgroup}/{swappable_param_subgroup}', + force=False) + + # Initialize the optimizer states with the flattened fp32 partition. + if is_adagrad: + self.optimizer = torch.optim.Adagrad(self.fp32_partitioned_groups_flat, **self.optimizer.defaults) + + self.timers(INIT_OPTIMIZER_TIMER).stop() + self.timers.log(timer_names) + + if self.swap_optimizer: + self.optimizer_swapper.log_timers() + + if not self.offload_optimizer: + for group in self.fp32_partitioned_groups_flat: + group.grad = None + + # Reset steps + return + + ######################################################################### + #########################ZeRO Partition Gradients######################## + ######################################################################### + + def get_first_param_index(self, group_id, param_group, partition_id): + for index, param in enumerate(param_group): + param_id = self.get_param_id(param) + if partition_id in self.param_to_partition_ids[group_id][param_id]: + return index + return None + + def initialize_gradient_partitioning_data_structures(self): + + total_partitions = dist.get_world_size(group=self.dp_process_group) + + for i, param_group in enumerate(self.fp16_groups): + + self.param_to_partition_ids[i] = {} + self.is_partition_reduced[i] = {} + self.total_grads_in_partition[i] = {} + self.remaining_grads_in_partition[i] = {} + self.is_grad_computed[i] = {} + self.grad_partition_insertion_offset[i] = {} + self.grad_start_offset[i] = {} + self.first_param_index_in_partition[i] = {} + + for partition_id in range(total_partitions): + self.is_grad_computed[i][partition_id] = {} + self.grad_partition_insertion_offset[i][partition_id] = {} + self.grad_start_offset[i][partition_id] = {} + self.initialize_gradient_partition(i, param_group, partition_id) + self.is_partition_reduced[i][partition_id] = False + self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( + i, param_group, partition_id) + + @instrument_w_nvtx + def independent_gradient_partition_epilogue(self): + self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) + for comm_dtype in sort_dtypes(self.ipg_buckets.keys()): + self.__reduce_and_partition_ipg_grads(comm_dtype) + self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) + + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + for param_id in self.params_already_reduced.keys(): + self.params_already_reduced[param_id] = False + + #in case of cpu offload, averaged gradients are already in fp32_partitioned_groups_flat.grad + #TODO: use a similar code path for both cpu_offload and non-cpu offload + if not self.offload_optimizer: + for i, sub_group in enumerate(self.fp16_groups): + #TODO: This is redundant + self.averaged_gradients[i] = [ + self.__param_id_to_grad_partition[param.ds_id] + if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group + ] + # this method gets called after every backward. need to increment + # here because if it gets incremented in backward() the micro step + # id will be off by one when we do the reduce and partition at the. + # start of this method. + # TODO. make this less error prone + self.micro_step_id += 1 + + def overlapping_partition_gradients_reduce_epilogue(self): + self.independent_gradient_partition_epilogue() + + def create_reduce_and_remove_grad_hooks(self): + print_rank_0(f'[Begin] Create gradient reduction hooks') + self.leaf_parameters = defaultdict(list) + for i, param_group in enumerate(self.fp16_groups): + for param in param_group: + if param.requires_grad: + #print_rank_0(f" Before all gather {param.device}, {param.shape}") + print_rank_0(f"Before all gather {param.device}, {param.shape}", force=False) + + # The hook must be created in un-partitioned parameter + param.all_gather() + + #print(f"After all gather {param.device}, {param.shape}") + def wrapper(param): + + @instrument_w_nvtx + def reduce_partition_and_remove_grads(*notneeded): + self.reduce_ready_partitions_and_remove_grads(param) + + self._grad_acc_hooks.append(register_grad_hook(param, reduce_partition_and_remove_grads)) + + #print(f"param grad fn {param.expand_as(param).grad_fn}") + if z3_leaf_parameter(param): + self.leaf_parameters[param.ds_z3_leaf_module].append(param) + else: + wrapper(param) + + # Partition the parameter after creating the hook + param.partition() + + # We delay reduce-scatter for all gradients in the leaf modules until the backward pass of the leaf module is done + for leaf_module, leaf_parameters in self.leaf_parameters.items(): + + def wrapper_pre_hook(params): + + def forward_pre_hook(module, input): + """Pre-forward hook to set backward hook on input tensors to the leaf module""" + module._leaf_module_inputs_remaining = 0 + + @instrument_w_nvtx + def reduce_leaf_module_grads(grad): + module._leaf_module_inputs_remaining -= 1 + # Make sure everything is done in the leaf module + if module._leaf_module_inputs_remaining == 0: + for param in params: + if param.grad is None: + param.grad = torch.zeros_like(param) + self.reduce_ready_partitions_and_remove_grads(param) + + def set_module_bwd_hook(tensor): + if tensor.requires_grad: + module._leaf_module_inputs_remaining += 1 + tensor.register_hook(reduce_leaf_module_grads) + return tensor + + output = apply_to_tensors_only(set_module_bwd_hook, input) + + return output + + return forward_pre_hook + + def wrapper_post_hook(): + + def forward_post_hook(module, input, output): + """Pre-forward hook to set backward hook on input tensors to the leaf module""" + module._leaf_output_required_grad_num = 0 + + def increment_rg_count_bwd_hook(tensor): + if tensor.requires_grad: + module._leaf_output_required_grad_num += 1 + return tensor + + apply_to_tensors_only(increment_rg_count_bwd_hook, output) + + if module._leaf_module_inputs_remaining == 0 and module._leaf_output_required_grad_num > 0: + raise RuntimeError( + "A module cannot be set as a leaf module when it does not have any input tensors that require gradients and has output tensors that require gradients. This is because the gradient reduction hook will not be called in this case." + ) + + return forward_post_hook + + self._leaf_module_hooks.append(leaf_module.register_forward_pre_hook(wrapper_pre_hook(leaf_parameters))) + self._leaf_module_hooks.append(leaf_module.register_forward_hook(wrapper_post_hook())) + + print_rank_0(f'[End] Create gradient reduction hooks') + + def get_param_id(self, param): + return OptimizerSwapper.parameter_id(param) + + ###############Independent Partition Gradient ######################## + def reduce_independent_p_g_buckets_and_remove_grads(self, param): + #print_rank_0(f"Inside reduce ipg buckets. {debug_param2name_id_shape(param)}, ipg elements {self.elements_in_ipg_bucket}, reduce bucket size {self.reduce_bucket_size}", force=True) + + # Because the ipg bucket is initialized with a random place holder tensor, we must + # explicitly check that the bucket has any real data in it (self.elements_in_ipg_bucket > + # 0). Otherwise if the incoming param.ds_numel is large, this branch may get triggered on a + # garbage data and `self.average_tensor()` will crash because its params_to_reduce will be + # empty, while reduction_list will have that garbage data. + comm_dtype = self.get_param_comm_dtype(param) + bucket = self.ipg_buckets[comm_dtype] + if bucket.elements + param.ds_numel > self.reduce_bucket_size and bucket.elements > 0: + self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.ds_numel) + self.__reduce_and_partition_ipg_grads(comm_dtype) + + # deal with a use-case of transient grads that will be generated in a loop for the same computation involving some model params - e.g. when performing a tiled memory calculation that shards the normal single sub-module call into a loop over a shards. + if getattr(param, "ds_grad_is_ready", True): + self.__add_grad_to_ipg_bucket(param) + + @instrument_w_nvtx + @torch.no_grad() + def __add_grad_to_ipg_bucket(self, param: Parameter) -> None: + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.wait_stream(get_accelerator().default_stream()) + + bucket = self.ipg_buckets[self.get_param_comm_dtype(param)] + if self.contiguous_gradients and bucket.elements + param.grad.numel() <= self.reduce_bucket_size: + # move the gradient to a contiguous buffer + with get_accelerator().stream(self.reduce_and_partition_stream): + # move the parameter's gradient to the contiguous flat buffer + new_grad_tensor = bucket.buffer.narrow(0, bucket.elements, param.grad.numel()).view_as(param.grad) + new_grad_tensor.copy_(param.grad, non_blocking=True) + if not get_accelerator().is_synchronized_device(): + param.grad.record_stream(get_accelerator().current_stream()) + param.grad.data = new_grad_tensor + + bucket.params.append(param) + bucket.elements += param.grad.numel() + + @instrument_w_nvtx + @torch.no_grad() + def __reduce_and_partition_ipg_grads(self, communication_data_type: torch.dtype, safe_mode: bool = False) -> None: + bucket = self.ipg_buckets[communication_data_type] + params_in_bucket = bucket.params + + if not params_in_bucket: + return + + for param in params_in_bucket: + if param.grad.numel() != param.ds_numel: + raise RuntimeError(f"{param.grad.numel()} != {param.ds_numel} Cannot reduce scatter " + f"gradients whose size is not same as the params") + + assert len(set(p.ds_id for p in params_in_bucket)) == len(params_in_bucket) + + while self.param_reduce_events and self.param_reduce_events[0].query(): + self.param_reduce_events.popleft() + if len(self.param_reduce_events) > self.max_param_reduce_events: + self.param_reduce_events.popleft().synchronize() + + with get_accelerator().stream(self.reduce_and_partition_stream): + if safe_mode: + assert_ints_same_as_other_ranks([p.ds_id for p in params_in_bucket]) + + if self.contiguous_gradients and bucket.elements <= self.reduce_bucket_size and not self.reduce_scatter: + grad_bucket = bucket.buffer.narrow(0, 0, bucket.elements) + grad_partitions = self.__avg_scatter_contiguous_grads(grad_bucket, communication_data_type) + else: + params_in_bucket.sort(key=lambda p: p.ds_id) + grad_partitions = self.__avg_scatter_grads(params_in_bucket, communication_data_type) + + self.partition_grads(params_in_bucket, grad_partitions) + + params_in_bucket.clear() + + if not get_accelerator().handles_memory_backpressure(): + event = get_accelerator().Event() + event.record() + self.param_reduce_events.append(event) + + @instrument_w_nvtx + def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor, + communication_data_type: torch.dtype) -> List[Tensor]: + dtype = buffer_to_reduce.dtype + if communication_data_type != dtype: + buffer_to_reduce = buffer_to_reduce.to(communication_data_type) + if self.postscale_gradients and self.gradient_predivide_factor != 1.0: + buffer_to_reduce = buffer_to_reduce.div_(self.gradient_predivide_factor) + + world_sz = dist.get_world_size(self.dp_process_group) + rank = dist.get_rank(self.dp_process_group) + buffer_to_reduce.div_(world_sz / float(self.sequence_parallel_size)) + + dist.all_reduce(buffer_to_reduce, group=self.dp_process_group) + + if self.postscale_gradients and self.gradient_predivide_factor != world_sz: + buffer_to_reduce = buffer_to_reduce.mul(self.gradient_predivide_factor) + + if communication_data_type != self.dtype: + buffer_to_reduce = buffer_to_reduce.to(self.dtype) + + grad_partitions = [] + grad_offset_in_buffer = 0 + for param in self.ipg_buckets[communication_data_type].params: + grad = param.grad + chunk_sz = math.ceil(grad.numel() / world_sz) + + start_offset = grad_offset_in_buffer + min(rank * chunk_sz, grad.numel()) + end_offset = grad_offset_in_buffer + min(rank * chunk_sz + chunk_sz, grad.numel()) + + partition = buffer_to_reduce[start_offset:end_offset] + if param.partition_numel() != partition.numel(): + padded_partition = torch.zeros(param.partition_numel(), device=grad.device, dtype=grad.dtype) + if partition.numel() > 0: + padded_partition[:partition.numel()] = partition + grad_partitions.append(padded_partition) + else: + grad_partitions.append(partition) + grad_offset_in_buffer += grad.numel() + + return grad_partitions + + @instrument_w_nvtx + def __avg_scatter_grads(self, params_to_reduce: List[Parameter], + communication_data_type: torch.dtype) -> List[Tensor]: + """average gradients and scatter partitions across ranks""" + + full_grads_for_rank = [p.grad for p in params_to_reduce] + if communication_data_type != self.dtype: + full_grads_for_rank = [g.to(communication_data_type) for g in full_grads_for_rank] + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0: + full_grads_for_rank = [g.div(self.gradient_predivide_factor) for g in full_grads_for_rank] + + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + if self.all2all_process_group is not None and num_nodes > 1: + grad_partitions_for_rank = (all_to_all_loco_quant_reduce(params_to_reduce, self.all2all_process_group, + self.zeropp_loco_param) + if self.zeropp_loco_param is not None else all_to_all_quant_reduce( + full_grads_for_rank, self.all2all_process_group)) + else: + grad_partitions_for_rank = reduce_scatter_coalesced(full_grads_for_rank, self.dp_process_group) + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0 and self.gradient_predivide_factor != dist.get_world_size( + self.dp_process_group): + grad_partitions_for_rank = [g.mul(self.gradient_predivide_factor) for g in grad_partitions_for_rank] + + if communication_data_type != self.dtype: + grad_partitions_for_rank = [g.to(self.dtype) for g in grad_partitions_for_rank] + + return grad_partitions_for_rank + + def set_grad_positions(self): + for i, group in enumerate(self.fp16_groups): + current_offset = 0 + for param in group: + param_id = self.get_param_id(param) + num_elements = param.partition_numel() + + self.grad_position[param_id] = [int(i), int(current_offset), int(num_elements)] + #print(f"param id {param_id} i:{i}, ds_tensor {num_elements} numel {param.numel()}") + current_offset += num_elements + see_memory_usage(f"After Set Grad positions", force=False) + + def _constant_buffered_norm2(self, input, buffer_size=250000000): + norm = None + for part in input.view(-1).split(buffer_size): + if norm is None: + norm = part.data.double().norm(2)**2.0 + else: + norm += part.data.double().norm(2)**2.0 + return norm**0.5 + + def set_norm_for_param_grad_in_gpu(self, param): + param_id = self.get_param_id(param) + #self.norm_for_param_grads[param_id] = param.grad.data.double().norm(2) + #Using a more memory efficient version + self.norm_for_param_grads[param_id] = self._constant_buffered_norm2(param.grad) + + def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param, fp32_grad_tensor): + with get_accelerator().stream(self.copy_grad_stream): + param_id = self.get_param_id(param) + src_tensor = param.grad.view(-1).float() + #print(f"src_tensor {src_tensor.size()} and fp32 grad {fp32_grad_tensor.size()}") + fp32_grad_tensor.copy_(src_tensor, non_blocking=True) + param.grad = None + + def complete_grad_norm_calculation_for_cpu_offload(self, params): + total_norm = 0.0 + norm_type = 2.0 + for p in params: + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + param_id = self.get_param_id(p) + if param_id in self.norm_for_param_grads.keys(): + param_norm = self.norm_for_param_grads[param_id] + total_norm += param_norm**2 + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) + + total_norm = total_norm_cuda[0]**(1. / norm_type) + + mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device) + + return total_norm.cpu() + + @instrument_w_nvtx + def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: + offload_fp32_gradients = {} + offload_fp32_offsets = {} + buffers = [] + for param, grad_partition in zip(params_to_release, grad_partitions): + + contains_real_data = param.partition_numel() * dist.get_rank(self.dp_process_group) < param.ds_numel + if not contains_real_data: + # this grad partition is empty - don't need to do anything + param.grad = None + continue + + # move or accumulate gradient partition to target buffer + grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel()) + buffers.append(grad_buffer) + if self.micro_step_id == 0: # don't accumulate + grad_buffer.copy_(grad_partition, non_blocking=True) + # ensure grad buffer is a CUDA buffer to speed up the next few + # operations and so it can be used asynchronously + grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) + elif get_accelerator().on_accelerator(grad_buffer): + grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(grad_buffer.shape)) + else: + # if dst is CPU, copy first to src device, do the addition + # there, then move back to dst. adding directly to cpu is very slow + cuda_grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) + cuda_grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(cuda_grad_buffer.shape)) + grad_buffer.copy_(cuda_grad_buffer, non_blocking=True) + # ensure grad buffer is a CUDA buffer to speed up the next few + # operations and so it can be used asynchronously + grad_buffer = cuda_grad_buffer + + # offload the gradient partition if applicable + if self.offload_optimizer: + i, dest_offset, _ = self.grad_position[self.get_param_id(param)] + + if self.is_gradient_accumulation_boundary: + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) + + if self._swappable_optimizer_subgroup(i): + if not i in offload_fp32_gradients.keys(): + offload_fp32_gradients[i] = [] + offload_fp32_offsets[i] = [] + + offload_fp32_gradients[i].append(grad_buffer.float()) + offload_fp32_offsets[i].append(dest_offset) + else: + fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( + 0, dest_offset, grad_buffer.numel()) + fp32_grad_tensor.copy_(grad_buffer.float()) + + # free the gradient + if not get_accelerator().is_synchronized_device(): + if param.grad is not None: + param.grad.record_stream(get_accelerator().current_stream()) + param.grad = None + + if self.offload_optimizer and self.swap_optimizer: + for i in offload_fp32_gradients.keys(): + self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i], + gradient_offsets=offload_fp32_offsets[i], + gradient_tensors=offload_fp32_gradients[i]) + return buffers + + def reduce_ready_partitions_and_remove_grads(self, param): + #print_rank_0(f"Backward {debug_param2name_id_shape(param)}", force=True) + self.reduce_independent_p_g_buckets_and_remove_grads(param) + + def zero_reduced_gradients(self, partition_id, i): + + def are_all_related_partitions_reduced(params_id): + for partition_id in self.param_to_partition_ids[i][params_id]: + if not self.is_partition_reduced[i][partition_id]: + return False + return True + + for params_id in self.is_grad_computed[i][partition_id]: + if are_all_related_partitions_reduced(params_id): + self.param_dict[params_id].grad = None + + def quantize_nontrainable_params(self): + """ In ZeRO-3, when the zero_quantized_nontrainable_weights flag is set, we quantize the non-trainable weights and also store them in quantized format. However, this check for trainable/non-trainable is done when deepspeed initializes the partitioning. So, if the user changes the trainable/non-trainable status of a parameter after the partitioning is done (e.g. LoRA), the user needs to re-quantize the non-trainable weights by calling this function. + """ + if not self.zero_quantized_nontrainable_weights: + print_rank_0( + f"Warning: quantize_nontrainable_params() called with zero_quantized_nontrainable_weights disabled, return without doing anything", + force=True) + return + quantizer_module = CUDAQuantizer() + + def quantize_dstensor(tensor): + assert tensor.dtype == torch.float16, f"quantize_dstensor() expects tensor.dtype == torch.float16, got {tensor.dtype}" + partition_size = tensor.ds_numel + ds_status = tensor.status + final_location = tensor.final_location + tensor, tensor.ds_quant_scale = quantizer_module.quantize(tensor) + tensor.ds_numel = partition_size + tensor.status = ds_status + tensor.final_location = final_location + tensor.requires_grad = False + return tensor + + for param in self.module.parameters(): + if hasattr(param, "ds_tensor") and (param.ds_tensor.numel() <= 2048 or param.ds_numel <= 500000): + # skip small parameters + continue + if hasattr(param, + "ds_tensor") and not param.requires_grad and not hasattr(param.ds_tensor, "ds_quant_scale"): + param.ds_tensor = quantize_dstensor(param.ds_tensor) + if hasattr(param, "ds_secondary_tensor") and not param.requires_grad and not hasattr( + param.ds_secondary_tensor, "ds_quant_scale") and param.ds_secondary_tensor is not None: + param.ds_secondary_tensor = quantize_dstensor(param.ds_secondary_tensor) + get_accelerator().synchronize() + + def flatten_and_print(self, message, tensors, start=0, n=5): + flatten_tensor = self.flatten(tensors) + + def print_func(): + logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) + + self.sequential_execution(print_func, message) + + def get_grads_to_reduce(self, i, partition_id): + + def get_reducible_portion(key): + grad = self.param_dict[key].grad + total_elements = grad.numel() + start = self.grad_start_offset[i][partition_id][key] + num_elements = min(total_elements - start, + self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) + if not pg_correctness_test: + if num_elements == total_elements: + return grad + else: + return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) + else: + if num_elements == total_elements: + return grad.clone() + else: + return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) + + grads_to_reduce = [] + for key in self.is_grad_computed[i][partition_id]: + grad = get_reducible_portion(key) + grads_to_reduce.append(grad) + return grads_to_reduce + + def sequential_execution(self, function, message, group=None): + if group is None: + group = self.dp_process_group + if dist.get_rank(group=group) == 0: + logger.info(message) + for id in range(dist.get_world_size(group=group)): + if id == dist.get_rank(group=group): + function() + dist.barrier(group=group) + + def set_none_gradients_to_zero(self, i, partition_id): + for param_id in self.is_grad_computed[i][partition_id]: + param = self.param_dict[param_id] + if param.grad is None: + param.grad = torch.zeros_like(param) + + ######################Reduction Related Methods############################## + + def allreduce_bucket(self, bucket, rank=None, log=None): + rank = None + tensor = self.flatten(bucket) + + tensor_to_allreduce = tensor + + if pg_correctness_test: + communication_data_type = torch.float32 + else: + communication_data_type = self.communication_data_type + + if communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(communication_data_type) + + tensor_to_allreduce.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) + + if rank is None: + # "All Reducing" + dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) + else: + global_rank = dist.get_global_rank(self.dp_process_group, rank) + dist.reduce(tensor_to_allreduce, global_rank, group=self.dp_process_group) + + if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + if rank is None or rank == dist.get_rank(group=self.dp_process_group): + tensor.copy_(tensor_to_allreduce) + + return tensor + + # if rank is specified do a reduction instead of an allreduce + def allreduce_and_copy(self, small_bucket, rank=None, log=None): + with get_accelerator().stream(self.reduction_stream): + allreduced = self.allreduce_bucket(small_bucket, rank=rank, log=log) + if rank is None or rank == dist.get_rank(group=self.dp_process_group): + for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + buf.copy_(synced) + + def allreduce_no_retain(self, bucket, numel_per_bucket=500000000, rank=None, log=None): + small_bucket = [] + numel = 0 + for tensor in bucket: + small_bucket.append(tensor) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy(small_bucket, rank=rank, log=None) + small_bucket = [] + if len(small_bucket) > 0: + self.allreduce_and_copy(small_bucket, rank=rank, log=log) + + ############################################################################# + ############################################################################# + ############################################################################# + + # views the tensor as multiple partitions and returns + # those partitions + def get_data_parallel_partitions(self, tensor): + partitions = [] + + dp = dist.get_world_size(group=self.dp_process_group) + dp_id = dist.get_rank(group=self.dp_process_group) + + total_num_elements = tensor.numel() + + base_size = total_num_elements // dp + remaining = total_num_elements % dp + + start = 0 + for id in range(dp): + partition_size = base_size + if id < remaining: + partition_size = partition_size + 1 + partitions.append(tensor.narrow(0, start, partition_size)) + start = start + partition_size + return partitions + + def get_partition_info(self, tensor_list, partition_size, partition_id): + params_in_partition = [] + params_not_in_partition = [] + + start_index = partition_size * partition_id + end_index = partition_size * (partition_id + 1) + + current_index = 0 + first_offset = 0 + + for tensor in tensor_list: + + tensor_size = tensor.numel() + + if start_index <= current_index < end_index: + params_in_partition.append(tensor) + + elif current_index < start_index < (current_index + tensor_size): + params_in_partition.append(tensor) + + assert (first_offset == 0 + ), "This can happen either zero or only once as this must be the first tensor in the partition" + first_offset = start_index - current_index + + else: + params_not_in_partition.append(tensor) + + current_index = current_index + tensor_size + + return params_in_partition, params_not_in_partition, first_offset + + @instrument_w_nvtx + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + self.micro_step_id = 0 + + # FP32 grad should never exist. + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if set_to_none: + if p.grad is not None and get_accelerator().on_accelerator(p.grad): + p.grad.record_stream(get_accelerator().current_stream()) + p.grad = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def _model_parallel_all_reduce(self, tensor, op): + """ Perform all reduce within model parallel group, if any. + """ + if self.model_parallel_group is None: + pass + else: + dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) + + @instrument_w_nvtx + def get_grad_norm_direct(self, gradients, params, norm_type=2): + """Clips gradient norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float or int): max norm of the gradients + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + """ + norm_type = float(norm_type) + if norm_type == inf: + total_norm = max(g.data.abs().max() for g in gradients) + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + # Take max across all GPUs. + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX) + total_norm = total_norm_cuda[0] + else: + # if dist.get_rank() == 0: + # logger.info(f"Total Norm beginning {total_norm}") + grad_norms = [] + for g, p in zip(gradients, params): + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + grad_norms.append(g.to(get_accelerator().device_name(), non_blocking=True).double().norm(2)) + + # Sum across all model parallel GPUs. + if len(grad_norms) == 0: + # FIX https://github.com/deepspeedai/DeepSpeed/issues/3564 + total_norm_cuda = torch.tensor(0, + dtype=gradients[0].dtype).to(get_accelerator().device_name()).double() + else: + total_norm_cuda = torch.sum(torch.pow(torch.stack(grad_norms), 2)) + + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) + + total_norm = total_norm_cuda**(1. / norm_type) + + norm_is_inf = total_norm.isinf() + norm_is_nan = total_norm.isnan() + inf_or_nan = norm_is_nan.logical_or(norm_is_inf) + + err = torch.tensor(-1.0, device=self.device, dtype=torch.float) + total_norm = torch.where(inf_or_nan, err, total_norm) + + return total_norm + + # creates a flat fused tensor from the tensor list starting at the first_offset + # in the first tensor of the list. If there are not enough elements in the tensor + # list then the flat tensor will be padded with zeros + def get_flat_partition(self, tensor_list, first_offset, partition_size, return_tensor_list=False): + flat_tensor_list = [] + current_size = 0 + for i, tensor in enumerate(tensor_list): + if tensor.grad is None: + tensor.grad = torch.zeros_like(tensor) + + tensor = tensor.grad + num_elements = tensor.numel() + tensor_offset = 0 + + # we need to offset to get to the right element + if i == 0 and first_offset > 0: + tensor_offset = first_offset + num_elements = num_elements - tensor_offset + + # we dont need all elements of the tensor + if num_elements > (partition_size - current_size): + num_elements = partition_size - current_size + + # we need a narrow view of the tensor based on the tensor offset and number of elements that + # we need from this tensor + if tensor_offset > 0 or num_elements < tensor.numel(): + flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) + else: + flat_tensor_list.append(tensor) + + current_size = current_size + num_elements + + # this means its the last partition and does not align with the dp boundary. We need to pad before flattening + if current_size < partition_size: + flat_tensor_list.append( + torch.zeros(int(partition_size - current_size), + dtype=tensor_list[0].dtype, + device=tensor_list[0].device)) + + if return_tensor_list: + return flat_tensor_list + + return self.flatten(flat_tensor_list) + + def free_grad_in_param_list(self, param_list): + for p in param_list: + p.grad = None + + def reset_cpu_buffers(self): + self.norm_for_param_grads = {} + + def _pre_step(self): + self.micro_step_id = 0 + + print_rank_0(f"Inside Step function") + see_memory_usage(f"In step before checking overflow", force=False) + + print_rank_0("Finished Tracing at Beginning of Step") + self._get_param_coordinator().hierarchy = 0 + + print_rank_0("Finished Tracing at Beginning of Step") + + @instrument_w_nvtx + def _get_norm_groups(self): + norm_groups = [] + for i, group in enumerate(self.fp16_groups): + if self.offload_optimizer: + norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(self.fp16_groups[i])) + else: + norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.fp16_groups[i])) + return norm_groups + + @instrument_w_nvtx + def _prepare_fp32_grad_for_sub_group(self, sub_group_id): + partition_id = dist.get_rank(group=self.dp_process_group) + + single_grad_partition = self.flatten(self.averaged_gradients[sub_group_id]).to( + self.fp32_partitioned_groups_flat[sub_group_id].dtype) + + assert single_grad_partition.numel() == self.fp32_partitioned_groups_flat[sub_group_id].numel(), \ + "averaged gradients have different number of elements that partition size {} {} {} {}".format( + single_grad_partition.numel(), self.fp32_partitioned_groups_flat[sub_group_id].numel(), sub_group_id, partition_id) + + self.fp32_partitioned_groups_flat[sub_group_id].grad = single_grad_partition + + # release all the gradient since we have already created a necessary copy in dp_grad_partition + self.zero_grad(set_to_none=True) + + if not get_accelerator().is_synchronized_device(): + for grad in filter(lambda g: get_accelerator().on_accelerator(g), self.averaged_gradients[sub_group_id]): + grad.record_stream(get_accelerator().current_stream()) + + self.averaged_gradients[sub_group_id] = None + + @instrument_w_nvtx + def _prepare_sub_group(self, sub_group_id, timer_names): + see_memory_usage(f'Before prepare optimizer sub group {sub_group_id}', force=False) + if self._swappable_optimizer_subgroup(sub_group_id): + self._optimizer_states_and_gradient_swap_in(sub_group_id, timer_names) + elif not self.offload_optimizer: + self._prepare_fp32_grad_for_sub_group(sub_group_id) + see_memory_usage(f'After prepare optimizer sub group {sub_group_id}', force=False) + + def _optimizer_states_and_gradient_swap_in(self, sub_group_id, timer_names=None): + param_length = self.fp16_partitioned_groups_flat_numel[sub_group_id] + fp32_param_id = self.get_param_id(self.fp32_partitioned_groups_flat[sub_group_id]) + assert self._swappable_optimizer_subgroup(sub_group_id), \ + f'Parameter {fp32_param_id} of numel={param_length} is not swappable' + + see_memory_usage(f'pre-step Before swapping in optimizer tensors {sub_group_id}', force=False) + if timer_names is not None: + timer_names.add(OPTIMIZER_SWAP_IN_STATE_TIMER) + self.timers(OPTIMIZER_SWAP_IN_STATE_TIMER).start() + + self.optimizer_swapper.swap_in_optimizer_state( + parameter=self.fp32_partitioned_groups_flat[sub_group_id], + async_parameter=self.next_swappable_fp32_partitioned_groups[sub_group_id]) + + if timer_names is not None: + self.timers(OPTIMIZER_SWAP_IN_STATE_TIMER).stop() + see_memory_usage(f'pre-step After swapping in optimizer tensors {sub_group_id}', force=False) + + @instrument_w_nvtx + def _release_sub_group(self, sub_group_id, timer_names): + see_memory_usage(f'Before release optimizer sub group {sub_group_id}', force=False) + # get rid of the fp32 gradients. Not needed anymore + if not self.offload_optimizer: + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + if self._swappable_optimizer_subgroup(sub_group_id): + self._optimizer_states_and_gradient_swap_out(sub_group_id, timer_names) + see_memory_usage(f'After release optimizer sub group {sub_group_id}', force=False) + + # create a flat tensor aligned at the alignment boundary + @instrument_w_nvtx + def flatten_dense_tensors_aligned(self, tensor_list, alignment): + num_elements = 0 + for tens in tensor_list: + num_elements = num_elements + tens.numel() + + remaining = num_elements % alignment + + if remaining: + elements_to_add = alignment - remaining + pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) + padded_tensor_list = tensor_list + [pad_tensor] + + num_elements = num_elements + elements_to_add + else: + padded_tensor_list = tensor_list + + return self.flatten(padded_tensor_list) + + def _optimizer_states_and_gradient_swap_out(self, sub_group_id, timer_names=None): + param_length = self.fp16_partitioned_groups_flat_numel[sub_group_id] + fp32_param_id = self.get_param_id(self.fp32_partitioned_groups_flat[sub_group_id]) + assert self._swappable_optimizer_subgroup(sub_group_id), \ + f'Parameter {fp32_param_id} of numel={param_length} is not swappable' + + see_memory_usage(f'post-step Before swapping out optimizer tensors {sub_group_id}', force=False) + if timer_names is not None: + timer_names.add(OPTIMIZER_SWAP_OUT_STATE_TIMER) + self.timers(OPTIMIZER_SWAP_OUT_STATE_TIMER).start() + + self.optimizer_swapper.swap_out_optimizer_state( + parameter=self.fp32_partitioned_groups_flat[sub_group_id], + async_swap=self.next_swappable_fp32_partitioned_groups[sub_group_id] is not None) + + if timer_names is not None: + self.timers(OPTIMIZER_SWAP_OUT_STATE_TIMER).stop() + see_memory_usage(f'post-step After swapping out optimizer tensors {sub_group_id}', force=False) + + # get rid of the fp32 gradients. Not needed anymore + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + def _release_swap_buffers(self, sub_group_id): + self.optimizer_swapper.release_swap_buffers(parameter=self.fp32_partitioned_groups_flat[sub_group_id]) + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + def _writeback_swap_state(self, sub_group_id, write_opt_state, write_gradients): + self.optimizer_swapper.writeback_optimizer_state_and_gradients(self.fp32_partitioned_groups_flat[sub_group_id], + write_opt_state, write_gradients) + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + def _unflatten_partitioned_parameters(self, sub_group_id): + updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]) + + for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + partitioned_param.data = q.data + + def _overflow_clean_up(self, prev_scale): + see_memory_usage('After overflow before clearing gradients', force=False) + self.zero_grad(set_to_none=True) + + if self.offload_optimizer: + self.reset_cpu_buffers() + else: + self.averaged_gradients = {} + + see_memory_usage('After overflow after clearing gradients', force=False) + + def _loco_err_buf_update(self, overflow: bool, scale=1.0): + """ + Loco Error Buffer update. + """ + if not overflow and scale == 1.0: return + if dist.get_rank() == 0: + logger.info(f"update loco-zero++ error buffer with overflow: {overflow}") + # FP32 grad should never exist. + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if hasattr(p, 'intra_ef_buf'): + if overflow: + del p.intra_ef_buf + del p.inter_ef_buf + else: + p.intra_ef_buf[1] *= scale + p.inter_ef_buf[1] *= scale + + @instrument_w_nvtx + def _overflow_check_and_loss_scale_update(self): + + # First compute norm for all group so we know if there is overflow + if self.dtype == torch.float16: + self.check_overflow() + + #loss scaling related computation + prev_scale = self.loss_scale + self._update_scale(self.overflow) + + if self.overflow: + self._overflow_clean_up(prev_scale) + + #update loco error buf + self._loco_err_buf_update(self.overflow, self.loss_scale / prev_scale) + + return self.overflow + + @instrument_w_nvtx + def _post_step(self, timer_names): + if self.offload_optimizer: + self.reset_cpu_buffers() + + #Gathering persisting parameters + if len(self.persistent_parameters) > 0: + self.persistent_parameters[0].all_gather(self.persistent_parameters) + + if self.swap_optimizer: + self.optimizer_swapper.log_timers() + + self.invalidate_secondary_tensor() + + self.timers.log(timer_names) + + see_memory_usage('After zero_optimizer step', force=False) + print_rank_0(f"------------------Finishing Step-----------------------") + + @instrument_w_nvtx + def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): + if self.fp16_partitioned_groups_flat[sub_group_id] is not None: + self.fp16_partitioned_groups_flat[sub_group_id].data.copy_( + self.fp32_partitioned_groups_flat[sub_group_id].data) + + #unflatten fp16 parameter subgroup + self._unflatten_partitioned_parameters(sub_group_id) + else: + self._partitioned_params_swap_out(sub_group_id) + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + @instrument_w_nvtx + def step(self, closure=None): + """ + Not supporting closure. + """ + self._pre_step() + self._partition_all_parameters() + + #checks for overflow, adjust the loss scale accordingly + if self._overflow_check_and_loss_scale_update(): + if self.swap_optimizer: + self.optimizer_swapper.log_timers() + return + + norm_groups = self._get_norm_groups() + scaled_global_grad_norm = torch.linalg.vector_norm(torch.stack(norm_groups)) + + # Stash unscaled gradient norm + self._global_grad_norm = scaled_global_grad_norm / self.loss_scale + + timer_names = set() + + timer_names.add(OPTIMIZER_STEP_TIMER) + self.timers(OPTIMIZER_STEP_TIMER).start() + + #update parameters one sub group at a time + for sub_group_id, group in enumerate(self.fp16_groups): + + #prepare optimizer states, gradients and fp32 parameters for update + self._prepare_sub_group(sub_group_id, timer_names) + + #scale the fp32 gradients + self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) + + #apply the optimizer step on the sub group and copy fp32 parameters to fp16 + self._optimizer_step(sub_group_id) + + #put fp16 parameters in appropriate location + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + + #release memory or swap out optimizer states of fp32 parameters + self._release_sub_group(sub_group_id, timer_names) + + self.timers(OPTIMIZER_STEP_TIMER).stop() + + self._post_step(timer_names) + + # warn user about caching allocator flushes + memory_stats = get_accelerator().memory_stats() + alloc_retries = memory_stats.get("num_alloc_retries") + if alloc_retries is None: + alloc_retries = 0 + if alloc_retries > self.n_caching_allocator_flushes: + if dist.get_rank() == 0: + logger.warning( + "%d pytorch allocator cache flushes since last step. this happens " + "when there is high memory pressure and is detrimental to " + "performance. if this is happening frequently consider adjusting " + "settings to reduce memory consumption. If you are unable to " + "make the cache flushes go away consider adding " + "get_accelerator().empty_cache() calls in your training loop to ensure " + "that all ranks flush their caches at the same time", + alloc_retries - self.n_caching_allocator_flushes) + self.n_caching_allocator_flushes = alloc_retries + + def dump_pre_step_gradients(self, debug_fp32_grads): + # Dump gradient norms for debugging + for i, _ in enumerate(self.fp16_groups): + print(f'Pre-Step Dump Norms for Group {i} FP16P, FP16G, FP32G, FP32GUC') + for fp16_param, fp32_grad in zip(self.fp16_groups[i], debug_fp32_grads[i]): + param_id = self.get_param_id(fp16_param) + fp16_grad_norm = self.debug_fp16_grads[i][param_id] + + fp32_grad_norm = [float(t.data.float().norm(2)) for t in fp32_grad] + norm_list = [fp16_grad_norm, fp32_grad_norm] + print(f'Pre-Step Norms {i} {param_id} = {norm_list}') + + def dump_post_step_gradients(self): + # Dump gradient norms for debugging + for i, group in enumerate(self.fp16_groups): + print(f'Post-Step Dump Norms for Group {i} FP16P, FP16DS, FP16FLAT, FP32FLAT') + unflat_fp16 = self.unflatten(self.fp16_groups_flat[i], self.fp16_groups[i]) + unflat_fp32 = self.unflatten(self.fp32_partitioned_groups_flat[i], self.fp16_groups[i]) + for j, p in enumerate(self.fp16_groups[i]): + param_id = self.get_param_id(p) + param_norm = float(p.data.float().norm(2)) + ds_norm = float(p.ds_tensor.data.float().norm(2)) + + unflat_norm = [float(t.data.float().norm(2)) for t in [unflat_fp16[j], unflat_fp32[j]]] + norm_list = [param_norm, ds_norm] + unflat_norm + print(f'Post-Step Norms {i} {param_id} = {norm_list}') + + @instrument_w_nvtx + def unscale_and_clip_grads(self, sub_group_id, total_norm): + # compute combined scale factor for this group + combined_scale = self.loss_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad + clip = torch.clamp(clip, min=1.0) + combined_scale = clip * self.loss_scale + + self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(1. / combined_scale) + + def _check_overflow(self, partition_gradients=True): + self.overflow = self.has_overflow(partition_gradients) + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params, is_grad_list=False): + for p in params: + if p.grad is not None and self._has_inf_or_nan(p.grad.data): + return True + + return False + + def has_overflow_partitioned_grads_serial(self): + for i in range(len(self.fp16_groups)): + for j, grad in enumerate(self.averaged_gradients[i]): + if grad is not None and self._has_inf_or_nan(grad.data, j): + return True + return False + + @instrument_w_nvtx + def has_overflow(self, partition_gradients=True): + if partition_gradients: + with get_accelerator().stream(self.reduce_and_partition_stream): + if hasattr(self.inf_or_nan_tracker, "logical_or_"): + self.inf_or_nan_tracker.logical_or_(torch.isinf(self.grad_partitions_flat_buffer).any()) + self.inf_or_nan_tracker.logical_or_(torch.isnan(self.grad_partitions_flat_buffer).any()) + else: + # logical_or_ not available in older versions of pytorch + self.inf_or_nan_tracker += torch.isinf(self.grad_partitions_flat_buffer).any() + self.inf_or_nan_tracker += torch.isnan(self.grad_partitions_flat_buffer).any() + self.inf_or_nan_tracker = self.inf_or_nan_tracker > 0 + + overflow_gpu = self.inf_or_nan_tracker.clone().to(get_accelerator().current_device_name()).to( + torch.uint8) + self.inf_or_nan_tracker.zero_() + + if not get_accelerator().resolves_data_dependency(): + get_accelerator().default_stream().wait_stream(self.reduce_and_partition_stream) + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + else: + params = [] + for group in self.fp16_groups: + for param in group: + params.append(param) + + overflow = self.has_overflow_serial(params, is_grad_list=partition_gradients) + overflow_gpu = get_accelerator().ByteTensor([overflow]) + + # Since each model parallel GPU carries only part of the model, + # make sure overflow flag is synced across all the model parallel GPUs + self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) + + overflow = overflow_gpu[0].item() + return bool(overflow) + + # `x` is a torch.Tensor + @staticmethod + def _has_inf_or_nan(x, j=None): + try: + # if x is half, the .float() incurs an additional deep copy, but it's necessary if + # Pytorch's .sum() creates a one-element tensor of the same type as x + # (which is true for some recent version of pytorch). + cpu_sum = float(x.float().sum()) + # More efficient version that can be used if .sum() returns a Python scalar + # cpu_sum = float(x.sum()) + except RuntimeError as instance: + # We want to check if inst is actually an overflow exception. + # RuntimeError could come from a different error. + # If so, we still want the exception to propagate. + if "value cannot be converted" not in instance.args[0]: + raise + return True + else: + if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: + return True + return False + + @instrument_w_nvtx + def backward(self, loss, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.swap_optimizer: + self.optimizer_swapper.pre_backward() + + see_memory_usage(f"Before backward", force=False) + + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + else: + self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) + + if self.swap_optimizer: + self.optimizer_swapper.post_backward() + + def get_fp32_grad_partitions(self) -> Dict[int, Dict[int, Tensor]]: + """get fp32 gradient partition dictionary + accessed as grad_dict[parameter_group_index][parameter_index] + """ + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + grad_dict = collections.defaultdict(dict) + if self.offload_optimizer: + for group in self.fp16_groups: + for param_idx, param in enumerate(group): + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + grad_dict[group_idx][param_idx] = fp32_grad + else: + for group_idx, group in self.averaged_gradients.items(): + for param_idx, gradient in enumerate(group): + grad_dict[group_idx][param_idx] = gradient.float() + + return grad_dict + + def _fp32_state_allgather(self, param, fp32_state_partition): + reduce_buffer = torch.empty(self.partition_count * fp32_state_partition.numel(), + dtype=torch.float32, + device=param.device) + my_rank = dist.get_rank(group=self.dp_process_group) + partition = reduce_buffer.narrow(0, fp32_state_partition.numel() * my_rank, fp32_state_partition.numel()) + partition.data.copy_(fp32_state_partition.data, non_blocking=False) + dist.all_gather_into_tensor(reduce_buffer, partition, group=self.dp_process_group) + return reduce_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape) + + def _get_fp32_grad_state_partition(self, param, release_swap_buffers): + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + if self.offload_optimizer: + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_in(group_idx) + + fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + + if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: + self._release_swap_buffers(group_idx) + else: + fp32_grad = self.__param_id_to_grad_partition[param.ds_id] + + return fp32_grad, group_idx + + def get_fp32_grad_for_param(self, param) -> Tensor: + if not param.requires_grad: + return None + + fp32_grad, _ = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=True) + fp32_grad = fp32_grad.to(get_accelerator().current_device_name()).float() + return self._fp32_state_allgather(param, fp32_grad) + + def set_fp32_grad_for_param(self, value, param): + if not param.requires_grad: + return + + # if not get_accelerator().resolves_data_dependency(): + # self.reduce_and_partition_stream.synchronize() + + # if self.offload_optimizer: + # group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + # fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + # else: + # fp32_grad = self.__param_id_to_grad_partition[param.ds_id] + + fp32_grad, group_idx = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=False) + # import pdb; pdb.set_trace() + my_rank = dist.get_rank(group=self.dp_process_group) + value_partition = value.flatten().narrow(0, fp32_grad.numel() * my_rank, fp32_grad.numel()) + fp32_grad.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._writeback_swap_state(group_idx, write_opt_state=False, write_gradients=True) + + def _get_fp32_opt_state_partition(self, param, release_swap_buffers, optim_state_key=None): + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_in(group_idx) + + fp32_param = self.fp32_partitioned_groups_flat[group_idx] + if optim_state_key is None: + fp32_opt_state = fp32_param.narrow(0, dest_offset, num_elements) + else: + fp32_opt_state = self.optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) + + if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: + self._release_swap_buffers(group_idx) + + return fp32_opt_state, group_idx + + def get_full_hp_param(self, param, optim_state_key=None) -> Tensor: + if not param.requires_grad: + return None + + # import pdb; pdb.set_trace() + fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=True, + optim_state_key=optim_state_key) + fp32_opt_state = fp32_opt_state.to(get_accelerator().current_device_name()) + hp_param = self._fp32_state_allgather(param, fp32_opt_state) + + return hp_param + + def set_full_hp_param(self, value, param, optim_state_key=None): + if not param.requires_grad: + return + + assert value.numel( + ) == param.ds_numel, f" Number of elements do not match: {value.numel()} != {param.ds_numel}" + + fp32_opt_state_partition, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=False, + optim_state_key=optim_state_key) + # print(f'{dist.get_rank()=} {fp32_opt_state_partition.shape=} -------- {value.shape=}') + # import pdb; pdb.set_trace() + my_rank = dist.get_rank(group=self.dp_process_group) + value_partition = value.flatten().narrow(0, + fp32_opt_state_partition.numel() * my_rank, + fp32_opt_state_partition.numel()) + fp32_opt_state_partition.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_out(group_idx) + + ### Local API START ### + def get_local_fp32_grad_for_param(self, param) -> Tensor: + if not param.requires_grad: + return None + + fp32_grad, _ = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=True) + fp32_grad = fp32_grad.to(get_accelerator().current_device_name()).float() + return fp32_grad + + def set_local_grad_for_param(self, value, param): + if not param.requires_grad: + return + + assert value.numel() == param.ds_tensor.numel( + ), f" Number of elements do not match: {value.numel()} != {param.ds_tensor.ds_numel}" + + # if not get_accelerator().resolves_data_dependency(): + # self.reduce_and_partition_stream.synchronize() + + # if self.offload_optimizer: + # group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + # fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + # else: + # fp32_grad = self.__param_id_to_grad_partition[param.ds_id] + + if self.offload_optimizer: + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(value) + + fp32_grad, group_idx = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=False) + fp32_grad.data.copy_(value.flatten().data) + + if self._swappable_optimizer_subgroup(group_idx): + self._writeback_swap_state(group_idx, write_opt_state=False, write_gradients=True) + + def get_local_fp32_param(self, param, optim_state_key=None) -> Tensor: + if not param.requires_grad: + return None + fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=True, + optim_state_key=optim_state_key) + fp32_opt_state = fp32_opt_state.to(get_accelerator().current_device_name()) + return fp32_opt_state + + def set_local_hp_param(self, value, param, optim_state_key=None): + if not param.requires_grad: + return + + assert hasattr(param, "ds_tensor"), f" The parameter does not contain the partitioned copy of the tensor." + assert value.numel() == param.ds_tensor.numel( + ), f" Number of elements do not match: {value.numel()} != {param.ds_tensor.ds_numel}" + + fp32_opt_state_partition, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=False, + optim_state_key=optim_state_key) + value_partition = value.flatten() + fp32_opt_state_partition.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_out(group_idx) + # logger.info(f"[set_local_hp_param][update the params' value successfully]") + + ### Local API END ### + + ### Vectorized API BEGIN ### + def update_fp32_grad_for_param_vectorized(self, update_func, param_list): + params_with_grad = [p for p in param_list if p.requires_grad] + if not params_with_grad: + return + + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + subgroups = {} + for p in params_with_grad: + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(p)] + param_entry = (p, dest_offset, num_elements) + if group_idx in subgroups.keys(): + subgroups[group_idx].append(param_entry) + else: + subgroups[group_idx] = [param_entry] + + for group_idx, group_params in subgroups.items(): + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_in(group_idx) + + for param, dest_offset, num_elements in group_params: + if self.offload_optimizer: + fp32_grad_part = self.fp32_partitioned_groups_flat[group_idx].grad.narrow( + 0, dest_offset, num_elements) + else: + fp32_grad_part = self.__param_id_to_grad_partition[param.ds_id] + + fp32_grad_full = self._fp32_state_allgather(param, fp32_grad_part) + new_fp32_grad_full = update_func(fp32_grad_full, param) + my_rank = dist.get_rank(group=self.dp_process_group) + value_partition = new_fp32_grad_full.flatten().narrow(0, + fp32_grad_part.numel() * my_rank, + fp32_grad_part.numel()) + fp32_grad_part.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._writeback_swap_state(sub_group_id=group_idx, write_opt_state=False, write_gradients=True) + + ### Vectorized API END ### + + ### Device API BEGIN ### + def get_hp_param_device(self, param, optim_state_key=None) -> torch.device: + if not param.requires_grad: + return None + + fp32_opt_state, _ = self._get_fp32_opt_state_partition(param, + release_swap_buffers=True, + optim_state_key=optim_state_key) + return fp32_opt_state.device + + ### Device API END ### + + @instrument_w_nvtx + def _partition_all_parameters(self): + self.parameter_offload.partition_all_parameters() + + def check_overflow(self, partition_gradients=True): + self._check_overflow(partition_gradients) + + def _update_scale(self, has_overflow=False): + self.loss_scaler.update_scale(has_overflow) + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + self.trainable_param_groups = self._get_trainable_parameter_groups() + + param_groups = property(_get_param_groups, _set_param_groups) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.loss_scaler.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) + cur_scale = property(_get_loss_scale, _set_loss_scale) + + def _get_lean_tensors(self, padded_flattened_tensor, group_tensors, paddings): + # Remove paddings from flattened tensor + individual_tensors = self.unflatten(padded_flattened_tensor, group_tensors) + lean_lengths = [t.numel() - pad for t, pad in zip(group_tensors, paddings)] + lean_tensors = [t[:len] for t, len in zip(individual_tensors, lean_lengths)] + #logger.info(f'rank {dist.get_rank()}: lean_tensors = {[t.numel() for t in lean_tensors]}') + return lean_tensors + + #TODO REVISIT this for stage 3 + def get_lean_optimizer_state(self): + # Return optimizer states after removing paddings. + # This method assumes that each param group contains a single flattened tensor. + optimizer_groups_state = [] + + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + lean_state = {} + for key, value in self.optimizer.state[p].items(): + if torch.is_tensor(value): + padded_lens = [t.numel() for t in self.fp16_partitioned_groups[i]] + lean_state[key] = self._get_lean_tensors(value, self.fp16_partitioned_groups[i], + self.groups_padding[i]) + lean_flat_len = sum([t.numel() for t in lean_state[key]]) + else: + lean_state[key] = value + + optimizer_groups_state.append(lean_state) + + return optimizer_groups_state + + def get_groups_without_padding(self, groups_with_padding): + # Return group tensor after removing paddings added for alignment to DP world size. + groups_without_padding = [] + for i, group in enumerate(groups_with_padding): + lean_group = self._get_lean_tensors(group, self.fp16_partitioned_groups[i], self.groups_padding[i]) + groups_without_padding.append(lean_group) + + return groups_without_padding + + def _set_fp32_optimizer_param_groups(self): + for sub_group_id, _ in enumerate(self.fp16_groups): + param_group_id = self.sub_group_to_group_id[sub_group_id] + self.optimizer.param_groups[param_group_id]['params'].append( + self.fp32_partitioned_groups_flat[sub_group_id]) + + def _clear_fp32_optimizer_param_groups(self): + for param_group in self.optimizer.param_groups: + param_group['params'] = [] + + def _rigid_state_dict(self): + state_dict = {} + state_dict[ZERO_STAGE] = ZeroStageEnum.weights + state_dict[LOSS_SCALER] = self.loss_scaler + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['overflow'] = self.overflow + state_dict[PARTITION_COUNT] = self.partition_count + + self._set_fp32_optimizer_param_groups() + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat + self._clear_fp32_optimizer_param_groups() + + return state_dict + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + if self.elastic_checkpoint: + raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.") + + return self._rigid_state_dict() + + +# Restore base optimizer fp32 weights from checkpoint by: +# 1) Merging fp32 weights from checkpoints of all partitions +# 2) Extracting fp32 weights for current partition from merged weights +# 3) Using extracted weights to update base optimizer weights directly. + + def _restore_from_fp32_weights(self, all_state_dict): + + flat_local_partition = [] + for i in range(len(self.fp32_partitioned_groups_flat)): + merged_partitions = [sd['fp32_groups'][i] for sd in all_state_dict] + flat_local_partition.append(self._get_flattened_partition(merged_partitions)) + + for current, saved in zip(self.fp32_partitioned_groups_flat, flat_local_partition): + current.data.copy_(saved.data) + + # Restore base optimizer fp32 weights from ZeRO fp16 weights + def _restore_from_bit16_weights(self): + for fp16_partitions, fp32_partition in zip(self.fp16_partitioned_groups_flat, + self.fp32_partitioned_groups_flat): + fp32_partition.data.copy_(fp16_partitions.data) + + # Refresh the fp32 master params from the fp16 copies. + def refresh_fp32_params(self): + self._restore_from_bit16_weights() + + # Extract flattened partition for current rank from all partitions + def _get_flattened_partition(self, all_partition_states): + partition_id = dist.get_rank(group=self.dp_process_group) + alignment = dist.get_world_size(group=self.dp_process_group) + + param_partitions = [[] for _ in range(len(all_partition_states[0]))] + for i, partition in enumerate(all_partition_states): + for j, param in enumerate(partition): + param_partitions[j].append(param) + + local_state_partitions = [] + for param_index, param_slices in enumerate(param_partitions): + flattened_merged_tensor = self.flatten_dense_tensors_aligned(param_slices, alignment) + new_partitions = self.get_data_parallel_partitions(flattened_merged_tensor) + local_state_partitions.append(new_partitions[partition_id]) + + if torch.is_tensor(local_state_partitions[0]): + return self.flatten_dense_tensors_aligned(local_state_partitions, alignment) + + # Assume non-tensor states are not partitioned and equal across ranks, so return first one + return local_state_partitions[0] + + # Restore base optimizer state from checkpoint by + # 1) Merging optimizer state from checkpoints of all partitions + # 2) Extracting optimizer state for current partition from the merged state + # 3) Using the extracted value to directly update the base optimizer. + def _restore_base_optimizer_state(self, all_state_dict): + base_optimizer_group_states = [] + for i in range(len(self.optimizer.param_groups)): + partition_states = {} + all_partition_group_states = [sd['base_optimizer_state'][i] for sd in all_state_dict] + for key in all_partition_group_states[0].keys(): + all_partition_states = [all_states[key] for all_states in all_partition_group_states] + partition_states[key] = self._get_flattened_partition(all_partition_states) + base_optimizer_group_states.append(partition_states) + + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + for key, saved in base_optimizer_group_states[i].items(): + if torch.is_tensor(self.optimizer.state[p][key]): + self.optimizer.state[p][key].data.copy_(saved.data) + else: + self.optimizer.state[p][key] = saved + + def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): + # I think it should actually be ok to reload the optimizer before the model. + self.loss_scaler = state_dict[LOSS_SCALER] + self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] + self.overflow = state_dict['overflow'] + + if load_optimizer_states: + self._set_fp32_optimizer_param_groups() + self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) + self._clear_fp32_optimizer_param_groups() + + if self.swap_optimizer: + # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint + self.optimizer_swapper.purge_state() + + if self.swap_optimizer: + # Touch all parameters to synchronize all buffers + timer_names = set() + self._partition_all_parameters() + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) + self._post_step(timer_names) + + # restore fp32 partitions + for curr_param, saved_param in zip(self.fp32_partitioned_groups_flat, state_dict[FP32_FLAT_GROUPS]): + curr_param.data.copy_(saved_param.data) + + # restore fp16 partitions from fp32 + for sub_group_id in range(len(self.fp32_partitioned_groups_flat)): + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + if sum(fp32_param.size()) > 0: + fp16_param = self.fp16_partitioned_groups_flat[sub_group_id] + fp16_param.data.copy_(fp32_param.data) + + # update fp16 unflattened params + for sub_group_id in range(len(self.fp16_partitioned_groups_flat)): + updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]) + + for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + partitioned_param.data = q.data + + # TODO: Support different/changing load/save DP degree. + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None, + param_shapes=None): + r"""Loading a ZeRO checkpoint + Arguments: + state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. + Note that the number of saved partitions may differ from number of loading partitions to support + changing GPU count, specifically DP world size, between saving and loading checkpoints. + load_optimizer_states: Boolean indicating whether or not to load base optimizer states + load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 + copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). + """ + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + + if self.elastic_checkpoint: + raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.") + + if checkpoint_folder: + self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights, + param_shapes) + else: + self._rigid_load_state_dict(state_dict_list[dist.get_rank(group=self.dp_process_group)], + load_optimizer_states=load_optimizer_states) + + # when use loading checkpoint serial, after finish loading, we need to + # delete the temp state_dict_list variable to save memory, then trigger + # the next rank's loading + if load_serial is not None: + load_serial += 1 + rank = dist.get_rank(group=self.dp_process_group) + local_rank = dist.get_local_rank() + del state_dict_list[rank] + rank_end = dist.get_world_size() - 1 + if local_rank != rank_end: + dist.send(tensor=load_serial, dst=rank + 1) + + if len(self.persistent_parameters) > 0: + self.persistent_parameters[0].partition(self.persistent_parameters) + # self.persistent_parameters[0].all_gather(self.persistent_parameters) # this will be done in checkpoint_event_epilogue() so remove it to prevent double all_gather + + def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights, + param_shapes): + self.load_hp_checkpoint_state_from_checkpoint_dir_stage3(checkpoint_folder, param_shapes) + + def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir, param_shapes): + """ Load optimizer and model states from the checkpoint directory. """ + checkpoint_dir = os.path.join(checkpoint_dir, "zero") + optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") + assert os.path.isfile( + optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' + + optim_sd = torch.load(optim_state_path, weights_only=False) + self._load_global_state_stage3(optim_sd) + + key_list = ["fp32", "exp_avg", "exp_avg_sq"] + + for key in key_list: + key_tensor = torch.empty(0) + for layer in param_shapes[0].keys(): + key_layer_state_partition = self.load_hp_checkpoint_state(os.path.join(checkpoint_dir, layer), key) + key_tensor = torch.cat((key_tensor, key_layer_state_partition)) + if key == "fp32": + self.fp32_partitioned_groups_flat[0].data.copy_(key_tensor) + self.optimizer.param_groups[0]['params'].append(self.fp32_partitioned_groups_flat[0]) + else: + optim_sd[OPTIMIZER_STATE_DICT]['state'][0][key] = key_tensor + + if self.swap_optimizer: + # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint + self.optimizer_swapper.purge_state() + + if self.swap_optimizer: + # Touch all parameters to synchronize all buffers + timer_names = set() + self._partition_all_parameters() + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) + self._post_step(timer_names) + + self.optimizer.load_state_dict(optim_sd[OPTIMIZER_STATE_DICT]) + for param_group in self.optimizer.param_groups: + param_group['params'] = [] + + for sub_group_id in range(len(self.fp32_partitioned_groups_flat)): + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + if sum(fp32_param.size()) > 0: + fp16_param = self.fp16_partitioned_groups_flat[sub_group_id] + fp16_param.data.copy_(fp32_param.data) + + for sub_group_id in range(len(self.fp16_partitioned_groups_flat)): + updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]) + + for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + partitioned_param.data = q.data + + def _load_global_state_stage3(self, sd): + self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) + self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) + self.overflow = sd.get('overflow', self.overflow) + + def load_hp_checkpoint_state(self, folder, key): + local_rank = dist.get_local_rank() + + # Load tensors from files and reshape them to flat vectors + loaded_checkpoint_state = torch.load(os.path.join(folder, f"{key}.pt"), weights_only=False).view(-1) + + # Partition the loaded data according to the local rank + world_size = dist.get_world_size(group=self.dp_process_group) + unpartitioned_numel = loaded_checkpoint_state.numel() + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + + if world_size * partitioned_numel != unpartitioned_numel: + padding_size = world_size * partitioned_numel - unpartitioned_numel + padding_tensor = torch.zeros(padding_size, dtype=loaded_checkpoint_state.dtype) + loaded_checkpoint_state = torch.cat([loaded_checkpoint_state, padding_tensor]) + checkpoint_state_partition = loaded_checkpoint_state.narrow(0, local_rank * partitioned_numel, + partitioned_numel) + + return checkpoint_state_partition + + def reset_swap_buffers(self): + timer_names = set() + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) + + def checkpoint_event_prologue(self): + self._partition_all_parameters() + + def checkpoint_event_epilogue(self): + if len(self.persistent_parameters) > 0: + self.persistent_parameters[0].all_gather(self.persistent_parameters) + + def empty_partition_cache(self): + self.parameter_offload.empty_partition_cache() + + def offload_states(self, + include: Container[OffloadStateTypeEnum] = None, + device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, + pin_memory: bool = True, + non_blocking: bool = False): + device = device.value + + self.empty_partition_cache() + + assert self.optimizer.__class__ == deepspeed.ops.adam.fused_adam.FusedAdam, f"Offloading is supported only for DeepSpeed FusedAdam." + + def needs_offload(target): + # return True + return target not in self.offloaded_states and (include == None or target in include) + + # HP param + if needs_offload(OffloadStateTypeEnum.hp_params): + if pin_memory: + if not hasattr(self, "hp_params_pin_buffers"): + self.hp_params_pin_buffers = [ + get_accelerator().pin_memory(torch.empty_like(t, device=device)) + for t in self.fp32_partitioned_groups_flat + ] + + for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers): + dest_buf.copy_(src_tensor, non_blocking=non_blocking) + src_tensor.data = dest_buf + else: + for buf in self.fp32_partitioned_groups_flat: + buf.data = buf.data.to(device, non_blocking=non_blocking) + self.offloaded_states.add(OffloadStateTypeEnum.hp_params) + + # LP param + if needs_offload(OffloadStateTypeEnum.lp_params): + if pin_memory: + if not hasattr(self, "lp_param_contiguous_pin_buffer"): + self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory( + torch.empty_like(self.lp_param_buffer, device=device)) + self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking) + cpu_buffer = self.lp_param_contiguous_pin_buffer + else: + cpu_buffer = self.lp_param_buffer.to(device, non_blocking=non_blocking) + + self.lp_param_buffer.data = cpu_buffer + for tensor, offset, tensor_numel in get_mapping_to_flat_buffer( + [p.ds_tensor for p in self.module.parameters()]): + tensor.data = cpu_buffer.narrow(0, offset, tensor_numel) + + self.fp16_partitioned_groups_flat.clear() + self.offloaded_states.add(OffloadStateTypeEnum.lp_params) + + # LP grad + if needs_offload(OffloadStateTypeEnum.lp_grads): + if pin_memory: + if not hasattr(self, "lp_grad_partitions_flat_pin_buffers"): + self.lp_grad_partitions_flat_pin_buffers = get_accelerator().pin_memory( + torch.empty_like(self.grad_partitions_flat_buffer, device=device)) + self.lp_grad_partitions_flat_pin_buffers.copy_(self.grad_partitions_flat_buffer, + non_blocking=non_blocking) + self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers + else: + self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to(device) + self.averaged_gradients = {} + + self.__param_id_to_grad_partition = {} + + self.offloaded_states.add(OffloadStateTypeEnum.lp_grads) + + # contiguous bucket + if needs_offload(OffloadStateTypeEnum.contiguous_grad_buffer): + for bucket in self.ipg_buckets.values(): + if bucket.buffer is not None: + # Record properties like shape, strides, etc. as a meta tensor + bucket.buffer_meta = bucket.buffer.to("meta") + bucket.buffer = None + self.offloaded_states.add(OffloadStateTypeEnum.contiguous_grad_buffer) + + # Adam + if needs_offload(OffloadStateTypeEnum.optim_states): + offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) + self.offloaded_states.add(OffloadStateTypeEnum.optim_states) + + gc.collect() + get_accelerator().empty_cache() + + def reload_states(self, non_blocking: bool = False): + + device = get_accelerator().current_device_name() + + # HP param + if OffloadStateTypeEnum.hp_params in self.offloaded_states: + if hasattr(self, "hp_params_pin_buffers"): + for src, dest in zip(self.hp_params_pin_buffers, self.fp32_partitioned_groups_flat): + dest.data = src.to(device, non_blocking=non_blocking) + else: + for buf in self.fp32_partitioned_groups_flat: + buf.data = buf.data.to(device, non_blocking=non_blocking) + self.offloaded_states.remove(OffloadStateTypeEnum.hp_params) + + # LP Param + if OffloadStateTypeEnum.lp_params in self.offloaded_states: + cpu_buffer = self.lp_param_contiguous_pin_buffer if hasattr( + self, "lp_param_contiguous_pin_buffer") else self.lp_param_buffer + self.lp_param_buffer.data = cpu_buffer.data.to(device, non_blocking=non_blocking) + self._set_fp16_partitioned_groups_flat() + + parameter_partitions = self._get_parameter_partitions() + for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(parameter_partitions): + tensor.data = self.lp_param_buffer.narrow(0, offset, tensor_numel) + self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) + + # LP grad + if OffloadStateTypeEnum.lp_grads in self.offloaded_states: + if hasattr(self, "lp_grad_partitions_flat_pin_buffers"): + self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers.to( + device, non_blocking=non_blocking) + else: + self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to( + device, non_blocking=non_blocking) + self.averaged_gradients = {} + + offset = 0 + all_params = list(itertools.chain.from_iterable(self.fp16_groups)) + for param in all_params: + self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( + 0, offset, param.partition_numel()) + offset += param.partition_numel() + + self.offloaded_states.remove(OffloadStateTypeEnum.lp_grads) + + # contiguous bucket + if OffloadStateTypeEnum.contiguous_grad_buffer in self.offloaded_states: + for bucket in self.ipg_buckets.values(): + if bucket.buffer_meta is not None: + # We don't restore the data + bucket.buffer = torch.empty_like(bucket.buffer_meta, device=device) + + self.offloaded_states.remove(OffloadStateTypeEnum.contiguous_grad_buffer) + + # Adam + if OffloadStateTypeEnum.optim_states in self.offloaded_states: + reload_adam_states(self.optimizer, device, non_blocking=non_blocking) + self.offloaded_states.remove(OffloadStateTypeEnum.optim_states) + + if non_blocking: + get_accelerator().synchronize() + + +def _handle_overflow(cpu_sum, x, i): + import math + rank = dist.get_rank() + if rank == 0: + t_i = -1 + for v_i, v in enumerate(x.data.contiguous().view(-1)): + if not math.isfinite(float(v)): + t_i = v_i + break + logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") + + +def estimate_zero3_model_states_mem_needs(total_params, + largest_layer_params, + num_gpus_per_node=1, + num_nodes=1, + cpu_offload=True, + cpu_offload_params=True, + zero_init=True, + additional_buffer_factor=1.5): + + total_gpus = num_nodes * num_gpus_per_node + gpus_factor = 1 / num_nodes + largest_layer_memory = (4 * largest_layer_params) + + if cpu_offload: + if cpu_offload_params: + gpu_mem = largest_layer_memory + + if zero_init: + cpu_mem = total_params * 18 * gpus_factor * additional_buffer_factor + else: + cpu_mem = total_params * max(4 * num_gpus_per_node, 18 * gpus_factor) * additional_buffer_factor + + else: + gpu_mem = largest_layer_memory + int(2 * total_params / total_gpus) + + if zero_init: + cpu_mem = total_params * 16 * gpus_factor * additional_buffer_factor + else: + cpu_mem = total_params * max(4 * num_gpus_per_node, 16 * gpus_factor) * additional_buffer_factor + else: + gpu_mem = largest_layer_memory + int(18 * total_params / total_gpus) + if zero_init: + cpu_mem = largest_layer_params * 4 * num_gpus_per_node * additional_buffer_factor + else: + cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor + + return int(cpu_mem), int(gpu_mem), largest_layer_memory + + +def model_to_params(model): + # shared params calculated only once + total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) + + largest_layer_params = 0 + for m in model.modules(): + # assuming no shared params within a single layer + layer_params = sum(p.numel() for p in m.parameters(recurse=False)) + largest_layer_params = max(largest_layer_params, layer_params) + + return total_params, largest_layer_params + + +def estimate_zero3_model_states_mem_needs_all_live(model, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients + for a given ``model`` and hardware setup. + + If you have an actual model object, use this function and everything will be derived + automatically. + + If it's a hypothetical model, use ``estimate_zero3_model_states_mem_needs_all_cold`` where you have to pass + the ``total_params`` and ``largest_layer_params`` explicitly. + + Args: + - ``model``: ``nn.Module`` object + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + total_params, largest_layer_params = model_to_params(model) + + estimate_zero3_model_states_mem_needs_all_cold(total_params=total_params, + largest_layer_params=largest_layer_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + additional_buffer_factor=additional_buffer_factor) + + +def estimate_zero3_model_states_mem_needs_all_cold(total_params, + largest_layer_params, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients + for a given ``model`` and hardware setup. + + If it's a hypothetical model, use this function where you have to pass + the ``total_params`` and ``largest_layer_params`` explicitly. + + If you have an actual model object, use ``estimate_zero3_model_states_mem_needs_all_live`` and everything + will be derived automatically. + + Args: + - ``total_params``: total model params + - ``largest_layer_params``: largest layer's params + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + def format_options(cpu_offload, cpu_offload_params, zero_init): + enabled = [] + padded_cpu_str = f'{OffloadDeviceEnum.cpu:4}' + param_device = padded_cpu_str if cpu_offload_params else "none" + enabled.append(f"offload_param={param_device}") + optimizer_device = padded_cpu_str if cpu_offload else "none" + enabled.append(f"offload_optimizer={optimizer_device}") + enabled.append(f"zero_init={1 if zero_init else 0}") + return ", ".join(enabled) + + nodes_str = "nodes" if num_nodes > 1 else "node" + gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" + print( + "Estimated memory needed for params, optim states and gradients for a:\n" + f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" + f"SW: Model with {int(total_params/1e6)}M total params, {int(largest_layer_params/1e6)}M largest layer params." + ) + print(" per CPU | per GPU | Options") + for cpu_offload in [True, False]: + for cpu_offload_params in [True, False]: + if not cpu_offload and cpu_offload_params: + continue + for zero_init in [True, False]: + cpu_mem, gpu_mem, largest_layer_memory = estimate_zero3_model_states_mem_needs( + total_params=total_params, + largest_layer_params=largest_layer_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + cpu_offload=cpu_offload, + cpu_offload_params=cpu_offload_params, + zero_init=zero_init, + additional_buffer_factor=additional_buffer_factor) + + options_str = format_options(cpu_offload=cpu_offload, + cpu_offload_params=cpu_offload_params, + zero_init=zero_init) + print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py new file mode 100644 index 0000000000000000000000000000000000000000..7339f538643107a7fd19e476f4f831962a405ff7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py @@ -0,0 +1,2642 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +from packaging import version as pkg_version +from collections import OrderedDict, defaultdict +from dataclasses import dataclass, field +from typing import List, Dict + +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + +from deepspeed.runtime.base_optimizer import ZeROOptimizer +from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler +from deepspeed.runtime.torch_autocast import get_all_autocast_dtypes, is_autocast_initialized, sort_dtypes +from deepspeed.runtime.utils import (empty_cache, see_memory_usage, inf, is_model_parallel_parameter, + align_dense_tensors, all_gather_dp_groups, mask_nan_or_inf_with_val_inplace) +from deepspeed.runtime.zero.config import ZeroStageEnum +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.utils import logger +from deepspeed.utils.torch import register_grad_hook +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank +from deepspeed.moe.utils import is_moe_param +from deepspeed.git_version_info import version + +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.accelerator import get_accelerator + +from deepspeed.checkpoint.constants import (DS_VERSION, GROUP_PADDINGS, PARTITION_COUNT, LOSS_SCALER, + SINGLE_PARTITION_OF_FP32_GROUPS, BASE_OPTIMIZER_STATE, + BASE_OPTIMIZER_STATE_STEP, CLIP_GRAD, ZERO_STAGE, PARAM_SLICE_MAPPINGS) +from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state +from deepspeed.checkpoint import enable_universal_checkpoint + +from deepspeed.utils import groups +# Toggle this to true to enable correctness test +# with gradient partitioning and without +pg_correctness_test = False + +OPTIMIZER_ALLGATHER_TIMER = 'optimizer_allgather' +OPTIMIZER_GRADIENTS_TIMER = 'optimizer_gradients' +OPTIMIZER_STEP_TIMER = 'optimizer_step' +OPTIMIZER_TIMERS = [OPTIMIZER_ALLGATHER_TIMER, OPTIMIZER_GRADIENTS_TIMER, OPTIMIZER_STEP_TIMER] +INITIAL_MICRO_STEP_ID = -1 + + +def input(msg): + return + + +def split_half_float_double(tensors): + device_type = get_accelerator().device_name() + dtypes = [ + "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type), + "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type) + ] + buckets = [] + for i, dtype in enumerate(dtypes): + bucket = [t for t in tensors if t.type() == dtype] + if bucket: + buckets.append(bucket) + return buckets + + +def isclose(a, b, rtol=1e-09, atol=0.0): + return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) + + +def lcm(x, y): + from fractions import gcd # or can import gcd from `math` in Python 3 + return x * y // gcd(x, y) + + +def get_alignment_padding(tensor_list, alignment): + num_elements = sum([tensor.numel() for tensor in tensor_list]) + remainder = num_elements % alignment + return (alignment - remainder) if remainder else remainder + + +def print_rank_msg(msg): + print(f"rank {dist.get_rank()} - {msg}") + + +def _get_padded_tensor(src_tensor, size): + if src_tensor.numel() >= size: + return src_tensor + padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) + slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) + slice_tensor.data.copy_(src_tensor.data) + return padded_tensor + + +def _pad_tensor_by_size(src_tensor, pad_size, dtype, device): + padded_tensor = torch.zeros(src_tensor.numel() + pad_size, dtype=dtype, device=device) + padded_tensor.data[:src_tensor.numel()].copy_(src_tensor.data) + return padded_tensor + + +@dataclass +class IPGBucket: + buffer: List[torch.Tensor] = field(default_factory=list) + params: List[torch.Tensor] = field(default_factory=list) + grads: List[torch.Tensor] = field(default_factory=list) + elements: int = 0 + index: int = 0 + has_moe_params: bool = False + + def clear(self): + self.params.clear() + self.grads.clear() + self.elements = 0 + self.index = 0 + self.has_moe_params = False + + +class DeepSpeedZeroOptimizer(ZeROOptimizer): + """ + DeepSpeedZeroOptimizer designed to reduce the memory footprint + required for training large deep learning models. + + For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models + https://arxiv.org/abs/1910.02054 + + For usage examples, refer to TODO: DeepSpeed Tutorial + + """ + + def __init__(self, + init_optimizer, + param_names, + timers, + static_loss_scale=1.0, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + use_multi_rank_bucket_allreduce=True, + allgather_bucket_size=5000000000, + dp_process_group=None, + expert_parallel_group=None, + expert_data_parallel_group=None, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + mpu=None, + clip_grad=0.0, + gradient_accumulation_dtype=torch.float32, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1.0, + gradient_accumulation_steps=1, + ignore_unused_parameters=True, + partition_grads=True, + round_robin_gradients=False, + has_moe_layers=False, + fp16_master_weights_and_gradients=False, + elastic_checkpoint=False, + check_grad_overflow=True): + + if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: + self.cpu_offload = True + self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory + else: + self.cpu_offload = False + self.cpu_offload_pin_memory = False + + if dist.get_rank() == 0: + logger.info(f"Reduce bucket size {reduce_bucket_size}") + logger.info(f"Allgather bucket size {allgather_bucket_size}") + logger.info(f"CPU Offload: {self.cpu_offload}") + logger.info(f'Round robin gradient partitioning: {round_robin_gradients}') + # The fused optimizer does all the work. We need this layer for two reason: + # 1. maintain same user API from apex.fp16_utils + # 2. keep common stuff here in case we need to add ne552w fused optimizer later + + self.elastic_checkpoint = elastic_checkpoint + self.check_grad_overflow = check_grad_overflow + self.param_names = param_names + self.mpu = mpu + # differences from apex.fp16_utils: + # - assume all model params in fp16 + # - assume all params requires grad + # - flat by groups, not keeping state. TODO: remove state explicitly? + # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? + if not get_accelerator().is_available(): + raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).") + self.optimizer = init_optimizer + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + + # ZeRO stage 1 (False) or 2 (True) + self.partition_gradients = partition_grads + self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1" + + self.timers = timers + + self.reduce_scatter = reduce_scatter + + self.overlap_comm = overlap_comm + + self.deepspeed_adam_offload = self.cpu_offload + + self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu' + + self.dp_process_group = dp_process_group + self.sequence_parallel_size = groups._get_sequence_parallel_world_size() + #expert parallel group + self.ep_process_group = expert_parallel_group + + #data parallel group for experts + self.expert_dp_process_group = expert_data_parallel_group + + #data parallel size for non-experts + dp_size = dist.get_world_size(group=self.dp_process_group) + + #For MoE models this maybe different for different param group + #It will be modified during MoE setup later in the init + self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] + self.partition_count = [dp_size for i in range(len(self.optimizer.param_groups))] + + self.is_gradient_accumulation_boundary = True + + # CPU-Offload requires contiguous gradients + self.contiguous_gradients = contiguous_gradients or self.cpu_offload + + self.has_moe_layers = has_moe_layers + if self.has_moe_layers: + self._configure_moe_settings() + self._global_grad_norm = 0. + + if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): + self.model_parallel_group = None + self.model_parallel_world_size = 1 + self.model_parallel_rank = 0 + else: + self.model_parallel_group = mpu.get_model_parallel_group() + self.model_parallel_world_size = mpu.get_model_parallel_world_size() + self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu) + + self.overflow = False + self.clip_grad = clip_grad + self.communication_data_type = communication_data_type + self.gradient_predivide_factor = gradient_predivide_factor + self.postscale_gradients = postscale_gradients + self.gradient_accumulation_steps = gradient_accumulation_steps + self.micro_step_id = INITIAL_MICRO_STEP_ID + self.ignore_unused_parameters = ignore_unused_parameters + self.round_robin_gradients = round_robin_gradients + + self.extra_large_param_to_reduce: Dict[int, torch.Tensor] = {} + self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients + + if self.fp16_master_weights_and_gradients: + assert self.cpu_offload and type(self.optimizer) in [DeepSpeedCPUAdam], \ + f"fp16_master_and_gradients requires optimizer to support keeping fp16 master and gradients while keeping the optimizer states in fp32."\ + f"Currently only supported using ZeRO-Offload with DeepSpeedCPUAdam. But current setting is ZeRO-Offload:{self.cpu_offload} and optimizer type {type(self.optimizer)}." \ + f"Either disable fp16_master_weights_and_gradients or enable {self.zero_stage_string} Offload with DeepSpeedCPUAdam." + + if self.reduce_scatter and self.partition_gradients: + valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) + assert self.communication_data_type in valid_reduce_scatter_dtypes, f"{self.zero_stage_string} supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" + assert self.gradient_predivide_factor == 1.0, f"gradient_predivide_factor != 1.0 is not yet supported with {self.zero_stage_string} with reduce scatter enabled" + assert self.postscale_gradients, f"pre-scale gradients is not yet supported with {self.zero_stage_string} with reduce scatter enabled" + + # param flattened by groups + self.bit16_groups = [] + self.bit16_groups_flat = [] + + # param partitioned by data parallel degree + # this will contain a list of equal sized tensors + # each of which will be updated by a different process + self.parallel_partitioned_bit16_groups = [] + + # a single 32-bit partition of the parallel partitioned parameters + # that this process will update + self.single_partition_of_fp32_groups = [] + + # a 16-bit CPU param buffer for cpu offload + if self.cpu_offload: + self.param_buffer_of_bit16_for_cpu_offload_groups = [] + + # param partition info + + # These are the parameters in each group that will not be updated by this process directly + self.params_not_in_partition = [] + + # These are the parameters that will be updated by this process directly + self.params_in_partition = [] + + # Offset from the first parameter in the self.params_in_partition + # the parameter boundaries may not align with partition boundaries + # so we need to keep track of the offset + self.first_offset = [] + + # number of elements per partition in each group + self.partition_size = [] + + # align nccl all-gather send buffers to 4-byte boundary + self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 + + assert ( + allgather_bucket_size % self.nccl_start_alignment_factor == 0 + ), f"allgather_bucket_size must be a multiple of nccl_start_alignment_factor, {self.nccl_start_alignment_factor} " + + self.all_reduce_print = False + self.dtype = self.optimizer.param_groups[0]['params'][0].dtype + self.gradient_accumulation_dtype = gradient_accumulation_dtype + + if self.dtype != self.gradient_accumulation_dtype: + self.use_separate_grad_accum = True + else: + self.use_separate_grad_accum = False + if self.use_separate_grad_accum and not self.partition_gradients: + self.use_grad_accum_attribute = True + else: + self.use_grad_accum_attribute = False + + self.round_robin_bit16_groups = [] + self.round_robin_bit16_indices = [] + self.round_robin_bit16_meta = [] + + # Use different parallel to do all_to_all_reduce related things + # padding on each partition for alignment purposes + self.groups_padding = [] + # loop to deal with groups + for i, param_group in enumerate(self.optimizer.param_groups): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + + # push this group to list before modify + # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group + trainable_parameters = [] + for param in param_group['params']: + if param.requires_grad: + param.grad_accum = None + param.param_idx_in_group = len(trainable_parameters) + trainable_parameters.append(param) + self.bit16_groups.append(trainable_parameters) + + # not sure why apex was cloning the weights before flattening + # removing cloning here + + see_memory_usage(f"Before moving param group {i} to CPU") + # move all the parameters to cpu to free up GPU space for creating flat buffer + + # Create temp CPU param copies, free accelerator tensors + orig_group_numel = 0 + for param in self.bit16_groups[i]: + orig_group_numel += param.numel() + param.cpu_data = param.data.cpu() + param.data = torch.empty(1).to(param.device) + + empty_cache() + see_memory_usage(f"After moving param group {i} to CPU", force=False) + + # Reorder group parameters for load balancing of gradient partitioning during backward among ranks. + # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks. + # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging + # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m). + if self.round_robin_gradients: + round_robin_tensors, round_robin_indices = self._round_robin_reorder( + self.bit16_groups[i], dist.get_world_size(group=self.real_dp_process_group[i])) + else: + round_robin_tensors = self.bit16_groups[i] + round_robin_indices = list(range(len(self.bit16_groups[i]))) + + self.round_robin_bit16_groups.append(round_robin_tensors) + self.round_robin_bit16_indices.append(round_robin_indices) + + # Create meta tensors list, ordered according to round_robin_tensors + meta_tensors = [] + for param in round_robin_tensors: + meta_tensors.append(torch.zeros_like(param.cpu_data, device="meta")) + self.round_robin_bit16_meta.append(meta_tensors) + + # create flat buffer in CPU + flattened_buffer = self.flatten_dense_tensors_aligned( + self.round_robin_bit16_groups[i], + self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i]), + use_cpu_data=True) + + # free temp CPU params + for param in self.bit16_groups[i]: + del param.cpu_data + + # Move CPU flat tensor to the accelerator memory. + self.bit16_groups_flat.append(flattened_buffer.to(get_accelerator().current_device_name())) + del flattened_buffer + + see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False) + + if dist.get_rank(group=self.real_dp_process_group[i]) == 0: + see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False) + + # set model bit16 weight to slices of flattened buffer + self._update_model_bit16_weights(i) + + # divide the flat weights into near equal partition equal to the data parallel degree + # each process will compute on a different part of the partition + data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i) + self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) + + # Record padding required for alignment + left_boundary = sum([t.numel() for t in data_parallel_partitions[:partition_id]]) + curr_partition_size = data_parallel_partitions[partition_id].numel() + + if orig_group_numel <= left_boundary: + padding = curr_partition_size + elif orig_group_numel < left_boundary + curr_partition_size: + padding = left_boundary + curr_partition_size - orig_group_numel + else: + padding = 0 + self.groups_padding.append(padding) + + # verify that data partition start locations are 4-byte aligned + for partitioned_data in data_parallel_partitions: + assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0) + + # A partition of the fp32 master weights that will be updated by this process. + # Note that the params in single_partition_of_fp32_groups is cloned and detached + # from the origin params of the model. + if not fp16_master_weights_and_gradients: + weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to( + self.device).clone().float().detach() + else: + weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to( + self.device).clone().half().detach() + + if self.cpu_offload: + weights_partition = get_accelerator().pin_memory(weights_partition) + temp_dtype = self.parallel_partitioned_bit16_groups[i][partition_id].dtype + temp_buffer_bit16 = torch.full(weights_partition.shape, + fill_value=0.0, + dtype=temp_dtype, + device=weights_partition.device) + if self.cpu_offload_pin_memory: + temp_pinned = get_accelerator().pin_memory(temp_buffer_bit16) + self.param_buffer_of_bit16_for_cpu_offload_groups.append(temp_pinned) + else: + self.param_buffer_of_bit16_for_cpu_offload_groups.append(temp_buffer_bit16) + + self.single_partition_of_fp32_groups.append(weights_partition) + + # Set local optimizer to have flat params of its own partition. + # After this, the local optimizer will only contain its own partition of params. + # In that case, the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params(zero stage1). + self.single_partition_of_fp32_groups[ + i].requires_grad = True # keep this in case internal optimizer uses it + param_group['params'] = [self.single_partition_of_fp32_groups[i]] + + partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_dp_process_group[i]) + params_in_partition, params_not_in_partition, first_offset = self.get_partition_info( + self.round_robin_bit16_groups[i], partition_size, partition_id) + + self.partition_size.append(partition_size) + self.params_in_partition.append(params_in_partition) + self.params_not_in_partition.append(params_not_in_partition) + self.first_offset.append(first_offset) + + self.reduce_bucket_size = int(reduce_bucket_size) + self.use_multi_rank_bucket_allreduce = use_multi_rank_bucket_allreduce + self.allgather_bucket_size = int(allgather_bucket_size) + + self.reduction_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream() + #self.copy_grad_stream = get_accelerator().Stream() + self.callback_queued = False + + self.param_dict = {} + + # map between param_id and bool to specify if a param is in this partition + self.is_param_in_current_partition = {} + + if is_autocast_initialized(): + comm_dtypes = get_all_autocast_dtypes([p for params in self.bit16_groups for p in params]) + self.torch_autocast_gradscaler = torch.amp.GradScaler(device=get_accelerator().device_name()) + else: + comm_dtypes = {self.communication_data_type} + self.torch_autocast_gradscaler = None + + self.ipg_buckets: Dict[torch.dtype, IPGBucket] = {dtype: IPGBucket() for dtype in comm_dtypes} + + self.params_already_reduced = [] + self._release_ipg_buffers() + self.previous_reduced_grads: Dict[int, List[torch.Tensor]] = defaultdict(list) + + # simplified param id + self.param_id = {} + + #interesting code: unique ids being assigned to individual parameters + largest_param_numel = 0 + count = 0 + for i, params_group in enumerate(self.bit16_groups): + for param in params_group: + unique_id = id(param) + self.param_id[unique_id] = count + self.param_dict[count] = param + self.params_already_reduced.append(False) + if param.numel() > largest_param_numel: + largest_param_numel = param.numel() + count = count + 1 + + for param_group in self.params_in_partition: + for param in param_group: + self.is_param_in_current_partition[self.get_param_id(param)] = True + + for param_group in self.params_not_in_partition: + for param in param_group: + self.is_param_in_current_partition[self.get_param_id(param)] = False + + if self.cpu_offload: + self.accumulated_grads_in_cpu = {} + self.norm_for_param_grads = {} + self.local_overflow = False + self.grad_position = {} + self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel, + device=self.device, + dtype=self.dtype) + if self.cpu_offload_pin_memory: + self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory( + self.temp_grad_buffer_for_cpu_offload) + self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel, + device=get_accelerator().current_device_name(), + dtype=self.dtype) + for i, params_group in enumerate(self.bit16_groups): + self.get_grad_position(i, self.params_in_partition[i], self.first_offset[i], self.partition_size[i]) + + # mapping from parameter to partition that it belongs to + self.param_to_partition_ids = {} + + # stores if a partition has been reduced in this step + self.is_partition_reduced = {} + + # number of grads in partition that still need to be computed + self.remaining_grads_in_partition = {} + + # total number of grads in partition + self.total_grads_in_partition = {} + + # stores if a grad in a partition has been computed or not + self.is_grad_computed = {} + + # stores the offset at which a parameter gradient needs to be inserted in a partition + self.grad_partition_insertion_offset = {} + + # the offset in the gradient at which it must be inserted at the beginning of the partition + self.grad_start_offset = {} + + # will store the averaged gradients required by this partition + self.averaged_gradients = {} + + # For cpu_offload, will store the averaged gradients required by this partition + self.offload_gradient_dict = {} + + # store index of first parameter in each partition + self.first_param_index_in_partition = {} + + # initializes all data structures for implementing gradient partitioning + self.initialize_gradient_partitioning_data_structures() + + # resets the data structure value for the next backward propagation + self.reset_partition_gradient_structures() + + # creates backward hooks for the following special handling of gradients + # 1. upcasting for fp32 gradient accumulation + # 2. gradient partitioning + # 3. overlapping backward and reduction + self._grad_acc_hooks = [] + + if (self.partition_gradients or self.overlap_comm or self.use_grad_accum_attribute + or self.contiguous_gradients): + self.create_gradient_handling_hooks() + + self.ready_for_gradients = False + self.custom_loss_scaler = False + self.external_loss_scale = None + + # we may have a way of fusing dynamic scale. Do not support for now + self.loss_scaler = CreateLossScaler(dtype=self.dtype, + static_loss_scale=static_loss_scale, + dynamic_scaling=dynamic_loss_scale, + dynamic_loss_args=dynamic_loss_args) + self.dynamic_loss_scale = self.loss_scaler.dynamic + + if self.dtype != torch.float16: + # Only fp16 should use dynamic loss scaling + assert self.loss_scaler.cur_scale == 1.0 + assert not self.dynamic_loss_scale + + see_memory_usage("Before initializing optimizer states", force=True) + self.initialize_optimizer_states() + see_memory_usage("After initializing optimizer states", force=True) + + if dist.get_rank() == 0: + logger.info(f"optimizer state initialized") + + if dist.get_rank(group=self.dp_process_group) == 0: + see_memory_usage(f"After initializing ZeRO optimizer", force=True) + + self._link_all_hp_params() + self._hp_optimizer_states_linked = False + + self._enable_universal_checkpoint() + self._param_slice_mappings = self._create_param_mapping() + if self.cpu_offload: + self._create_optimizer_mapping() + + def destroy(self): + for i, _ in enumerate(self.optimizer.param_groups): + for p in self.bit16_groups[i]: + if getattr(p, '_hp_mapping', None): + p._hp_mapping = None + for hook in self._grad_acc_hooks: + hook.remove() + self.print_rank_0("Removed grad acc hooks") + + def _enable_universal_checkpoint(self): + for lp_param_group in self.bit16_groups: + enable_universal_checkpoint(param_list=lp_param_group) + + def _create_param_mapping(self): + param_mapping = [] + for i, _ in enumerate(self.optimizer.param_groups): + param_mapping_per_group = OrderedDict() + for lp in self.bit16_groups[i]: + if lp._hp_mapping is not None: + lp_name = self.param_names[lp] + param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() + param_mapping.append(param_mapping_per_group) + + return param_mapping + + def _create_optimizer_mapping(self): + for i, _ in enumerate(self.optimizer.param_groups): + for lp in self.bit16_groups[i]: + if lp._hp_mapping is not None: + lp._zero_optimizer = self + + def _link_all_hp_params(self): + if self.cpu_offload: + self._get_offload_gradient_dict() + + for i, _ in enumerate(self.optimizer.param_groups): + # Link bit16 and fp32 params in partition + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + partition_size = self.bit16_groups_flat[i].numel() // dist.get_world_size( + group=self.real_dp_process_group[i]) + flat_hp_partition = self.single_partition_of_fp32_groups[i] + link_hp_params(lp_param_list=self.bit16_groups[i], + flat_hp_partition=flat_hp_partition, + gradient_dict=self.averaged_gradients, + offload_gradient_dict=self.offload_gradient_dict, + use_offload=self.cpu_offload, + param_group_index=i, + partition_start=partition_id * partition_size, + partition_size=partition_size, + dp_group=self.real_dp_process_group[i]) + + def _lazy_init_hp_params_optimizer_state(self): + if not self._hp_optimizer_states_linked: + for i, _ in enumerate(self.optimizer.param_groups): + lazy_init_hp_params_optimizer_state(self.bit16_groups[i], self.single_partition_of_fp32_groups[i], + self.optimizer.state) + self._hp_optimizer_states_linked = True + + def is_moe_group(self, group): + return 'moe' in group and group['moe'] + + def _configure_moe_settings(self): + # if we're using ZeRO stage 2, ensure contiguous gradients are used + if self.partition_gradients: + assert self.contiguous_gradients, "Contiguous Gradients in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" + # NOTE: To run ZeRO stage 1 with MoE, we need to set self.contiguous_gradients to True or ignore the assertion + if not self.partition_gradients and not self.contiguous_gradients: + logger.warning( + "ZeRO Stage 1 has not been thoroughly tested with MoE. This configuration is still experimental.") + assert self.reduce_scatter, "Reduce Scatter in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" + + assert any( + [self.is_moe_group(group) for group in self.optimizer.param_groups] + ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" + self.is_moe_param_group = [] + for i, group in enumerate(self.optimizer.param_groups): + if self.is_moe_group(group): + assert all([is_moe_param(param) + for param in group['params']]), "All params in MoE group must be MoE params" + self.real_dp_process_group[i] = self.expert_dp_process_group[group['name']] + self.partition_count[i] = dist.get_world_size(group=self.expert_dp_process_group[group['name']]) + self.is_moe_param_group.append(True) + else: + self.is_moe_param_group.append(False) + + assert self.expert_dp_process_group is not None, "Expert data parallel group should be configured with MoE" + assert self.ep_process_group is not None, "Expert parallel group should be configured with MoE" + + def _update_model_bit16_weights(self, group_index): + updated_params = self.unflatten(self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index]) + for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params): + p.data = q.data + + # set model fp16 weight to slices of reordered flattened buffer + for param_index, param in enumerate(self.bit16_groups[group_index]): + new_index = self.round_robin_bit16_indices[group_index][param_index] + param.data = self.round_robin_bit16_groups[group_index][new_index].data + + def _round_robin_reorder(self, tensor_list, num_partitions): + + # disable round robin if need to debug something + # return tensor_list, list(range(len(tensor_list))) + + partition_tensors = {} + + for i, tensor in enumerate(tensor_list): + j = i % num_partitions + if not j in partition_tensors: + partition_tensors[j] = [] + partition_tensors[j].append((i, tensor)) + + reordered_tensors = [] + reordered_indices = {} + + for partition_index in partition_tensors.keys(): + for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]): + reordered_indices[original_index] = len(reordered_tensors) + reordered_tensors.append(tensor) + + return reordered_tensors, reordered_indices + + def _release_ipg_buffers(self): + if self.contiguous_gradients: + for bucket in self.ipg_buckets.values(): + bucket.buffer.clear() + + self.grads_in_partition = None + self.grads_in_partition_offset = 0 + self.ready_for_gradients = False + + def initialize_optimizer_states(self): + + for i, group in enumerate(self.bit16_groups): + single_grad_partition = torch.zeros(int(self.partition_size[i]), + dtype=self.single_partition_of_fp32_groups[i].dtype, + device=self.device) + self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory( + single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition + + # Initialize the optimizer states with the flattened fp32 partition. + # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers + # which do lazy initialization of the state at the first call to step. + if isinstance(self.optimizer, torch.optim.Adagrad): + self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults) + + if not self.cpu_offload: + for group in self.single_partition_of_fp32_groups: + group.grad = None #class init + + return + + ######################################################################### + #################### ZeRO Stage 1 - reduce gradients #################### + ######################################################################### + def reduce_gradients(self, pipeline_parallel=False): + world_size = dist.get_world_size(self.dp_process_group) + my_rank = dist.get_rank(self.dp_process_group) + + # with PP we must create ipg buffer, since backward is handled outside zero + if pipeline_parallel and self.contiguous_gradients: + for dtype, bucket in self.ipg_buckets.items(): + bucket.buffer.append( + torch.empty(int(self.reduce_bucket_size), + dtype=dtype, + device=get_accelerator().current_device_name())) + bucket.index = 0 + + if not self.overlap_comm: + for i, group in enumerate(self.bit16_groups): + for param in group: + grad_reduc = self.get_gradient_for_reduction(param) + if grad_reduc is not None: + self.reduce_ready_partitions_and_remove_grads(param, i) + # reduce any pending grads in either hook/non-hook case + self.overlapping_partition_gradients_reduce_epilogue() + + ######################################################################### + #########################ZeRO Partition Gradients######################## + ######################################################################### + + def get_first_param_index(self, group_id, param_group, partition_id): + for index, param in enumerate(param_group): + param_id = self.get_param_id(param) + if group_id in self.param_to_partition_ids and param_id in self.param_to_partition_ids[group_id]: + if partition_id in self.param_to_partition_ids[group_id][param_id]: + return index + return None + + def initialize_gradient_partitioning_data_structures(self): + + for i, param_group in enumerate(self.round_robin_bit16_groups): + total_partitions = dist.get_world_size(group=self.real_dp_process_group[i]) + + self.param_to_partition_ids[i] = {} + self.is_partition_reduced[i] = {} + self.total_grads_in_partition[i] = {} + self.remaining_grads_in_partition[i] = {} + self.is_grad_computed[i] = {} + self.grad_partition_insertion_offset[i] = {} + self.grad_start_offset[i] = {} + self.first_param_index_in_partition[i] = {} + + for partition_id in range(total_partitions): + self.is_grad_computed[i][partition_id] = {} + self.grad_partition_insertion_offset[i][partition_id] = {} + self.grad_start_offset[i][partition_id] = {} + self.total_grads_in_partition[i][partition_id] = 0 + self.initialize_gradient_partition(i, param_group, partition_id) + self.is_partition_reduced[i][partition_id] = False + self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( + i, param_group, partition_id) + + def independent_gradient_partition_epilogue(self): + self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) + self.reduce_ipg_grads() + self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) + + # if dist.get_rank() == 0: + # logger.info("Params already reduced %s", self.params_already_reduced) + for i in range(len(self.params_already_reduced)): + self.params_already_reduced[i] = False + + if self.overlap_comm: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + # It is safe to clear previously reduced grads of other partitions + self._clear_previous_reduced_grads() + + if self.cpu_offload is False: + for i, _ in enumerate(self.bit16_groups): + + if not i in self.averaged_gradients or self.averaged_gradients[i] is None: + self.averaged_gradients[i] = self.get_flat_partition( + self.params_in_partition[i], + self.first_offset[i], + self.partition_size[i], + dtype=self.gradient_accumulation_dtype, + device=get_accelerator().current_device_name(), + return_tensor_list=True) + else: + avg_new = self.get_flat_partition(self.params_in_partition[i], + self.first_offset[i], + self.partition_size[i], + dtype=self.gradient_accumulation_dtype, + device=get_accelerator().current_device_name(), + return_tensor_list=True) + + for accumulated_grad, new_avg_grad in zip(self.averaged_gradients[i], avg_new): + accumulated_grad.add_(new_avg_grad) + + self._release_ipg_buffers() + + # No need to keep the gradients anymore. + # All gradients required by the step + # are in self.averaged_gradients + self.zero_grad(set_to_none=True) + see_memory_usage(f"End ipg_epilogue") + + # resets all partition to no reduced + # sets remaining grads to the total number of grads in each partition + # set is grad computed to false for all grads in partition + def reset_partition_gradient_structures(self): + for i, _ in enumerate(self.bit16_groups): + total_partitions = dist.get_world_size(group=self.real_dp_process_group[i]) + for partition_id in range(total_partitions): + self.is_partition_reduced[i][partition_id] = False + self.remaining_grads_in_partition[i][partition_id] = self.total_grads_in_partition[i][partition_id] + + for param_id in self.is_grad_computed[i][partition_id]: + self.is_grad_computed[i][partition_id][param_id] = False + + def initialize_gradient_partition(self, i, param_group, partition_id): + + def set_key_value_list(dictionary, key, value): + if key in dictionary: + dictionary[key].append(value) + else: + dictionary[key] = [value] + + def increment_value(dictionary, key): + if key in dictionary: + dictionary[key] += 1 + else: + dictionary[key] = 1 + + partition_size = self.partition_size[i] + + start_index = partition_size * partition_id + end_index = partition_size * (partition_id + 1) + + current_index = 0 + first_offset = 0 + + for param in param_group: + + param_size = param.numel() + param_id = self.get_param_id(param) + + if start_index <= current_index < end_index: + set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) + increment_value(self.total_grads_in_partition[i], partition_id) + + self.is_grad_computed[i][partition_id][param_id] = False + + self.grad_partition_insertion_offset[i][partition_id][param_id] = current_index - start_index + self.grad_start_offset[i][partition_id][param_id] = 0 + + elif current_index < start_index < (current_index + param_size): + assert (first_offset == 0 + ), "This can happen either zero or only once as this must be the first tensor in the partition" + first_offset = start_index - current_index + + set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) + increment_value(self.total_grads_in_partition[i], partition_id) + + self.is_grad_computed[i][partition_id][param_id] = False + + self.grad_partition_insertion_offset[i][partition_id][param_id] = 0 + self.grad_start_offset[i][partition_id][param_id] = first_offset + + current_index = current_index + param_size + + def overlapping_partition_gradients_reduce_epilogue(self): + self.independent_gradient_partition_epilogue() + + def _fill_param_grad_accum_attribute(self, param): + if param.grad is not None: + if param.grad_accum is None: + param.grad_accum = param.grad.to(self.gradient_accumulation_dtype) + else: + param.grad_accum.add_(param.grad.to(self.gradient_accumulation_dtype).view(param.grad_accum.shape)) + param.grad = None + + def fill_grad_accum_attribute(self): + for group in self.bit16_groups: + for param in group: + self._fill_param_grad_accum_attribute(param) + + def get_gradient_for_reduction(self, param): + if self.use_grad_accum_attribute: + return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None + else: + return param.grad + + def get_param_gradient_attribute(self, param): + return param.grad_accum if self.use_grad_accum_attribute else param.grad + + # Clear the tensor the reduction gradient attribute is pointing to + def clear_grad_attribute(self, param): + if self.use_grad_accum_attribute: + param.grad_accum = None + else: + param.grad = None + + def create_gradient_handling_hooks(self): + for i, param_group in enumerate(self.bit16_groups): + for param in param_group: + if param.requires_grad: + + def wrapper(param, i): + + def grad_handling_hook(*notneeded): + self.process_gradients(param, i) + + self._grad_acc_hooks.append(register_grad_hook(param, grad_handling_hook)) + + wrapper(param, i) + + def get_param_id(self, param): + unique_id = id(param) + return self.param_id[unique_id] + + # create a flat tensor aligned at the alignment boundary + def flatten_dense_tensors_aligned(self, tensor_list, alignment, use_cpu_data=False): + tensor_list = [param.cpu_data for param in tensor_list] if use_cpu_data else tensor_list + return self.flatten(align_dense_tensors(tensor_list, alignment)) + + ############### Independent Partition Gradient ######################## + def reduce_independent_p_g_buckets_and_remove_grads(self, param, i): + + grad_reduc = self.get_gradient_for_reduction(param) + bucket = self.ipg_buckets[self.get_param_comm_dtype(param)] + if bucket.elements + param.numel() > self.reduce_bucket_size: + self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel(), param.dtype) + self.reduce_ipg_grads() + if self.contiguous_gradients and self.overlap_comm: + # Swap index between 0 and 1 + bucket.index = 1 - bucket.index + self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel(), param.dtype) + + param_id = self.get_param_id(param) + assert self.params_already_reduced[param_id] == False, \ + f"The parameter {param_id} has already been reduced. \ + Gradient computed twice for this partition. \ + Multiple gradient reduction is currently not supported" + + if self.contiguous_gradients: + if param.numel() > self.reduce_bucket_size: + self.extra_large_param_to_reduce[param.dtype] = param + else: + # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening + new_grad_tensor = bucket.buffer[bucket.index].narrow(0, bucket.elements, param.numel()) + new_grad_tensor.copy_(grad_reduc.view(-1)) + grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) + + bucket.elements += param.numel() + + assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient" + + # deal with a use-case of transient grads that will be generated in a loop for the same computation involving some model params - e.g. when performing a tiled memory calculation that shards the normal single sub-module call into a loop over a shards. + if getattr(param, "ds_grad_is_ready", True): + bucket.grads.append(grad_reduc) + bucket.params.append((i, param.param_idx_in_group, param_id)) + + #make sure the average tensor function knows how to average the gradients + if is_moe_param(param): + bucket.has_moe_params = True + + self.report_ipg_memory_usage("End ipg_remove_grads", 0) + + def print_rank_0(self, message): + if dist.get_rank() == 0: + logger.info(message) + + def gradient_reduction_w_predivide(self, tensor, communication_data_type: torch.dtype): + if tensor.size().numel() == 0: + return tensor + + dp_world_size = dist.get_world_size(group=self.dp_process_group) + + tensor_to_allreduce = tensor + + if communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(communication_data_type) + + if self.postscale_gradients: + if self.gradient_predivide_factor != 1.0: + tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor) + + dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) + + if self.gradient_predivide_factor != dp_world_size: + tensor_to_allreduce.mul_(self.gradient_predivide_factor / + (dp_world_size / float(self.sequence_parallel_size))) + else: + tensor_to_allreduce.div_(dp_world_size / float(self.sequence_parallel_size)) + dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) + + if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + tensor.copy_(tensor_to_allreduce) + + return tensor + + def allreduce_and_copy_with_multiple_ranks(self, + small_bucket, + communication_data_type: torch.dtype, + log=None, + divide=True, + process_group=None, + bucket_ranks=None): + process_group = self.dp_process_group if process_group is None else process_group + allreduced = self.allreduce_bucket(small_bucket, + communication_data_type, + log=log, + divide=divide, + process_group=process_group) + for buf, synced, bucket_rank in zip(small_bucket, self.unflatten(allreduced, small_bucket), bucket_ranks): + if dist.get_rank(group=process_group) == bucket_rank: + buf.copy_(synced) + + def allreduce_and_scatter(self, + bucket, + communication_data_type: torch.dtype, + numel_per_bucket=500000000, + log=None, + divide=True, + process_group=None): + small_bucket = [] + small_bucket_ranks = [] + numel = 0 + allreduce_sizes = [] + + for i, bucket_elem in enumerate(bucket): + rank, tensor = bucket_elem + small_bucket.append(tensor) + small_bucket_ranks.append(rank) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy_with_multiple_ranks(small_bucket, + communication_data_type, + log=None, + divide=divide, + process_group=process_group, + bucket_ranks=small_bucket_ranks) + small_bucket = [] + small_bucket_ranks = [] + numel = 0 + + if len(small_bucket) > 0: + self.allreduce_and_copy_with_multiple_ranks(small_bucket, + communication_data_type, + log=None, + divide=divide, + process_group=process_group, + bucket_ranks=small_bucket_ranks) + + def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dtype): + if self.overlap_comm: + stream = self.reduction_stream + if not get_accelerator().resolves_data_dependency(): + stream.wait_stream(get_accelerator().current_stream()) + get_accelerator().current_stream().wait_stream(stream) + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + if not self.reduce_scatter: + self.gradient_reduction_w_predivide(tensor, communication_data_type) + return + + # Accumulate destination ranks and bucket offsets for each gradient slice. + # Note: potential future optimization, record access pattern of parameters + # in backward pass and partition gradients w.r.t. access pattern so that our + # bucket is guaranteed to be contiguous w.r.t. ranks + rank_and_offsets = [] + real_dp_process_group = [] + curr_size = 0 + prev_id, prev_process_group = -1, None + + process_group = self.dp_process_group + # count = 0 + bucket = self.ipg_buckets[communication_data_type] + for i, param_idx_in_group, param_id in bucket.params: + param = self.bit16_groups[i][param_idx_in_group] + + process_group = self.dp_process_group + + if bucket.has_moe_params: + process_group = self.expert_dp_process_group[param.group_name] if is_moe_param( + param) else self.dp_process_group + + partition_ids = self.param_to_partition_ids[i][param_id] + assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids + ]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}" + partition_size = self.partition_size[i] + # Get all partition ids + their offsets + partition_ids_w_offsets = [] + for partition_id in partition_ids: + offset = self.grad_start_offset[i][partition_id][param_id] + partition_ids_w_offsets.append((partition_id, offset)) + partition_ids_w_offsets.sort(key=lambda t: t[1]) + + # Calculate rank and offsets for grad slices + for idx in range(len(partition_ids_w_offsets)): + partition_id, offset = partition_ids_w_offsets[idx] + + # if dist.get_rank() == 0 and count < 100: + # print(f"Rank {dist.get_rank()} rank offset id {idx} calculated dp size {dist.get_world_size(group=process_group)} real dp size {dist.get_world_size(self.real_dp_process_group[i])} and dst: {partition_id}") + # count += 1 + + # Calculate numel for grad slice depending on partition location + if idx == len(partition_ids_w_offsets) - 1: + # Last partition_id uses its own offset + numel = param.numel() - offset + else: + # Set numel to next partition's offset + numel = partition_ids_w_offsets[idx + 1][1] - offset + + # Merge bucket ranges if they belong to the same rank + if partition_id == prev_id and process_group == prev_process_group: + prev_pid, prev_size, prev_numel = rank_and_offsets[-1] + rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel) + else: + rank_and_offsets.append((partition_id, curr_size, numel)) + real_dp_process_group.append(process_group) + curr_size += numel + prev_id, prev_process_group = partition_id, process_group + + tensor.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) + + buckets = {} + for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets): + grad_slice = tensor.narrow(0, int(bucket_offset), int(numel)) + bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else ( + dst, real_dp_process_group[i]) + if bucket_key not in buckets: + buckets[bucket_key] = [] + if self.use_multi_rank_bucket_allreduce: + buckets[bucket_key].append((dst, grad_slice)) + else: + buckets[bucket_key].append(grad_slice) + + for bucket_key in buckets: + if self.use_multi_rank_bucket_allreduce: + self.allreduce_and_scatter(buckets[bucket_key], + communication_data_type, + numel_per_bucket=self.reduce_bucket_size, + divide=False, + process_group=bucket_key) + else: + dst, process_group = bucket_key + self.allreduce_no_retain(buckets[bucket_key], + communication_data_type, + numel_per_bucket=self.reduce_bucket_size, + rank=dst, + divide=False, + process_group=process_group) + + ############################################################################## + ############################# CPU Offload Methods############################# + ############################################################################## + def get_grad_position(self, group_id, tensor_list, first_offset, partition_size): + current_offset = 0 + + for i, tensor in enumerate(tensor_list): + param_id = self.get_param_id(tensor) + param_start_offset = 0 + + num_elements = tensor.numel() + + # we need to offset to get to the right element + if i == 0 and first_offset > 0: + tensor_offset = first_offset + num_elements = num_elements - tensor_offset + param_start_offset = first_offset + + # we dont need all elements of the tensor + if num_elements > (partition_size - current_offset): + num_elements = partition_size - current_offset + + self.grad_position[param_id] = [ + int(group_id), int(param_start_offset), + int(current_offset), int(num_elements) + ] + current_offset += num_elements + + def update_offload_overflow_tracker(self, grad): + if grad is not None and self._has_inf_or_nan(grad.data): + self.local_overflow = True + + def update_offload_overflow_tracker_for_param_grad(self, param): + grad_accum = self.get_param_gradient_attribute(param) + self.update_offload_overflow_tracker(grad_accum) + + def _get_offload_gradient_dict(self): + for param_group_index, _ in enumerate(self.optimizer.param_groups): + self.offload_gradient_dict[param_group_index] = [] + for lp_param in self.params_in_partition[param_group_index]: + param_id = self.get_param_id(lp_param) + [_, _, dest_offset, num_elements] = self.grad_position[param_id] + dest_tensor = self.single_partition_of_fp32_groups[param_group_index].grad.view(-1).narrow( + 0, dest_offset, num_elements) + self.offload_gradient_dict[param_group_index].append(dest_tensor) + + def async_accumulate_grad_in_cpu_via_gpu(self, param): + param_id = self.get_param_id(param) + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + # copy to a preexisiting buffer to avoid memory allocation penalty + dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel()) + + #buffer for storing gradients for this parameter in CPU + def buffer_to_accumulate_to_in_cpu(): + if not self.fp16_master_weights_and_gradients: + buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device) + return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer + else: + return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) + + #accumulate gradients into param.grad_accum or parts of it that belongs to this partition + def accumulate_gradients(): + grad_accum = self.get_param_gradient_attribute(param) + if not self.fp16_master_weights_and_gradients: + dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True) + grad_accum.data.view(-1).add_(dest_buffer) + else: + dest_buffer.narrow(0, source_offset, + num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1), + non_blocking=True) + grad_accum.data.view(-1).narrow(0, source_offset, + num_elements).add_(dest_buffer.narrow(0, source_offset, num_elements)) + + #move accumulated gradients back to CPU + def copy_gradients_to_cpu(): + grad_accum = self.get_param_gradient_attribute(param) + if not self.fp16_master_weights_and_gradients: + self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True) + else: + self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1).narrow( + 0, source_offset, num_elements), + non_blocking=True) + + if param_id not in self.accumulated_grads_in_cpu: + self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu() + + if self.micro_step_id > 0: + accumulate_gradients() + else: + copy_gradients_to_cpu() + + def set_norm_for_param_grad(self, param): + param_id = self.get_param_id(param) + grad_accum = self.get_param_gradient_attribute(param) + accumulated_grad = self.accumulated_grads_in_cpu[ + param_id] if self.gradient_accumulation_steps > 1 else grad_accum + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + start = source_offset + accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) + + self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) + + def set_norm_for_param_grad_in_gpu(self, param): + param_id = self.get_param_id(param) + grad_accum = self.get_param_gradient_attribute(param) + if grad_accum is None: + accumulated_grad = param.grad + else: + accumulated_grad = grad_accum + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + start = source_offset + accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) + + self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) + + def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param): + param_id = self.get_param_id(param) + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + dest_tensor = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) + + grad_accum = self.get_param_gradient_attribute(param) + if grad_accum is None: + src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) + else: + src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) + if not self.fp16_master_weights_and_gradients: + src_tensor = src_tensor.float() + + dest_tensor.copy_(src_tensor, non_blocking=True) + self.clear_grad_attribute(param) #offload only + + def complete_grad_norm_calculation_for_cpu_offload(self, params): + total_norm = 0.0 + norm_type = 2.0 + for p in params: + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + param_id = self.get_param_id(p) + # as some model have trainable parameters but skipped in training, + # their backward hooks in self.create_gradient_handling_hooks() will not run, + # so they have no norm_for_param_grads + if param_id in self.norm_for_param_grads: + param_norm = self.norm_for_param_grads[param_id] + total_norm += param_norm.item()**2 + else: + # As unused parameters in modules may not be expected sometimes, + # add an explicit error msg when it occurred and an option to + # avoid the error + assert self.ignore_unused_parameters, """ + This assert indicates that your module has parameters that + were not used in producing loss. + You can avoid this assert by + (1) enable ignore_unused_parameters option in zero_optimization config; + (2) making sure all trainable parameters and `forward` function + outputs participate in calculating loss. + """ + + # Sum across all model parallel GPUs. + total_dev_norm = get_accelerator().FloatTensor([float(total_norm)]) + dist.all_reduce(total_dev_norm, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_dev_norm, op=dist.ReduceOp.SUM) + + total_norm = total_dev_norm[0].item()**(1. / norm_type) + + if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: + total_norm = -1.0 + + return torch.tensor(total_norm, device=self.device, dtype=torch.float) + + ############################################################################################ + def copy_grads_in_partition(self, param): + if self.cpu_offload: + + if self.gradient_accumulation_steps > 1: + self.async_accumulate_grad_in_cpu_via_gpu(param) + + if self.is_gradient_accumulation_boundary: + self.set_norm_for_param_grad_in_gpu(param) + + self.update_offload_overflow_tracker_for_param_grad(param) + + self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param) + + return + #print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}") + if self.grads_in_partition is None: + self.grads_in_partition_offset = 0 + total_size = 0 + for group in self.params_in_partition: + for param_in_partition in group: + total_size += param_in_partition.numel() + + see_memory_usage(f"before copying {total_size} gradients into partition") + self.grads_in_partition = torch.empty(int(total_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + see_memory_usage(f"after copying {total_size} gradients into partition") + + grad_reduc = self.get_gradient_for_reduction(param) + # The allreduce buffer will be rewritten. Copy the gradients in partition to a new buffer + new_grad_tensor = self.grads_in_partition.view(-1).narrow(0, self.grads_in_partition_offset, param.numel()) + new_grad_tensor.copy_(grad_reduc.view(-1)) + grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) + #print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}") + self.grads_in_partition_offset += param.numel() + + def reduce_ipg_grads(self): + for comm_dtype in sort_dtypes(self.ipg_buckets.keys()): + bucket = self.ipg_buckets[comm_dtype] + + if self.contiguous_gradients: + if comm_dtype in self.extra_large_param_to_reduce: + assert len(bucket.params) == 1, "more than 1 param in ipg bucket, this shouldn't happen" + _, _, param_id = bucket.params[0] + assert self.get_param_id(self.extra_large_param_to_reduce[comm_dtype] + ) == param_id, "param in ipg bucket does not match extra-large param" + extra_large_grad_reduc = self.get_gradient_for_reduction( + self.extra_large_param_to_reduce[comm_dtype]) + self.average_tensor(extra_large_grad_reduc.view(-1), comm_dtype) + del self.extra_large_param_to_reduce[comm_dtype] + else: + self.average_tensor(bucket.buffer[bucket.index].narrow(0, 0, bucket.elements), comm_dtype) + else: + self.buffered_reduce_fallback(None, bucket.grads, comm_dtype, elements_per_buffer=bucket.elements) + + if self.overlap_comm: + stream = self.reduction_stream + elif self.cpu_offload: + # TODO: copy_grad_stream is disabled because of race with reduce. This hurts perf and should be fixed. + # get_accelerator().synchronize() + # stream = self.copy_grad_stream + stream = get_accelerator().current_stream() + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + for comm_dtype in sort_dtypes(self.ipg_buckets.keys()): + bucket = self.ipg_buckets[comm_dtype] + + for group_idx, param_idx_in_group, param_id in bucket.params: + param = self.bit16_groups[group_idx][param_idx_in_group] + + assert self.params_already_reduced[param_id] == False, \ + f"The parameter {param_id} has already been reduced. \ + Gradient computed twice for this partition. \ + Multiple gradient reduction is currently not supported" + + self.params_already_reduced[param_id] = True + if self.partition_gradients: + if not self.is_param_in_current_partition[param_id]: + if self.overlap_comm and self.contiguous_gradients is False: + # Clear grads of other partitions during the next reduction + # to avoid clearing them before the reduction is complete. + self.previous_reduced_grads[comm_dtype].append(param) + else: + self.clear_grad_attribute(param) + elif self.contiguous_gradients: + self.copy_grads_in_partition(param) + else: # zero stage 1 - partition only optimizer state + if self.contiguous_gradients and self.is_param_in_current_partition[param_id]: + self.copy_grads_in_partition(param) + bucket.clear() + ##################################################################### + + def process_gradients(self, param, i): + self.backward_prologue() + if self.use_grad_accum_attribute: + self._fill_param_grad_accum_attribute(param) + if self.partition_gradients or self.overlap_comm: + self.reduce_ready_partitions_and_remove_grads(param, i) + + def reduce_ready_partitions_and_remove_grads(self, param, i): + if self.partition_gradients or self.is_gradient_accumulation_boundary: + self.reduce_independent_p_g_buckets_and_remove_grads(param, i) + + def zero_reduced_gradients(self, partition_id, i): + + def are_all_related_partitions_reduced(params_id): + for partition_id in self.param_to_partition_ids[i][params_id]: + if not self.is_partition_reduced[i][partition_id]: + return False + return True + + for params_id in self.is_grad_computed[i][partition_id]: + if are_all_related_partitions_reduced(params_id): + self.param_dict[params_id].grad = None # dead code + + def flatten_and_print(self, message, tensors, start=0, n=5): + flatten_tensor = self.flatten(tensors) + + def print_func(): + logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) + + self.sequential_execution(print_func, message) + + def get_grads_to_reduce(self, i, partition_id): + + def get_reducible_portion(key): + grad = self.param_dict[key].grad + total_elements = grad.numel() + start = self.grad_start_offset[i][partition_id][key] + num_elements = min(total_elements - start, + self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) + if not pg_correctness_test: + if num_elements == total_elements: + return grad + else: + return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) + else: + if num_elements == total_elements: + return grad.clone() + else: + return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) + + grads_to_reduce = [] + for key in self.is_grad_computed[i][partition_id]: + grad = get_reducible_portion(key) + grads_to_reduce.append(grad) + return grads_to_reduce + + def sequential_execution(self, function, message, group=None): + if group is None: + group = self.dp_process_group + if dist.get_rank(group=group) == 0: + logger.info(message) + for id in range(dist.get_world_size(group=group)): + if id == dist.get_rank(group=group): + function() + dist.barrier(group=group) + + def set_none_gradients_to_zero(self, i, partition_id): + for param_id in self.is_grad_computed[i][partition_id]: + param = self.param_dict[param_id] + if param.grad is None: + param.grad = torch.zeros_like(param) + + ######################Reduction Related Methods############################## + def allreduce_bucket(self, + bucket, + communication_data_type: torch.dtype, + rank=None, + log=None, + divide=True, + process_group=None): + + tensor = self.flatten(bucket) + + process_group = self.dp_process_group if process_group is None else process_group + + tensor_to_allreduce = tensor + + if pg_correctness_test or self.sequence_parallel_size > 1: + communication_data_type = torch.float32 + + if communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(communication_data_type) + + if divide: + tensor_to_allreduce.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size)) + + if rank is None: + # "All Reducing" + dist.all_reduce(tensor_to_allreduce, group=process_group) + else: + global_rank = dist.get_global_rank(process_group, rank) + dist.reduce(tensor_to_allreduce, global_rank, group=process_group) + + if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + if rank is None or rank == dist.get_rank(group=process_group): + tensor.copy_(tensor_to_allreduce) + + return tensor + + def _clear_previous_reduced_grads(self): + for dtype in self.previous_reduced_grads: + for param in self.previous_reduced_grads[dtype]: + self.clear_grad_attribute(param) + self.previous_reduced_grads[dtype].clear() + + # if rank is specified do a reduction instead of an allreduce + def allreduce_and_copy(self, + small_bucket, + communication_data_type: torch.dtype, + rank=None, + log=None, + divide=True, + process_group=None): + process_group = self.dp_process_group if process_group is None else process_group + if self.overlap_comm: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + # It is safe to clear the previously reduced grads of other partitions + self._clear_previous_reduced_grads() + stream = self.reduction_stream + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + allreduced = self.allreduce_bucket( + small_bucket, + communication_data_type, + rank=rank, + log=log, + divide=divide, + process_group=process_group, + ) + if rank is None or rank == dist.get_rank(group=self.dp_process_group): + for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + buf.copy_(synced) + + def allreduce_no_retain( + self, + bucket, + communication_data_type: torch.dtype, + numel_per_bucket=500000000, + rank=None, + log=None, + divide=True, + process_group=None, + ): + small_bucket = [] + numel = 0 + for tensor in bucket: + small_bucket.append(tensor) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy(small_bucket, + communication_data_type, + rank=rank, + log=None, + divide=divide, + process_group=process_group) + small_bucket = [] + numel = 0 + + if len(small_bucket) > 0: + self.allreduce_and_copy(small_bucket, + communication_data_type, + rank=rank, + log=log, + divide=divide, + process_group=process_group) + + # allows using reduction of gradients instead of using all_reduce + + def buffered_reduce_fallback(self, + rank, + grads, + communication_data_type: torch.dtype, + elements_per_buffer=500000000, + log=None): + split_buckets = split_half_float_double(grads) + + for i, bucket in enumerate(split_buckets): + self.allreduce_no_retain(bucket, + communication_data_type, + numel_per_bucket=elements_per_buffer, + rank=rank, + log=log) + + ############################################################################# + ############################################################################# + ############################################################################# + + # views the tensor as multiple partitions and returns + # those partitions + def get_data_parallel_partitions(self, tensor, group_id): + partitions = [] + + dp = dist.get_world_size(group=self.real_dp_process_group[group_id]) + # dp_id = dist.get_rank(group=self.real_dp_process_group[group_id]) + + total_num_elements = tensor.numel() + + base_size = total_num_elements // dp + remaining = total_num_elements % dp + + start = 0 + for id in range(dp): + partition_size = base_size + if id < remaining: + partition_size = partition_size + 1 + partitions.append(tensor.narrow(0, start, partition_size)) + start = start + partition_size + return partitions + + def get_partition_info(self, tensor_list, partition_size, partition_id): + params_in_partition = [] + params_not_in_partition = [] + + start_index = partition_size * partition_id + end_index = partition_size * (partition_id + 1) + + current_index = 0 + first_offset = 0 + + for tensor in tensor_list: + + tensor_size = tensor.numel() + + if start_index <= current_index < end_index: + params_in_partition.append(tensor) + + elif current_index < start_index < (current_index + tensor_size): + params_in_partition.append(tensor) + + assert (first_offset == 0 + ), "This can happen either zero or only once as this must be the first tensor in the partition" + first_offset = start_index - current_index + + else: + params_not_in_partition.append(tensor) + + current_index = current_index + tensor_size + + return params_in_partition, params_not_in_partition, first_offset + + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + # FP32 grad should never exist. + # For speed, set model fp16 grad to None by default + # zero all pointers to grad tensors + for group in self.bit16_groups: + for p in group: + if set_to_none: + p.grad = None # epilogue and in step + p.grad_accum = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def _model_parallel_all_reduce(self, tensor, op): + """ Perform all reduce within model parallel group, if any. + """ + if self.model_parallel_group is None or self.model_parallel_world_size == 1: + pass + else: + dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) + + def get_grad_norm_direct(self, gradients, params, norm_type=2): + """Clips gradient norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float or int): max norm of the gradients + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + """ + norm_type = float(norm_type) + all_norms = [] + if norm_type == inf: + for g in gradients: + all_norms.append(g.data.abs().max().float()) + total_norm = torch.stack(all_norms).max() + dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + # Take max across all GPUs. + self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.MAX) + else: + # if dist.get_rank() == 0: + # logger.info(f"Total Norm beginning {total_norm}") + for g, p in zip(gradients, params): + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + all_norms.append( + torch.linalg.vector_norm(g.data.double().detach(), + ord=norm_type).to(get_accelerator().current_device_name())) + if len(all_norms) > 0: + total_norm = torch.stack(all_norms).square().sum().float() + else: + total_norm = torch.tensor(0.0, dtype=torch.float32).to(self.device) + # Sum across all model parallel Device. + dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.SUM) + + total_norm = total_norm.pow(1. / norm_type) + + mask_nan_or_inf_with_val_inplace(total_norm, device=self.device) + + return total_norm + + # creates a flat fused tensor from the tensor list starting at the first_offset + # in the first tensor of the list. If there are not enough elements in the tensor + # list then the flat tensor will be padded with zeros + def get_flat_partition(self, tensor_list, first_offset, partition_size, dtype, device, return_tensor_list=False): + flat_tensor_list = [] + current_size = 0 + + for i, tensor in enumerate(tensor_list): + grad_accum = self.get_param_gradient_attribute(tensor) + if grad_accum is None: + grad_accum = torch.zeros_like(tensor, dtype=dtype) + + tensor = grad_accum + num_elements = tensor.numel() + tensor_offset = 0 + + # we need to offset to get to the right element + if i == 0 and first_offset > 0: + tensor_offset = first_offset + num_elements = num_elements - tensor_offset + + # we dont need all elements of the tensor + if num_elements > (partition_size - current_size): + num_elements = partition_size - current_size + + # we need a narrow view of the tensor based on the tensor offset and number of elements that + # we need from this tensor + if tensor_offset > 0 or num_elements < tensor.numel(): + flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) + else: + flat_tensor_list.append(tensor) + + current_size = current_size + num_elements + + # this means its the last partition and does not align with the dp boundary. We need to pad before flattening + if current_size < partition_size: + flat_tensor_list.append(torch.zeros(int(partition_size - current_size), dtype=dtype, device=device)) + + if return_tensor_list: + return flat_tensor_list + + return self.flatten(flat_tensor_list) + + def free_grad_in_param_list(self, param_list): + for p in param_list: + p.grad = None # in step + p.grad_accum = None + + def reset_cpu_buffers(self): + self.norm_for_param_grads = {} + self.local_overflow = False + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + def scaled_global_norm(self, norm_type=2): + assert norm_type == 2, "only L2 norm supported" + norm_groups = [] + for i, group in enumerate(self.bit16_groups): + if self.cpu_offload: + norm = self.complete_grad_norm_calculation_for_cpu_offload(self.params_in_partition[i]) + norm_groups.append(norm) + else: + norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.params_in_partition[i])) + + if self.has_moe_layers: + self._average_expert_grad_norms(norm_groups) + + # calculating L2 norm + return torch.linalg.vector_norm(torch.stack(norm_groups), ord=norm_type) + + def get_bit16_param_group(self, group_no): + bit16_partitions = self.parallel_partitioned_bit16_groups[group_no] + partition_id = dist.get_rank(group=self.real_dp_process_group[group_no]) + return [bit16_partitions[dist.get_rank(group=self.real_dp_process_group[group_no])]] + + def _optimizer_step(self, group_no): + original_param_groups = self.optimizer.param_groups + self.optimizer.param_groups = [original_param_groups[group_no]] + # Disabling this as the C++ side copy & synchronize is not working correctly + #from deepspeed.ops.adam import DeepSpeedCPUAdam + #if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half: + # self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)]) + #else: + # self.optimizer.step() + if self.torch_autocast_gradscaler: + self.torch_autocast_gradscaler.step(self.optimizer) + self.torch_autocast_gradscaler.update() + else: + self.optimizer.step() + self.optimizer.param_groups = original_param_groups + + # We need to link optimizer state after the first step() call + self._lazy_init_hp_params_optimizer_state() + + def step(self, closure=None): + """ + Not supporting closure. + """ + self.micro_step_id = INITIAL_MICRO_STEP_ID + + see_memory_usage(f"In step before checking overflow") + + # First compute norm for all group so we know if there is overflow + if self.check_grad_overflow: + self.check_overflow(partition_gradients=self.partition_gradients) + + prev_scale = self.loss_scale + self._update_scale(self.overflow) + if self.overflow: + see_memory_usage('After overflow before clearing gradients') + self.zero_grad(set_to_none=True) + if self.cpu_offload: + self.reset_cpu_buffers() + else: + for k in self.averaged_gradients.keys(): + self.averaged_gradients[k] = None + + see_memory_usage('After overflow after clearing gradients') + + for timer in OPTIMIZER_TIMERS: + self.timers(timer).start() + self.timers(timer).stop() + return + + # Step 1:- Calculate gradient norm using bit-16 grads + see_memory_usage('Before norm calculation') + scaled_global_grad_norm = self.scaled_global_norm() + self._global_grad_norm = scaled_global_grad_norm / prev_scale + see_memory_usage('After norm before optimizer') + + # Step 2:- run optimizer and upscaling simultaneously + for i, group in enumerate(self.bit16_groups): + self.timers(OPTIMIZER_GRADIENTS_TIMER).start() + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + if self.cpu_offload: + single_grad_partition = self.single_partition_of_fp32_groups[i].grad + self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) + + self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() + self.timers(OPTIMIZER_STEP_TIMER).start() + self._optimizer_step(i) + + # Disabled, this is not currently working + #from deepspeed.ops.adam import DeepSpeedCPUAdam + #if not (type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half): + # bit16_partitions = self.parallel_partitioned_bit16_groups[i] + # fp32_partition = self.single_partition_of_fp32_groups[i] + # bit16_partitions[partition_id].data.copy_(fp32_partition.data) + bit16_partitions = self.parallel_partitioned_bit16_groups[i] + fp32_partition = self.single_partition_of_fp32_groups[i] + bit16_partition_buffer = self.param_buffer_of_bit16_for_cpu_offload_groups[i] + bit16_partition_buffer.data.copy_(fp32_partition.data) + bit16_partitions[partition_id].data.copy_(bit16_partition_buffer.data, non_blocking=True) + + self.timers(OPTIMIZER_STEP_TIMER).stop() + else: + # free gradients for all the parameters that are not updated by this process(ZeRO stage2) + self.free_grad_in_param_list(self.params_not_in_partition[i]) + + # create a flat gradients for parameters updated by this process + # If we are last partition, ensure we have same size grads and partition size, if not pad with zero tensors + if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: + single_grad_partition = self.flatten_dense_tensors_aligned( + self.averaged_gradients[i], + int(self.partition_size[i])).to(self.single_partition_of_fp32_groups[i].dtype) + else: + single_grad_partition = self.flatten(self.averaged_gradients[i]).to( + self.single_partition_of_fp32_groups[i].dtype) + assert single_grad_partition.numel() == self.partition_size[i], \ + "averaged gradients have different number of elements that partition size {} {} {} {}".format( + single_grad_partition.numel(), self.partition_size[i], i, partition_id) + + self.single_partition_of_fp32_groups[i].grad = single_grad_partition + # release all the gradient since we have already created a necessary copy in dp_grad_partition(ZeRO stage2) + self.free_grad_in_param_list(self.params_in_partition[i]) + + self.averaged_gradients[i] = None + + self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) + + self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() + + # Step 3:- run the optimizer if no offloading + self.timers(OPTIMIZER_STEP_TIMER).start() + self._optimizer_step(i) + # Step 4:- get rid of the fp32 gradients. Not needed anymore + self.single_partition_of_fp32_groups[i].grad = None + del single_grad_partition + bit16_partitions = self.parallel_partitioned_bit16_groups[i] + fp32_partition = self.single_partition_of_fp32_groups[i] + bit16_partitions[partition_id].data.copy_(fp32_partition.data) + self.timers(OPTIMIZER_STEP_TIMER).stop() + + see_memory_usage('After optimizer before all-gather') + if self.cpu_offload: + self.reset_cpu_buffers() + + self.timers(OPTIMIZER_ALLGATHER_TIMER).start() + # Gather the updated weights from everyone. + # Then all partitions of the model parameters are updated and ready for next round forward. + all_gather_dp_groups(groups_flat=self.bit16_groups_flat, + partitioned_param_groups=self.parallel_partitioned_bit16_groups, + dp_process_group=self.real_dp_process_group, + start_alignment_factor=self.nccl_start_alignment_factor, + allgather_bucket_size=self.allgather_bucket_size) + self.timers(OPTIMIZER_ALLGATHER_TIMER).stop() + + # TODO: we probably don't need this? just to be safe + for i in range(len(self.bit16_groups)): + self._update_model_bit16_weights(i) + + self.timers.log(OPTIMIZER_TIMERS) + see_memory_usage('After zero_optimizer step') + + return + + @torch.no_grad() + def update_lp_params(self): + for i, (bit16_partitions, fp32_partition) in enumerate( + zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + bit16_partitions[partition_id].data.copy_(fp32_partition.data) + + all_gather_dp_groups(groups_flat=self.bit16_groups_flat, + partitioned_param_groups=self.parallel_partitioned_bit16_groups, + dp_process_group=self.real_dp_process_group, + start_alignment_factor=self.nccl_start_alignment_factor, + allgather_bucket_size=self.allgather_bucket_size) + + def _average_expert_grad_norms(self, norm_groups): + for i, norm in enumerate(norm_groups): + if self.is_moe_param_group[i]: + scaled_norm_tensor = norm * 1.0 / dist.get_world_size(group=self.real_dp_process_group[i]) + if self.device == 'cpu': + scaled_norm_tensor = scaled_norm_tensor.to(get_accelerator().current_device_name()) + dist.all_reduce(scaled_norm_tensor, group=self.real_dp_process_group[i]) + norm_groups[i] = scaled_norm_tensor.to(self.device) + + def unscale_and_clip_grads(self, grad_groups_flat, total_norm): + # compute combined scale factor for this group + combined_scale = self.loss_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad + clip = torch.clamp(clip, min=1.0) + combined_scale = clip * self.loss_scale + + for grad in grad_groups_flat: + if isinstance(grad, list): + sub_partitions = grad + for g in sub_partitions: + g.data.mul_(1. / combined_scale) + else: + grad.data.mul_(1. / combined_scale) + + def _check_overflow(self, partition_gradients=True): + self.overflow = self.has_overflow(partition_gradients) + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params): + invalid_grad_count = torch.zeros([1], dtype=torch.float, device=get_accelerator().current_device_name()) + for p in params: + if p.grad is not None: + invalid_grad_count += self._has_inf_or_nan(p.grad) + return invalid_grad_count.bool() + + def has_overflow_partitioned_grads_serial(self): + invalid_grad_count = torch.zeros([1], dtype=torch.float, device=get_accelerator().current_device_name()) + for i in range(len(self.bit16_groups)): + for j, grad in enumerate(self.averaged_gradients[i]): + if grad is not None: + invalid_grad_count += self._has_inf_or_nan(grad) + return invalid_grad_count.bool() + + def has_overflow(self, partition_gradients=True): + overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial() + overflow_gpu = get_accelerator().ByteTensor([overflow]) if self.cpu_offload else overflow.byte().to( + get_accelerator().current_device_name()) + + if partition_gradients: + '''This will capture overflow across all data parallel and expert parallel process + Since expert parallel process are a subset of data parallel process''' + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + # Since each model parallel GPU carries only part of the model, + # make sure overflow flag is synced across all the model parallel GPUs + self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) + + overflow = overflow_gpu[0].item() + return bool(overflow) + + # `x` is a torch.Tensor + @staticmethod + def _has_inf_or_nan(x, j=None): + float_x = x.float() + nan = float_x.isnan() + inf = float_x.isinf() + inf_or_nan = nan.logical_or(inf) + return inf_or_nan.float().max() + + def backward_prologue(self): + if not self.ready_for_gradients: + self.micro_step_id += 1 + + if self.contiguous_gradients: + for _, bucket in self.ipg_buckets.items(): + bucket.buffer.clear() + + # Buffer's dtype is the same as the dtype of optimizer, not dtype for autocast + buf_0 = torch.empty(int(self.reduce_bucket_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + bucket.buffer.append(buf_0) + bucket.index = 0 + + # Use double buffers to avoid data access conflict when overlap_comm is enabled. + if self.overlap_comm: + for _, bucket in self.ipg_buckets.items(): + buf_1 = torch.empty(int(self.reduce_bucket_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + bucket.buffer.append(buf_1) + + self.ready_for_gradients = True + + def backward_epilogue(self): + # Only for Stage 1, Mode 2 + if self.use_grad_accum_attribute: + self.fill_grad_accum_attribute() + + def backward(self, loss, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + elif self.torch_autocast_gradscaler: + self.torch_autocast_gradscaler.scale(loss).backward(retain_graph=retain_graph) + else: + self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) + self.backward_epilogue() + + def check_overflow(self, partition_gradients=True): + self._check_overflow(partition_gradients) + + def _update_scale(self, has_overflow=False): + self.loss_scaler.update_scale(has_overflow) + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + + param_groups = property(_get_param_groups, _set_param_groups) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.loss_scaler.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) + cur_scale = property(_get_loss_scale, _set_loss_scale) + + # Return group tensor after removing paddings that are added for alignment to DP world size. + # This method works on the assumption that each group contains a single flattened tensor. + def _get_groups_without_padding(self, groups_with_padding): + groups_without_padding = [] + for i, group in enumerate(groups_with_padding): + lean_length = group.numel() - self.groups_padding[i] + groups_without_padding.append(group[:lean_length]) + + return groups_without_padding + + # Return optimizer state after removing paddings that are added for alignment. + def _get_state_without_padding(self, state_with_padding, padding): + lean_state = {} + for key, value in state_with_padding.items(): + if torch.is_tensor(value): + lean_length = value.numel() - padding + lean_state[key] = value[:lean_length] + else: + lean_state[key] = value + + return lean_state + + # Return base optimizer states. + # This method assumes that each param group contains a single flattened tensor. + def _get_base_optimizer_state(self): + optimizer_groups_state = [] + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + lean_optimizer_state = self._get_state_without_padding(self.optimizer.state[p], self.groups_padding[i]) + optimizer_groups_state.append(lean_optimizer_state) + + return optimizer_groups_state + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + state_dict = {} + state_dict[LOSS_SCALER] = self.loss_scaler + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['overflow'] = self.overflow + state_dict[CLIP_GRAD] = self.clip_grad + + if self.elastic_checkpoint: + state_dict[BASE_OPTIMIZER_STATE] = self._get_base_optimizer_state() + + if "step" in self.optimizer.param_groups[0]: + # Assuming "step" is the only item that changes through training iterations + assert all(group["step"] == self.optimizer.param_groups[0]["step"] + for group in self.optimizer.param_groups), "All param groups must have the same step value" + state_dict[BASE_OPTIMIZER_STATE_STEP] = self.optimizer.param_groups[0]["step"] + else: + state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() + + # Remove paddings for DP alignment to enable loading for other alignment values + fp32_groups_without_padding = self._get_groups_without_padding(self.single_partition_of_fp32_groups) + state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding + + state_dict[ + ZERO_STAGE] = ZeroStageEnum.gradients if self.partition_gradients else ZeroStageEnum.optimizer_states + state_dict[GROUP_PADDINGS] = self.groups_padding + state_dict[PARTITION_COUNT] = self.partition_count + + state_dict[DS_VERSION] = version + state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings + + return state_dict + + # Restore base optimizer fp32 weights from elastic checkpoint by: + # 1) Merging fp32 weights from checkpoints of all partitions + # 2) Extracting fp32 weights for current partition from merged weights + # 3) Using extracted weights to update base optimizer weights directly. + def _restore_from_elastic_fp32_weights(self, all_state_dict): + merged_single_partition_of_fp32_groups = [] + + for i in range(len(self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + merged_partitions = [sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict] + if self.is_moe_group(self.optimizer.param_groups[i]): + ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) + merged_partitions = [merged_partitions[i] for i in ranks] + flat_merged_partitions = self.flatten_dense_tensors_aligned( + merged_partitions, + self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i])) + dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, i) + merged_single_partition_of_fp32_groups.append(dp_partitions[partition_id]) + + for current, saved in zip(self.single_partition_of_fp32_groups, merged_single_partition_of_fp32_groups): + current.data.copy_(saved.data) + + # Restore base optimizer fp32 weights from ZeRO fp16 or bfloat16 weights + def _restore_from_bit16_weights(self): + for group_id, (bit16_partitions, fp32_partition) in enumerate( + zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) + fp32_partition.data.copy_(bit16_partitions[partition_id].data) + + # Refresh the fp32 master params from the fp16 or bfloat16 copies. + def refresh_fp32_params(self): + self._restore_from_bit16_weights() + + # Extract optimizer state for current partition from merged states of all partitions + def _partition_base_optimizer_state(self, state_key, all_partition_states, group_id): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) + alignment = self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[group_id]) + if torch.is_tensor(all_partition_states[0]): + flat_merged_partitions = self.flatten_dense_tensors_aligned(all_partition_states, alignment) + dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, group_id) + return dp_partitions[partition_id] + else: + # Assume non-tensor states are not partitioned and equal across ranks, so return first one + return all_partition_states[0] + + def _restore_step_from_elastic_checkpoint(self, all_state_dict): + assert BASE_OPTIMIZER_STATE_STEP in all_state_dict[0] + assert all(sd[BASE_OPTIMIZER_STATE_STEP] == all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] + for sd in all_state_dict), "State dicts of all partitions must have the same step value" + return all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] + + def _restore_base_optimizer_state(self, base_optimizer_group_states, base_optimizer_state_step, group_paddings): + if type(base_optimizer_group_states) == dict: + base_optimizer_group_states = base_optimizer_group_states['state'] + + saved_keys = base_optimizer_group_states[0].keys() + + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + padding = 0 if group_paddings is None else group_paddings[i] + for key in saved_keys: + saved = base_optimizer_group_states[i][key] + + if torch.is_tensor(saved): + if key in self.optimizer.state[p]: + dst_tensor = self.optimizer.state[p][key] + src_tensor = _get_padded_tensor(saved, dst_tensor.numel()) + self.optimizer.state[p][key].data.copy_(src_tensor.data) + else: + self.optimizer.state[p][key] = _pad_tensor_by_size( + saved, padding, torch.float32, + torch.device('cpu') if self.cpu_offload else self.device) + else: + self.optimizer.state[p][key] = saved + + for param_group in self.optimizer.param_groups: + param_group['step'] = base_optimizer_state_step + + def get_ep_ranks(self, rank=0, group_name=None): + from deepspeed.utils import groups + expert_parallel_size_ = groups._get_expert_parallel_world_size(group_name) + world_size = groups._get_data_parallel_world_size() + rank = groups._get_expert_parallel_rank(group_name) + ranks = range(rank, world_size, expert_parallel_size_) + return list(ranks) + + # Restore base optimizer state from elastic checkpoint by + # 1) Merging optimizer state from checkpoints of all partitions + # 2) Extracting optimizer state for current partition from the merged state + # 3) Using the extracted value to directly update the base optimizer. + def _restore_elastic_base_optimizer_state(self, all_state_dict): + base_optimizer_group_states = [] + for i in range(len(self.optimizer.param_groups)): + partition_states = {} + all_partition_group_states = [sd[BASE_OPTIMIZER_STATE][i] for sd in all_state_dict] + + if self.is_moe_group(self.optimizer.param_groups[i]): + ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) + all_partition_group_states = [all_partition_group_states[i] for i in ranks] + + for key in all_partition_group_states[0].keys(): + all_partition_states = [all_states[key] for all_states in all_partition_group_states] + partition_states[key] = self._partition_base_optimizer_state(key, all_partition_states, i) + base_optimizer_group_states.append(partition_states) + + self._restore_base_optimizer_state(base_optimizer_group_states, + self._restore_step_from_elastic_checkpoint(all_state_dict), None) + + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None, + param_shapes=None): + if checkpoint_folder: + self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) + else: + self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) + + def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): + self.load_hp_checkpoint_state_from_checkpoint_dir("bit16_groups", checkpoint_folder) + + def _load_global_state(self, sd): + self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) + self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) + self.overflow = sd.get('overflow', self.overflow) + self.clip_grad = sd.get(CLIP_GRAD, self.clip_grad) + + ckpt_version = sd.get(DS_VERSION, False) + assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" + ckpt_version = pkg_version.parse(ckpt_version) + + # zero stage 1 mode + if not self.partition_gradients: + required_version = pkg_version.parse("0.3.17") + error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \ + "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \ + "please use an older version of DeepSpeed (<= 0.5.8) and set 'legacy_stage1': true in your zero config json." + assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}" + + def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): + r"""Loading ZeRO checkpoint + + Arguments: + state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. + Note that the number of saved partitions may differ from number of loading partitions to support + changing GPU count, specifically DP world size, between saving and loading checkpoints. + load_optimizer_states: Boolean indicating whether or not to load base optimizer states + load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 + copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). + """ + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + + # I think it should actually be ok to reload the optimizer before the model. + dp_rank = dist.get_rank(group=self.dp_process_group) + current_rank_sd = state_dict_list[dp_rank] + self._load_global_state(current_rank_sd) + + ckpt_is_rigid = isinstance(current_rank_sd[BASE_OPTIMIZER_STATE], dict) + + # padding is always at the last rank/partition + # if DP=1024 and param-group elems=16 -> padding will be 1024-16 across all but one rank + # scenario-1 (shrink): saving w. 4 gpus -> loading w. 2 gpus + # scenario-2 (expand): saving w. 2 gpus -> loading w. 4 gpus + # if load_optimizer_states: + # if new_dp_size: + # self.strip_padding() + # self.add_padding_w_new_dp_size() + # self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) + + if load_optimizer_states: + if ckpt_is_rigid: + # loading rigid ckpt into either rigid or elastic exec + self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) + else: + if self.elastic_checkpoint: + # loading elastic into elastic exec + self._restore_elastic_base_optimizer_state(state_dict_list) + else: + # loading an elastic checkpoint into rigid exec + self._restore_base_optimizer_state(current_rank_sd[BASE_OPTIMIZER_STATE], + current_rank_sd[BASE_OPTIMIZER_STATE_STEP], + current_rank_sd[GROUP_PADDINGS]) + + # At this point, the optimizer's references to the model's fp32 parameters are up to date. + # The optimizer's hyperparameters and internal buffers are also up to date. + # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still + # out of date. There are two options. + # 1: Refresh the master params from the model's fp16 params. + # This requires less storage but incurs precision loss. + # 2: Save and restore the fp32 master copies separately. + # We choose option 1 if changing DP degree and option 2 otherwise. + # + # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device + # of their associated parameters, because it's possible those buffers might not exist yet in + # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been + # constructed in the same way as the one whose state_dict we are loading, the same master params + # are guaranteed to exist, so we can just copy_() from the saved master params. + + if load_from_fp32_weights: + # option 2 from above + if self.elastic_checkpoint and not ckpt_is_rigid: + self._restore_from_elastic_fp32_weights(state_dict_list) + else: + # For non-elastic checkpoint, simply copying from saved weights of current rank is sufficient. + for current, saved in zip(self.single_partition_of_fp32_groups, + current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): + src_tensor = _get_padded_tensor(saved, current.numel()) + current.data.copy_(src_tensor.data) + else: + # option 1 from above + self._restore_from_bit16_weights() + + if load_optimizer_states: + self._link_all_hp_params() + + +def _handle_overflow(cpu_sum, x, i): + import math + rank = dist.get_rank() + if rank == 0: + t_i = -1 + for v_i, v in enumerate(x.data.contiguous().view(-1)): + if not math.isfinite(float(v)): + t_i = v_i + break + logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") + + +def estimate_zero2_model_states_mem_needs(total_params, + num_gpus_per_node=1, + num_nodes=1, + cpu_offload=True, + additional_buffer_factor=1.5): + + total_gpus = num_nodes * num_gpus_per_node + + if cpu_offload: + gpu_mem = 2 * total_params + cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor + else: + # GPU's total_params multipliers: 2 = params_16bit, + # 18 = 2_grads_16bit + 4_grads_32bit + 4_params_32bit + 8_optimizer_states_32bit(momentum and variance) + gpu_mem = 2 * total_params + int(18 * total_params / total_gpus) + cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor + + return int(cpu_mem), int(gpu_mem) + + +def model_to_params(model): + # shared params calculated only once + total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) + return total_params + + +def estimate_zero2_model_states_mem_needs_all_live(model, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients + for a given ``model`` and hardware setup. + + If you have an actual model object, use this function and everything will be derived + automatically. + + If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass + the ``total_params`` explicitly. + + Args: + - ``model``: ``nn.Module`` object + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + total_params = model_to_params(model) + + estimate_zero2_model_states_mem_needs_all_cold(total_params=total_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + additional_buffer_factor=additional_buffer_factor) + + +def estimate_zero2_model_states_mem_needs_all_cold(total_params, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients + for a given ``model`` and hardware setup. + + If it's a hypothetical model, use this function where you have to pass + the ``total_params`` and ``largest_layer_params`` explicitly. + + If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything + will be derived automatically. + + Args: + - ``total_params``: total model params + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + def format_options(cpu_offload): + enabled = [] + device = f'{OffloadDeviceEnum.cpu:4}' if cpu_offload else "none" + enabled.append(f"offload_optimizer={device}") + return ", ".join(enabled) + + nodes_str = "nodes" if num_nodes > 1 else "node" + gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" + print("Estimated memory needed for params, optim states and gradients for a:\n" + f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" + f"SW: Model with {int(total_params/1e6)}M total params.") + print(" per CPU | per GPU | Options") + for cpu_offload in [True, False]: + cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params=total_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + cpu_offload=cpu_offload, + additional_buffer_factor=additional_buffer_factor) + + options_str = format_options(cpu_offload=cpu_offload) + print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/test.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/test.py new file mode 100644 index 0000000000000000000000000000000000000000..1904a8d395be80a4322bc25bd1008a80bdd8f982 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/test.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.runtime.zero.contiguous_memory_allocator import ContiguousMemoryAllocator + + +def test1(): + mem = ContiguousMemoryAllocator(1024, torch.half, 'cpu') + mem.print_allocation(resolution=100) + a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0) + mem.print_allocation(resolution=100) + mem.release_tensor(a1) + mem.print_allocation(resolution=100) + a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0) + a3 = mem.allocate_tensor(256).mul_(0.0).add_(3.0) + a4 = mem.allocate_tensor(128).mul_(0.0).add_(4.0) + mem.print_allocation(resolution=100) + mem.release_tensor(a3) + mem.print_allocation(resolution=100) + a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0) + a6 = mem.allocate_tensor(256).mul_(0.0).add_(6.0) + a7 = mem.allocate_tensor(128).mul_(0.0).add_(7.0) + mem.print_allocation(resolution=100) + a8 = mem.allocate_tensor(256).mul_(0.0).add_(8.0) + a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0) + mem.print_allocation(resolution=100) + mem.release_tensor(a9) + mem.release_tensor(a6) + mem.release_tensor(a2) + mem.release_tensor(a5) + + a10 = mem.allocate_tensor(512).mul_(0.0).add_(10.0) + mem.print_allocation(resolution=100) + #print(f"a4:{a4}") + #print(f"a7:{a7}") + #print(f"a8:{a8}") + #print(f"a10:{a10}") + assert (a4.norm() + a7.norm() + a8.norm() + a10.norm()).item() == 474.50, "Test failed" + + +def test2(): + mem = ContiguousMemoryAllocator(512, torch.half, 'cpu') + a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0) + a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0) + a3 = mem.allocate_tensor(64).mul_(0.0).add_(3.0) + a4 = mem.allocate_tensor(64).mul_(0.0).add_(4.0) + a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0) + a6 = mem.allocate_tensor(64).mul_(0.0).add_(6.0) + a7 = mem.allocate_tensor(64).mul_(0.0).add_(7.0) + a8 = mem.allocate_tensor(64).mul_(0.0).add_(8.0) + mem.release_tensor(a2) + mem.release_tensor(a4) + mem.release_tensor(a6) + mem.release_tensor(a8) + mem.print_allocation(resolution=100) + + a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0) + a10 = mem.allocate_tensor(64).mul_(0.0).add_(10.0) + a11 = mem.allocate_tensor(64).mul_(0.0).add_(11.0) + mem.release_tensor(a1) + mem.release_tensor(a5) + mem.print_allocation(resolution=100) + a12 = mem.allocate_tensor(128).mul_(0.0).add_(12.0) + mem.print_allocation(resolution=100) + print(f"a7:{a7}") + print(f"a9:{a9}") + print(f"a10:{a10}") + print(f"a11:{a11}") + print(f"a12:{a12}") + assert (a7.norm() + a9.norm() + a10.norm() + a11.norm() + a12.norm()) == 460.75, "TestFailed" + + +test1() +test2() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/tiling.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/tiling.py new file mode 100644 index 0000000000000000000000000000000000000000..af1c5e54737a75cb66b39736718bf62169952a80 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/tiling.py @@ -0,0 +1,296 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import deepspeed +from deepspeed.runtime.utils import partition_uniform as partition + + +def split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False): + """Split a tensor along its last dimension. Adapted from Megatron-LM. + + Arguments: + tensor: input tensor. + partitions: list of partition sizes to supply to torch.split + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + + # Split. + tensor_list = torch.split(tensor, partitions, dim=last_dim) + # Note: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +class TiledLinear(torch.nn.Module): + + def __init__(self, + in_features, + out_features, + bias=True, + in_splits=1, + out_splits=1, + input_is_already_split=False, + combine_out_splits=True, + linear_cls=torch.nn.Linear, + init_linear=None, + **kwargs): + """A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce + memory requirements via tiling. + + TiledLinear breaks the input and output dimensions of a linear layer + into tiles that are processed in sequence. This class enables huge + linear layers when combined with ZeRO-3 because inactive tiles can be + partitioned and offloaded. + + .. note:: + We recommend using as few tiles as necessary. Tiling + significantly reduces memory usage, but can reduce throughput + for inexpensive layers. This due to the smaller kernels having + less parallelism and lower arithmetic intensity, while + introducing more frequent synchronization and communication. + + Args: + in_features (int): See ``torch.nn.Linear`` + out_features (int): See ``torch.nn.Linear`` + bias (bool, optional): See ``torch.nn.Linear`` + in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1. + out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1. + input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in + to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``. + combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs + into a single tensor. Defaults to ``True``. + linear_cls (class, optional): The underlying class to build individual tiles. + Defaults to ``torch.nn.Linear``. + init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of + ``init_linear``. Useful for debugging. Defaults to ``None``. + kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``. + + Raises: + RuntimeError: ``in_splits`` must be within the range [1, in_features). + RuntimeError: ``out_splits`` must be within the range of [1, out_features). + """ + + super().__init__() + + if (in_splits < 1) or (in_splits > in_features): + raise RuntimeError('in splits must be in range [1, in_features].') + if (out_splits < 1) or (out_splits > out_features): + raise RuntimeError('out splits must be in range [1, out_features].') + + # global, not necessarily local + self.in_features = in_features + self.out_features = out_features + self.use_bias = bias + + self.out_splits = out_splits + self.in_splits = in_splits + self.input_is_already_split = input_is_already_split + self.combine_out_splits = combine_out_splits + + # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features] + # For example, row_parts[p] gives the start of partition p and row_parts[p+1] + # is the exclusive end. + self.in_parts = partition(num_items=in_features, num_parts=in_splits) + self.out_parts = partition(num_items=out_features, num_parts=out_splits) + + assert len(self.out_parts) == out_splits + 1 + assert len(self.in_parts) == in_splits + 1 + assert self.out_parts[0] == 0 + assert self.out_parts[out_splits] == out_features + assert self.in_parts[in_splits] == in_features + + self.linears = torch.nn.ModuleList() + for out_id in range(out_splits): + self.linears.append(torch.nn.ModuleList()) + + local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id] + + for in_id in range(in_splits): + #if input_size is split, we only need one bias + local_bias = bias if in_id == (in_splits - 1) else False + + local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id] + local = linear_cls(local_in_dim, local_out_dim, bias=local_bias, **kwargs) + self.linears[out_id].append(local) + + # Optionally initialize with a known tensor + if init_linear is not None: + self.copy_params_from(init_linear) + + def forward(self, input_): + if self.in_splits > 1 and not self.input_is_already_split: + input_parts = partition(input_.shape[-1], self.in_splits) + split_sizes = [input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)] + inputs = self._split_global_input(input_, split_sizes) + elif self.in_splits > 1: + inputs = input_ + assert len( + inputs) == self.in_splits, f"Col splits {self.in_splits} does not match input splits {len(inputs)}" + else: + # no splits + inputs = [input_] + + outputs = [None] * self.out_splits + for out_id in range(self.out_splits): + for in_id in range(self.in_splits): + local_output = self.linears[out_id][in_id](inputs[in_id]) + + outputs[out_id] = self._reduce_local_output(in_id=in_id, + out_id=out_id, + current_out=outputs[out_id], + new_out=local_output) + + if self.combine_out_splits: + return self._combine_output_splits(outputs) + + return outputs + + def _split_global_input(self, input, split_sizes): + """Partition an input tensor along the last dimension, aligned with given splits. + + Subclasses should override this method to account for new input types. + + Args: + input (List[Tensor]): The tensor to partition along the last dimension. + split_sizes (List[int]): The size of each partition. + + Returns: + List[Any]: A list of the chunks of ``input``. + """ + return split_tensor_along_last_dim(input, split_sizes) + + def _reduce_local_output(self, in_id, out_id, current_out, new_out): + """Reduce (sum) a new local result into the existing local results. + + Subclasses should override this method. + + For a given ``out_id``, this method is called ``in_id-1`` times. The first input + split is a simple assignment. + + Args: + in_id (int): The input split that produced ``new_out``. + out_id (int): The output split that produced ``new_out``. + current_out (Any): The reduced form of all previous ``out_id`` results. + new_out (Any): The local result from forward (``in_id``, ``out_id``)e + + Returns: + Any: The combined result of ``current_out`` and ``new_out``. + """ + + if current_out is None: + #this clone is necessary to preserve auto grad + #there is some issue with inplace update for outputs that are views + return new_out.clone() + else: + return current_out + new_out + + def _combine_output_splits(self, outputs): + """Join the splits of the output into a single result. + + Args: + outputs (List[Any]): The reduced outputs for each output split. + + Returns: + Any: The combined outputs. + """ + assert len(outputs) == self.out_splits + return torch.cat(outputs, dim=-1) + + @torch.no_grad() + def copy_params_from(self, other): + """Copy the weight and bias data from ``other``. + + This is especially useful for reproducible initialization and testing. + + Equivalent to: + + .. code-block:: python + + with torch.no_grad(): + self.weight.copy_(other.weight) + if self.bias is not None: + self.bias.copy_(other.bias) + + .. note:: + If ZeRO-3 is enabled, this is a collective operation and the + updated parameters of data-parallel rank 0 will be visible on all + ranks. See :class:`deepspeed.zero.GatheredParameters` for more + information. + + + Args: + other (``torch.nn.Linear``): the linear layer to copy from. + """ + assert hasattr(other, 'weight') + assert other.weight.size() == (self.out_features, self.in_features) + if self.use_bias: + assert hasattr(other, 'bias') + assert other.bias is not None + assert other.bias.size() == (self.out_features, ) + else: + assert other.bias is None + + for row in range(self.out_splits): + rstart = self.out_parts[row] + rstop = self.out_parts[row + 1] + + for col in range(self.in_splits): + cstart = self.in_parts[col] + cstop = self.in_parts[col + 1] + + local = self.linears[row][col] + global_weight = other.weight[rstart:rstop, cstart:cstop] + with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0): + local.weight.copy_(global_weight) + + if local.bias is not None: + with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0): + local.bias.data.copy_(other.bias[rstart:rstop].data) + + +class TiledLinearReturnBias(TiledLinear): + """Wrapper for a Linear class that returns its own bias parameter, such as + used by Megatron-LM. + """ + + def _reduce_local_output(self, in_id, out_id, current_out, new_out): + """Reduces output tensors, but not the returned bias. """ + if current_out is not None: + old_tensor, old_bias = current_out + else: + old_tensor, old_bias = None, None + + assert isinstance(new_out, tuple) + assert len(new_out) == 2 + + tensor, bias = new_out + assert tensor is not None + + tensor = super()._reduce_local_output(in_id=in_id, out_id=out_id, current_out=old_tensor, new_out=tensor) + + if bias is None: + bias = old_bias + + return tensor, bias + + def _combine_output_splits(self, outputs): + # stack output tensors + tensors = [o[0] for o in outputs] + tensor = super()._combine_output_splits(tensors) + + # stack biases if applicable + biases = [o[1] for o in outputs if o[1] is not None] + if len(biases) > 0: + bias = super()._combine_output_splits(biases) + else: + bias = None + + return tensor, bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2d1cf17962d8a59fee5cb1ac790129fcb6c53a10 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/utils.py @@ -0,0 +1,175 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from typing import List, Tuple + +import torch +from deepspeed import comm as dist +from deepspeed.utils import logger +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad +from deepspeed.ops.adam import FusedAdam +from deepspeed.ops.lion import DeepSpeedCPULion, FusedLion +from deepspeed.utils.nvtx import instrument_w_nvtx +from deepspeed.accelerator import get_accelerator + +# ensure we only warn once, otherwise every iteration will trigger a warning +warned = False + + +def _initialize_parameter_parallel_groups(parameter_parallel_size=None): + data_parallel_size = int(dist.get_world_size()) + parameter_parallel_size = parameter_parallel_size or data_parallel_size + logger.info("data_parallel_size: %s, parameter_parallel_size: %s", data_parallel_size, parameter_parallel_size) + assert data_parallel_size % parameter_parallel_size == 0, \ + 'world size should be divisible by parameter parallel size' + rank = dist.get_rank() + my_group = None + for i in range(data_parallel_size // parameter_parallel_size): + ranks = range(i * parameter_parallel_size, (i + 1) * parameter_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + my_group = group + return my_group + + +class ZeRORuntimeException(Exception): + pass + + +ZERO_SUPPORTED_OPTIMIZERS = [ + torch.optim.Adam, torch.optim.AdamW, FusedAdam, DeepSpeedCPUAdam, torch.optim.Adagrad, DeepSpeedCPUAdagrad, + DeepSpeedCPULion, FusedLion +] + +# Add apex FusedAdam to supported list if apex is installed +try: + import apex + if hasattr(apex, 'optimizers') and hasattr(apex.optimizers, 'FusedAdam'): + ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam) +except ImportError: + pass + + +def is_zero_supported_optimizer(optimizer): + if dist.get_rank() == 0: + logger.info(f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}') + return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS + + +def get_lst_from_rank0(lst: List[int]) -> None: + """ + NOTE: creates both communication and synchronization overhead so should be used + sparingly + """ + lst_tensor = torch.tensor( + lst if dist.get_rank() == 0 else [-1] * len(lst), + dtype=int, + device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])), + requires_grad=False, + ) + dist.broadcast(lst_tensor, src=0, async_op=False) + + return list(lst_tensor.cpu().numpy()) + + +@instrument_w_nvtx +def assert_ints_same_as_other_ranks(ints: List[int]) -> None: + """ + NOTE: creates both communication and synchronization overhead so should be + used sparingly + + takes a list of ints from each rank and ensures that they are the same + across ranks, throwing an exception if they are not. + """ + rank0_ints = get_lst_from_rank0(ints) + if ints != rank0_ints: + raise RuntimeError(f"disagreement between rank0 and rank{dist.get_rank()}: " + f"rank0: {rank0_ints}, rank{dist.get_rank()}: {ints}") + + +def is_builtin_type(obj): + # https://stackoverflow.com/a/17795199 + return obj.__class__.__module__ == '__builtin__' or obj.__class__.__module__ == "builtins" + + +def isinstance_namedtuple(obj: object) -> bool: + """ + Is this an instance of namedtuple/NamedTuple? + From: https://stackoverflow.com/a/62692640 + + Args: + obj (object): An object. + + Returns: + bool: True if namedtuple/NamedTuple else False. + """ + return isinstance(obj, tuple) and hasattr(obj, '_asdict') and hasattr(obj, '_fields') + + +def is_zero_param(parameter): + if not torch.is_tensor(parameter): + return False + return hasattr(parameter, 'ds_id') + + +def apply_to_tensors_only(function, value, warning_msg_fn=None): + """ + Apply `function` to every Tensor in `value`. + + Args: + functional: The function class to apply. + value (Any): Target object to apply `function` to. + + Returns: + Any: Output of `function`. + """ + if isinstance(value, (tuple, list)): + touched_outputs = [] + for elem in value: + touched_output = apply_to_tensors_only(function, elem) + touched_outputs.append(touched_output) + + if isinstance_namedtuple(value): + # namedtuples require a slightly different syntax. + return value.__class__(*touched_outputs) + + return value.__class__(touched_outputs) + elif isinstance(value, dict): + # apply inplace to avoid recreating dict inherited objects + for key in value.keys(): + value[key] = apply_to_tensors_only(function, value[key]) + return value + + elif isinstance(value, torch.Tensor): + # this also applies to torch.Tensor's subclasses like torch.nn.parameter.Parameter + touched_output = function(value) + + # restore zero param attributes if those get stripped by `backward_function` + if not is_zero_param(touched_output) and is_zero_param(value): + touched_output.ds_param_alias = value + + return touched_output + else: + if not is_builtin_type(value): + global warned + if warning_msg_fn and not warned and dist.get_rank() == 0: + logger.warning(warning_msg_fn(value)) + warned = True + return value + + +def get_mapping_to_flat_buffer(tensors: List[torch.Tensor]) -> List[Tuple[torch.Tensor, int, int]]: + tensor_infos: List[Tuple[torch.Tensor, int, int]] = [] + + offset = 0 + for tensor in tensors: + tensor_numel = tensor.numel() + # record some data so we can restore the device tensor later + tensor_infos.append((tensor, offset, tensor_numel)) + offset += tensor_numel + + return tensor_infos diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/sequence/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ee6e3bc5a29137ed25c58a1a576d03bf32623e8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..605d658129f85945b2b2206640df80c428f3bc6d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..085469b669ddea5c734839cadd4edaba3b7588b3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/layer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54065dbe50bd295f3671478eed9b98da3fb26c16 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/layer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/cross_entropy.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..baa7bc1ea7a803055f2cca20141ea6888cc3f8c7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/sequence/cross_entropy.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +import deepspeed.comm as dist + + +class _VocabSequenceParallelCrossEntropy(torch.autograd.Function): + + @staticmethod + def forward(ctx, vocab_seq_parallel_logits, target, sp_group): + # vocab_seq_parallel_logits: [S/P, B, V] + # target: [S/P, B] + # return: [S, B] + + # Need softmax for backward + softmax = torch.nn.functional.softmax(vocab_seq_parallel_logits, dim=-1) + ctx.vocab_size = vocab_seq_parallel_logits.size(2) + loss = torch.nn.functional.nll_loss(softmax.log().view(-1, ctx.vocab_size), target.view(-1), reduction='none') + + sp_world_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + ctx.sp_world_size = sp_world_size + ctx.sp_rank = sp_rank + ctx.seqlen = vocab_seq_parallel_logits.size(0) * sp_world_size + batch_size = vocab_seq_parallel_logits.size(1) + + loss_all = torch.empty(ctx.seqlen, + batch_size, + dtype=vocab_seq_parallel_logits.dtype, + device=vocab_seq_parallel_logits.device) + dist.all_gather_into_tensor(loss_all, loss, group=sp_group) + + ctx.save_for_backward(softmax, target) + + return loss_all + + @staticmethod + def backward(ctx, grad_output): + softmax, target = ctx.saved_tensors + + step_seqlen = ctx.seqlen // ctx.sp_world_size + sp_rank = ctx.sp_rank + grad_output_part = grad_output[step_seqlen * sp_rank:step_seqlen * (sp_rank + 1), :] + + grad_input = softmax + grad_2d = grad_input.view(-1, ctx.vocab_size) + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) + + grad_2d[arange_1d, target.view(-1)] -= 1 + grad_input.mul_(grad_output_part.unsqueeze(dim=-1)) + + return grad_input, None, None, None + + +def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, target, sp_group): + return _VocabSequenceParallelCrossEntropy.apply(vocab_parallel_logits, target, sp_group) diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/fpdt_layer.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/fpdt_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..4fa2cc988a19490f91c90867fe9a01d8a120e2e8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/sequence/fpdt_layer.py @@ -0,0 +1,1225 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from typing import Optional, Any, Tuple +from torch import Tensor +from packaging import version +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator + +try: + import flash_attn + from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward + flash_attn_version = version.parse(flash_attn.__version__) +except ImportError: + _flash_attn_forward = None + _flash_attn_backward = None + +from einops import rearrange +from .layer import single_all_to_all, apply_rotary_pos_emb + + +def _rotate_half_backward(x): + x = rearrange(x, '... (j d) -> ... j d', j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((x2, -x1), dim=-1) + + +def apply_rotary_pos_emb_backward(grad_output, freqs_cos, freqs_sin): + rot_dim = freqs_cos.shape[-1] + grad, grad_pass = grad_output[..., :rot_dim], grad_output[..., rot_dim:] + grad_t = (grad * freqs_cos) + (_rotate_half_backward(grad * freqs_sin)) + grad = grad_t if grad_pass.shape[-1] == 0 else torch.cat((grad_t, grad_pass), dim=-1) + return grad + + +def _update_out_and_lse( + out: torch.Tensor, + lse: torch.Tensor, + block_out: torch.Tensor, + block_lse: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + + block_out = block_out.to(torch.float32) + block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1) + + new_lse = lse + torch.log1p(torch.exp(block_lse - lse)) + + out = torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out + + lse = new_lse + return out, lse + + +def update_out_and_lse( + out: Optional[torch.Tensor], + lse: Optional[torch.Tensor], + block_out: torch.Tensor, + block_lse: torch.Tensor, + slice_=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if out is None: + if slice_ is not None: + raise RuntimeError("first update_out_and_lse should not pass slice_ args") + out = block_out.to(torch.float32) + lse = block_lse.permute(0, 2, 1).contiguous().unsqueeze(dim=-1).contiguous() + elif slice_ is not None: + slice_out, slice_lse = out[slice_], lse[slice_] + slice_out, slice_lse = _update_out_and_lse(slice_out, slice_lse, block_out, block_lse) + out[slice_], lse[slice_] = slice_out, slice_lse + else: + out, lse = _update_out_and_lse(out, lse, block_out, block_lse) + return out, lse + + +class FPDT_InputConstruct(torch.nn.Module): + + def __init__(self, tokens, labels, loss_mask, attention_mask, position_ids, args, sp_size, sp_rank) -> None: + + super(FPDT_InputConstruct, self).__init__() + self.tokens = tokens + self.labels = labels + self.loss_mask = loss_mask + self.attention_mask = attention_mask + self.position_ids = position_ids + global_seq_len = tokens.shape[1] + batch_size = tokens.shape[0] + assert global_seq_len % sp_size == 0 + assert global_seq_len % args.ds_sequence_parallel_fpdt_chunk_size == 0 + num_chunk_per_gpu = global_seq_len // args.ds_sequence_parallel_fpdt_chunk_size + local_seq_len = global_seq_len // sp_size + assert local_seq_len % num_chunk_per_gpu == 0 + + self.num_chunk_per_gpu = num_chunk_per_gpu + self.chunk_size = local_seq_len // num_chunk_per_gpu + self.sp_size = sp_size + self.sp_rank = sp_rank + self.global_seq_len = global_seq_len + self.local_seq_len = local_seq_len + self.batch_size = batch_size + self.device = tokens.device + + def generate(self): + device = self.device + totalChunks = self.global_seq_len // self.chunk_size + token_chunk_idx = torch.arange(self.global_seq_len, device=device, dtype=torch.int) // self.chunk_size + chunk_to_gpu = torch.arange(totalChunks, device=device, dtype=torch.int) + chunk_to_gpu = chunk_to_gpu.reshape(self.num_chunk_per_gpu, -1).t().contiguous() + + gather_chunk = chunk_to_gpu.flatten().unsqueeze(1).contiguous() + mask = gather_chunk == token_chunk_idx + + indices = mask.nonzero(as_tuple=False) + gather_indices = indices[:, 0] + token_chunk_indices = indices[:, 1] + indices = torch.cat([token_chunk_indices[gather_indices == i] for i in range(gather_chunk.shape[0])]) + load_balanced_loss_mask = self.loss_mask[:, indices] if self.loss_mask is not None else self.loss_mask + + indices = indices.reshape(-1, self.chunk_size)[self.num_chunk_per_gpu * self.sp_rank:self.num_chunk_per_gpu * + (self.sp_rank + 1)].flatten().contiguous() + load_balanced_tokens = self.tokens[:, indices] + load_balanced_labels = self.labels[:, indices] if self.labels is not None else self.labels + + load_balanced_attention_mask = self.attention_mask if self.attention_mask is not None else self.attention_mask + load_balanced_position_ids = self.position_ids[:, + indices] if self.position_ids is not None else self.position_ids + + return load_balanced_tokens, load_balanced_labels, load_balanced_loss_mask, load_balanced_attention_mask, load_balanced_position_ids + + +class _FPDTGPUAttentionImpl_(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, + layernorm_output, + attention_mask, + inference_params, + rotary_pos_emb, + spg, + scatter_idx, + gather_idx, + hidden_size, + projection_size, + hidden_size_per_attention_head, + kv_projection_size, + qkv_linear_weight, + qkv_linear_bias, + dropout, + num_chunks=8, + cpu_offloading=True): + + do_save = layernorm_output.requires_grad + + if rotary_pos_emb is not None: + pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) + ctx.pos_emb_cos = pos_emb_cos + ctx.pos_emb_sin = pos_emb_sin + else: + ctx.pos_emb_cos = None + ctx.pos_emb_sin = None + + with torch.no_grad(): + per_gpu_seq_len = layernorm_output.shape[0] + chunk_size = per_gpu_seq_len // num_chunks + assert chunk_size * num_chunks == per_gpu_seq_len + assert attention_mask is None + ctx.num_chunks = num_chunks + ctx.cpu_offloading = cpu_offloading + ctx.spg = spg + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + device = get_accelerator().current_device_name() + ctx.device = device + ctx.dtype = layernorm_output.dtype + ctx.projection_size = projection_size + ctx.kv_projection_size = kv_projection_size + + global_q = [] + global_k = [] + global_v = [] + + ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) + + ctx.dropout_p = dropout + ctx.window_size = (-1, -1) + ctx.alibi_slopes = None + + batch_size = layernorm_output.shape[1] + + global_o = [None for _ in range(num_chunks)] + global_lse = [None for _ in range(num_chunks)] + + for i in range(num_chunks): + + st = chunk_size * i + ed = st + chunk_size + + qkv_chunk = torch.matmul(layernorm_output[st:ed], qkv_linear_weight.t()) + qkv_linear_bias + + q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) + global_q_chunk_len = q_chunk.shape[1] + if rotary_pos_emb is not None: + q_chunk = apply_rotary_pos_emb(q_chunk, + pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], + pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) + global_q.append(q_chunk) + + k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) + if rotary_pos_emb is not None: + k_chunk = apply_rotary_pos_emb(k_chunk, + pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], + pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) + global_k.append(k_chunk) + + v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) + global_v.append(v_chunk) + + for k_i in range(len(global_k)): + causal_chunk = i == k_i + if flash_attn_version >= version.parse("2.6.0"): + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], + global_k[k_i], + global_v[k_i], + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + softcap=0.0, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + else: + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], + global_k[k_i], + global_v[k_i], + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + + global_o[i], global_lse[i] = update_out_and_lse(global_o[i], global_lse[i], block_out, block_lse) + + global_o[i] = global_o[i].to(q_chunk.dtype) + + output = [None for i in range(num_chunks)] + + for i in range(num_chunks): + global_lse[i] = global_lse[i][:, :, :, 0].permute(0, 2, 1).contiguous() + output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) + output = torch.cat(output, dim=1) + + head_dim = output.shape[-1] + + if do_save: + ctx.save_for_backward(layernorm_output) + ctx.global_q = global_q + ctx.global_k = global_k + ctx.global_v = global_v + ctx.attn_output = global_o + ctx.attn_lse = global_lse + ctx.head_dim = head_dim + ctx.batch_size = batch_size + + ctx.qkv_linear_weight = qkv_linear_weight + ctx.qkv_linear_bias = qkv_linear_bias + + return output + + @staticmethod + def backward(ctx, grad_output): + + num_chunks = ctx.num_chunks + device = ctx.device + dtype = ctx.dtype + spg = ctx.spg + scatter_idx = ctx.scatter_idx + gather_idx = ctx.gather_idx + softmax_scale = ctx.softmax_scale + dropout_p = ctx.dropout_p + window_size = ctx.window_size + alibi_slopes = ctx.alibi_slopes + + projection_size = ctx.projection_size + kv_projection_size = ctx.kv_projection_size + + layernorm_output = ctx.saved_tensors[0] + + global_q = ctx.global_q + global_k = ctx.global_k + global_v = ctx.global_v + attn_output = ctx.attn_output + lse = ctx.attn_lse + + qkv_linear_weight = ctx.qkv_linear_weight + qkv_linear_bias = ctx.qkv_linear_bias + + input_chunk_size = layernorm_output.shape[0] // num_chunks + grad_layernorm_output = [ + torch.zeros((input_chunk_size, layernorm_output.shape[1], layernorm_output.shape[2]), + device=device, + dtype=dtype) for _ in range(num_chunks) + ] + + grad_global_attn_output = [] + chunk_size = grad_output.shape[1] // num_chunks + + for i in range(num_chunks): + st = chunk_size * i + ed = st + chunk_size + grad_global_attn_output.append( + single_all_to_all(grad_output[:, st:ed].contiguous(), scatter_idx, gather_idx, 0, spg)) + + del grad_output + + dq = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] + dk = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] + dv = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] + + grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, + device=qkv_linear_weight.device, + dtype=torch.float) + grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, device=qkv_linear_weight.device, dtype=torch.float) + + for i in range(num_chunks): + k_chunk = global_k[i] + v_chunk = global_v[i] + + for q_i in range(num_chunks): + no_computation = q_i < i + if no_computation: + continue + + causal_chunk = q_i == i + + q_chunk = global_q[q_i] + attn_output_chunk = attn_output[q_i] + lse_chunk = lse[q_i] + d_out = grad_global_attn_output[q_i] + + dq_this = torch.zeros(global_q[0].shape, dtype=dtype, device=device) + dk_this = torch.zeros(global_k[0].shape, dtype=dtype, device=device) + dv_this = torch.zeros(global_v[0].shape, dtype=dtype, device=device) + + if flash_attn_version >= version.parse("2.6.0"): + _flash_attn_backward(d_out, + q_chunk, + k_chunk, + v_chunk, + attn_output_chunk, + lse_chunk, + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + softcap=0.0, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + else: + _flash_attn_backward(d_out, + q_chunk, + k_chunk, + v_chunk, + attn_output_chunk, + lse_chunk, + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + + dq[q_i].add_(dq_this.to(torch.float)) + dk[i].add_(dk_this.to(torch.float)) + dv[i].add_(dv_this.to(torch.float)) + + dk_seq_len = dk[i].shape[1] + + if ctx.pos_emb_cos is not None: + dk[i] = apply_rotary_pos_emb_backward(dk[i].to(dtype), + ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) + else: + dk[i] = dk[i].to(dtype) + dv[i] = dv[i].to(dtype) + dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg) + dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg) + + input_st = i * input_chunk_size + input_ed = input_st + input_chunk_size + + input_chunk = layernorm_output[input_st:input_ed].reshape(-1, layernorm_output.shape[-1]) + + dk[i] = dk[i].flatten(2).permute(1, 0, 2) + dv[i] = dv[i].flatten(2).permute(1, 0, 2) + l, b = dk[i].shape[0], dk[i].shape[1] + grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( + torch.matmul(dk[i].reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( + torch.matmul(dv[i].reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk[i].sum(0).sum(0)) + grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv[i].sum(0).sum(0)) + + grad_layernorm_output[i].add_( + torch.matmul(dk[i], qkv_linear_weight[projection_size:projection_size + kv_projection_size])) + grad_layernorm_output[i].add_(torch.matmul(dv[i], + qkv_linear_weight[projection_size + kv_projection_size:])) + + dk[i] = None + dv[i] = None + + for i in range(num_chunks): + dq_seq_len = dq[i].shape[1] + if ctx.pos_emb_cos is not None: + dq[i] = apply_rotary_pos_emb_backward(dq[i].to(dtype), + ctx.pos_emb_cos[:, dq_seq_len * i:dq_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dq_seq_len * i:dq_seq_len * (i + 1)]) + else: + dq[i] = dq[i].to(dtype) + dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(), gather_idx, scatter_idx, 0, spg) + + input_chunk = layernorm_output[:input_chunk_size].reshape(-1, layernorm_output.shape[-1]) + layernorm_output = layernorm_output[input_chunk_size:] + + dq[i] = dq[i].flatten(2).permute(1, 0, 2) + l, b = dq[i].shape[0], dq[i].shape[1] + grad_qkv_linear_weight[:projection_size].add_(torch.matmul(dq[i].reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_bias[:projection_size].add_(dq[i].sum(0).sum(0)) + + grad_layernorm_output[i].add_(torch.matmul(dq[i], qkv_linear_weight[:projection_size])) + + dq[i] = None + + return torch.cat( + grad_layernorm_output, + dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( + dtype), grad_qkv_linear_bias.to(dtype), None, None, None + + +class SequenceChunk: + + def __init__(self, chunk: torch.Tensor, device=None, is_in_use=False): + + self.chunk_shape = chunk.shape + self.chunk_dtype = chunk.dtype + self.device = chunk.device if device is None else device + + cpu_chunk = torch.empty(chunk.shape, dtype=chunk.dtype, device='cpu', pin_memory=True) + + if get_accelerator().on_accelerator(chunk): + cpu_chunk.copy_(chunk, non_blocking=True) + else: + cpu_chunk = chunk + + self.cpu_chunk = cpu_chunk + + self.gpu_chunk = chunk if is_in_use else None + + def load_to_gpu(self): + assert self.gpu_chunk is None + if self.gpu_chunk is not None: + pass + else: + gpu_chunk = torch.empty(self.chunk_shape, device=self.device, dtype=self.chunk_dtype) + gpu_chunk.copy_(self.cpu_chunk, non_blocking=True) + self.gpu_chunk = gpu_chunk + + def get_gpu_chunk(self): + assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device + return self.gpu_chunk + + def check_gpu_chunk(self, ): + assert (self.gpu_chunk is not None) and ( + self.gpu_chunk.device == self.device + ), f"gpu_chunk {self.gpu_chunk is not None} shound be on {self.device}, but it is now on {self.gpu_chunk.device}" + return True + + def offload(self): + assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device + del self.gpu_chunk + self.gpu_chunk = None + + def overwrite_to_cpu(self): + assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device + self.cpu_chunk.copy_(self.gpu_chunk, non_blocking=True) + + +class _FPDTGPUOffloadingAttentionImpl_(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, + layernorm_output, + attention_mask, + inference_params, + rotary_pos_emb, + spg, + scatter_idx, + gather_idx, + hidden_size, + projection_size, + hidden_size_per_attention_head, + kv_projection_size, + qkv_linear_weight, + qkv_linear_bias, + dropout, + num_chunks=8, + cpu_offloading=True): + + do_save = layernorm_output.requires_grad + + if rotary_pos_emb is not None: + pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) + ctx.pos_emb_cos = pos_emb_cos + ctx.pos_emb_sin = pos_emb_sin + else: + ctx.pos_emb_cos = None + ctx.pos_emb_sin = None + with torch.no_grad(): + per_gpu_seq_len = layernorm_output.shape[0] + chunk_size = per_gpu_seq_len // num_chunks + assert chunk_size * num_chunks == per_gpu_seq_len + assert attention_mask is None + ctx.num_chunks = num_chunks + ctx.cpu_offloading = cpu_offloading + ctx.spg = spg + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + ctx.chunk_size = chunk_size + device = get_accelerator().current_device_name() + ctx.device = device + ctx.dtype = layernorm_output.dtype + ctx.projection_size = projection_size + ctx.kv_projection_size = kv_projection_size + + global_q = [] + global_k = [] + global_v = [] + + ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) + + ctx.dropout_p = dropout + ctx.window_size = (-1, -1) + ctx.alibi_slopes = None + + batch_size = layernorm_output.shape[1] + + global_o = [] + global_lse = [] + + layernorm_output_cpu = [] + final_output = [] + + offload_stream = get_accelerator().Stream() + general_offload_stream = get_accelerator().Stream() + compute_stream = get_accelerator().default_stream() + + q_compute_chunk_idx = 0 + kv_compute_chunk_idx = 0 + for i in range(num_chunks): + + qkv_chunk = torch.matmul(layernorm_output[:chunk_size], + qkv_linear_weight.t()) + qkv_linear_bias # torch.Size([18126, 1, 12288]) + + with get_accelerator().stream(general_offload_stream): + layernorm_output_cpu.append(SequenceChunk(layernorm_output[:chunk_size])) + + layernorm_output = layernorm_output[chunk_size:] + + q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) + global_q_chunk_len = q_chunk.shape[1] + + k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) + + v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) + + dist.barrier() + + if ctx.pos_emb_cos is not None: + pos_emb_cos_chunk = pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] + pos_emb_sin_chunk = pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] + + q_chunk = apply_rotary_pos_emb(q_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) + k_chunk = apply_rotary_pos_emb(k_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + with get_accelerator().stream(offload_stream): + global_q.append(SequenceChunk(q_chunk, is_in_use=True)) + global_k.append(SequenceChunk(k_chunk, is_in_use=True)) + global_v.append(SequenceChunk(v_chunk, is_in_use=True)) + + del qkv_chunk + + cur_attn_output = None + cur_attn_lse = None + for k_i in range(len(global_k)): + causal_chunk = i == k_i + with get_accelerator().stream(compute_stream): + if flash_attn_version >= version.parse("2.6.0"): + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + softcap=0.0, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + else: + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + cur_attn_output, cur_attn_lse = update_out_and_lse(cur_attn_output, cur_attn_lse, block_out, + block_lse) + + can_offload_kv = True + if k_i != (len(global_k) - 1) or i != (num_chunks - 1): + if k_i != (len(global_k) - 1): + next_kv_compute_chunk_idx = k_i + 1 + else: + next_kv_compute_chunk_idx = 0 + + if next_kv_compute_chunk_idx == kv_compute_chunk_idx: + can_offload_kv = False + else: + if next_kv_compute_chunk_idx != (len(global_k) - 1): + with get_accelerator().stream(offload_stream): + global_k[next_kv_compute_chunk_idx].load_to_gpu() + global_v[next_kv_compute_chunk_idx].load_to_gpu() + + if i == num_chunks - 1 and k_i == num_chunks - 1: + with get_accelerator().stream(offload_stream): + global_q[0].load_to_gpu() + global_k[0].load_to_gpu() + global_v[0].load_to_gpu() + global_o[0].load_to_gpu() + global_lse[0].load_to_gpu() + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + if can_offload_kv: + global_k[kv_compute_chunk_idx].offload() + global_v[kv_compute_chunk_idx].offload() + kv_compute_chunk_idx = next_kv_compute_chunk_idx + + global_q[q_compute_chunk_idx].offload() + q_compute_chunk_idx += 1 + + all2all_output = single_all_to_all( + cur_attn_output.to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) + final_output.append(all2all_output) + with get_accelerator().stream(general_offload_stream): + global_o.append(SequenceChunk(cur_attn_output.to(ctx.dtype))) + global_lse.append(SequenceChunk(cur_attn_lse[:, :, :, 0].permute(0, 2, 1).contiguous())) + + compute_stream.wait_stream(general_offload_stream) + compute_stream.synchronize() + + final_output = torch.cat(final_output, dim=1) + + head_dim = final_output.shape[-1] + + if do_save: + ctx.layernorm_output = layernorm_output_cpu + ctx.global_q = global_q + ctx.global_k = global_k + ctx.global_v = global_v + ctx.attn_output = global_o + ctx.attn_lse = global_lse + ctx.head_dim = head_dim + ctx.batch_size = batch_size + + ctx.qkv_linear_weight = qkv_linear_weight + ctx.qkv_linear_bias = qkv_linear_bias + + return final_output + + @staticmethod + def backward(ctx, grad_output): + num_chunks = ctx.num_chunks + device = grad_output.device + dtype = ctx.dtype + spg = ctx.spg + scatter_idx = ctx.scatter_idx + gather_idx = ctx.gather_idx + softmax_scale = ctx.softmax_scale + dropout_p = ctx.dropout_p + window_size = ctx.window_size + alibi_slopes = ctx.alibi_slopes + + projection_size = ctx.projection_size + kv_projection_size = ctx.kv_projection_size + + layernorm_output = ctx.layernorm_output + + global_q = ctx.global_q + global_k = ctx.global_k + global_v = ctx.global_v + attn_output = ctx.attn_output + lse = ctx.attn_lse + + qkv_linear_weight = ctx.qkv_linear_weight + qkv_linear_bias = ctx.qkv_linear_bias + + offload_stream = get_accelerator().Stream() + general_offload_stream = get_accelerator().Stream() + compute_stream = get_accelerator().default_stream() + + chunk_size = grad_output.shape[1] // num_chunks + assert chunk_size == layernorm_output[0].cpu_chunk.shape[0] + + grad_layernorm_output = [ + torch.zeros(layernorm_output[0].chunk_shape, device=device, dtype=dtype) for _ in range(num_chunks) + ] + + grad_global_attn_output = [None for _ in range(num_chunks)] + + q_compute_chunk_idx = 0 + kv_compute_chunk_idx = 0 + last_q_accum_idx = 0 + + with get_accelerator().stream(general_offload_stream): + layernorm_output[0].load_to_gpu() + grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, + device=qkv_linear_weight.device, + dtype=torch.float) + grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, + device=qkv_linear_weight.device, + dtype=torch.float) + + grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx, + gather_idx, 0, spg) + get_accelerator().synchronize() + grad_output = grad_output[:, chunk_size:] + + with get_accelerator().stream(offload_stream): + grad_global_attn_output[0] = SequenceChunk(grad_global_attn_output_chunk, is_in_use=True) + dq = [ + SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device=device), is_in_use=True) + ] + [ + SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device='cpu', pin_memory=True), + device) for _ in range(num_chunks - 1) + ] + dk_accum = torch.zeros(global_k[0].chunk_shape, dtype=torch.float, device=device) + dv_accum = torch.zeros(global_v[0].chunk_shape, dtype=torch.float, device=device) + + for i in range(num_chunks): + for q_i in range(num_chunks): + no_computation = q_i < i + if no_computation: + continue + + causal_chunk = q_i == i + + dq_this = torch.zeros(global_q[0].chunk_shape, dtype=dtype, device=device) + dk_this = torch.zeros(global_k[0].chunk_shape, dtype=dtype, device=device) + dv_this = torch.zeros(global_v[0].chunk_shape, dtype=dtype, device=device) + + with get_accelerator().stream(compute_stream): + if flash_attn_version >= version.parse("2.6.0"): + _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + attn_output[q_compute_chunk_idx].get_gpu_chunk(), + lse[q_compute_chunk_idx].get_gpu_chunk(), + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + softcap=0.0, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + else: + _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + attn_output[q_compute_chunk_idx].get_gpu_chunk(), + lse[q_compute_chunk_idx].get_gpu_chunk(), + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + + if i != (len(global_k) - 1): + if q_i != (len(global_q) - 1): + next_q_compute_chunk_idx = q_i + 1 + else: + next_q_compute_chunk_idx = i + 1 + + can_offload_q = True + + if next_q_compute_chunk_idx == q_compute_chunk_idx: + can_offload_q = False + else: + with get_accelerator().stream(offload_stream): + if i > 0 or q_i > 0: + if can_offload_q and last_q_accum_idx != i: # the first q chunk calculate in the loop will be sent out, therefore we do not offload it + dq[last_q_accum_idx].offload() + dq[next_q_compute_chunk_idx].load_to_gpu() + global_q[next_q_compute_chunk_idx].load_to_gpu() + attn_output[next_q_compute_chunk_idx].load_to_gpu() + lse[next_q_compute_chunk_idx].load_to_gpu() + if grad_global_attn_output[next_q_compute_chunk_idx] is not None: + grad_global_attn_output[next_q_compute_chunk_idx].load_to_gpu() + + if grad_global_attn_output[next_q_compute_chunk_idx] is None: + grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), + scatter_idx, gather_idx, 0, spg) + dist.barrier() + grad_output = grad_output[:, chunk_size:] + grad_global_attn_output[next_q_compute_chunk_idx] = SequenceChunk( + grad_global_attn_output_chunk, is_in_use=True) + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + with get_accelerator().stream(compute_stream): + dq[q_compute_chunk_idx].check_gpu_chunk() + dq[q_compute_chunk_idx].gpu_chunk.add_(dq_this) + dk_accum.add_(dk_this) + dv_accum.add_(dv_this) + + offload_stream.wait_stream(compute_stream) + with get_accelerator().stream(offload_stream): + dq[q_compute_chunk_idx].overwrite_to_cpu() + + if can_offload_q: + global_q[q_compute_chunk_idx].offload() + attn_output[q_compute_chunk_idx].offload() + lse[q_compute_chunk_idx].offload() + grad_global_attn_output[q_compute_chunk_idx].offload() + + last_q_accum_idx = q_compute_chunk_idx + q_compute_chunk_idx = next_q_compute_chunk_idx + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + dk_seq_len = dk_accum.shape[1] + + if ctx.pos_emb_cos is not None: + dq_accum = apply_rotary_pos_emb_backward(dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype), + ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) + dk_accum = apply_rotary_pos_emb_backward(dk_accum.to(dtype), + ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) + else: + dq_accum = dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype) + dk_accum = dk_accum.to(dtype) + dv_accum = dv_accum.to(dtype) + + dq_accum = single_all_to_all(dq_accum.contiguous(), gather_idx, scatter_idx, 0, spg) + dk_accum = single_all_to_all(dk_accum.contiguous(), gather_idx, scatter_idx, 0, spg) + dv_accum = single_all_to_all(dv_accum.contiguous(), gather_idx, scatter_idx, 0, spg) + + general_offload_stream.synchronize() + compute_stream.wait_stream(general_offload_stream) + dist.barrier() + + with get_accelerator().stream(compute_stream): + input_chunk = layernorm_output[i].get_gpu_chunk().reshape(-1, layernorm_output[i].chunk_shape[-1]) + + dq_accum = dq_accum.flatten(2).permute(1, 0, 2) + dk_accum = dk_accum.flatten(2).permute(1, 0, 2) + dv_accum = dv_accum.flatten(2).permute(1, 0, 2) + + l, b = dk_accum.shape[0], dk_accum.shape[1] + + grad_qkv_linear_weight[:projection_size].add_( + torch.matmul(dq_accum.reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( + torch.matmul(dk_accum.reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( + torch.matmul(dv_accum.reshape(l * b, -1).t(), input_chunk)) + + grad_qkv_linear_bias[:projection_size].add_(dq_accum.sum(0).sum(0)) + grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk_accum.sum(0).sum(0)) + grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv_accum.sum(0).sum(0)) + + grad_layernorm_output[i].add_(torch.matmul(dq_accum, qkv_linear_weight[:projection_size])) + grad_layernorm_output[i].add_( + torch.matmul(dk_accum, qkv_linear_weight[projection_size:projection_size + kv_projection_size])) + grad_layernorm_output[i].add_( + torch.matmul(dv_accum, qkv_linear_weight[projection_size + kv_projection_size:])) + + del dq_accum, dk_accum, dv_accum + dk_accum = torch.zeros(global_k[i].chunk_shape, dtype=torch.float, device=device) + dv_accum = torch.zeros(global_v[i].chunk_shape, dtype=torch.float, device=device) + dq[kv_compute_chunk_idx].offload() + dq[kv_compute_chunk_idx] = None + + if i != (len(global_k) - 1): + next_kv_compute_chunk_idx = kv_compute_chunk_idx + 1 + with get_accelerator().stream(offload_stream): + global_k[next_kv_compute_chunk_idx].load_to_gpu() + global_v[next_kv_compute_chunk_idx].load_to_gpu() + + with get_accelerator().stream(general_offload_stream): + layernorm_output[next_kv_compute_chunk_idx].load_to_gpu() + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + layernorm_output[kv_compute_chunk_idx].offload() + global_k[kv_compute_chunk_idx].offload() + global_v[kv_compute_chunk_idx].offload() + kv_compute_chunk_idx = next_kv_compute_chunk_idx + + return torch.cat( + grad_layernorm_output, + dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( + dtype), grad_qkv_linear_bias.to(dtype), None, None, None + + +class FPDT_Attention(torch.nn.Module): + + def __init__(self, + config, + first_weight, + first_bias, + second_weight, + second_bias, + sequence_process_group, + gather_idx: int = 0, + scatter_idx: int = 2, + return_bias=True, + chunk_size=65536, + enable_offloading=True) -> None: + + super(FPDT_Attention, self).__init__() + if _flash_attn_forward is None or _flash_attn_backward is None: + raise ImportError( + "DeepSpeed FPDT requires flash-attn 2.6.3. Please install it with `pip install flash-attn --no-build-isolation`." + ) + + self.spg = sequence_process_group + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + self.config = config + + self.projection_size = config.kv_channels * config.num_attention_heads + self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads + self.kv_projection_size = config.kv_channels * config.num_key_value_heads + self.hidden_size = config.hidden_size + + self.qkv_linear_weight = first_weight + self.qkv_linear_bias = first_bias + self.qkv_dense_weight = second_weight + self.qkv_dense_bias = second_bias + + self.reture_bias = return_bias + self.dropout = config.attention_dropout + + self.chunk_size = chunk_size + self.double_buffer = enable_offloading + + def forward(self, + layernorm_output, + attention_mask, + inference_params, + rotary_pos_emb, + cpu_offloading=True) -> Tensor: + self.num_chunks_attn = layernorm_output.shape[0] * dist.get_world_size(self.spg) // self.chunk_size + + if not cpu_offloading or self.num_chunks_attn == 1: + output = _FPDTGPUAttentionImpl_.apply(layernorm_output, attention_mask, inference_params, rotary_pos_emb, + self.spg, self.scatter_idx, self.gather_idx, self.hidden_size, + self.projection_size, self.hidden_size_per_attention_head, + self.kv_projection_size, self.qkv_linear_weight, + self.qkv_linear_bias, self.dropout, self.num_chunks_attn, + cpu_offloading) + else: + output = _FPDTGPUOffloadingAttentionImpl_.apply( + layernorm_output, attention_mask, inference_params, rotary_pos_emb, self.spg, self.scatter_idx, + self.gather_idx, self.hidden_size, self.projection_size, self.hidden_size_per_attention_head, + self.kv_projection_size, self.qkv_linear_weight, self.qkv_linear_bias, self.dropout, + self.num_chunks_attn, cpu_offloading) + + output = output.flatten(2).permute(1, 0, 2).contiguous() + + output = torch.matmul(output, self.qkv_dense_weight.t()) + if not self.reture_bias: + output += self.qkv_dense_bias + return output, self.qkv_dense_bias if self.reture_bias else None + + +@torch.jit.script +def bias_gelu(x): + return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + + +@torch.jit.script +def bias_gelu_back(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) + return ff * g + + +class FPDT_FFN(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, x, w1, b1, w2, b2, add_bias, chunk_size): + do_save = x.requires_grad + ctx.add_bias = add_bias + device = x.device + + with torch.no_grad(): + num_chunk = x.shape[0] // chunk_size + ctx.num_chunk = num_chunk + result = torch.empty(x.shape, device=device, dtype=x.dtype) + assert chunk_size * num_chunk == x.shape[0] + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + x_ = torch.matmul(x[st:ed], w1.t()) + b1 + x_ = bias_gelu(x_) + if add_bias: + result[st:ed] = torch.matmul(x_, w2.t()) + b2 + else: + result[st:ed] = torch.matmul(x_, w2.t()) + + del x_ + + if do_save: + ctx.device = device + ctx.dtype = x.dtype + ctx.save_for_backward(x, w1, b1, w2, b2) + ctx.grad_x_shape = x.shape + return result.to(x.dtype), b2 if not add_bias else None + + @staticmethod + def backward(ctx, grad_output, grad_bias): + x, w1, b1, w2, b2 = ctx.saved_tensors + device = ctx.device + dtype = ctx.dtype + add_bias = ctx.add_bias + + num_chunk = ctx.num_chunk + chunk_size = x.shape[0] // num_chunk + assert chunk_size * num_chunk == grad_output.shape[0] + + grad_w2 = torch.zeros(w2.shape, device=device, dtype=torch.float) + grad_b2 = torch.zeros(b2.shape, device=device, dtype=torch.float) + grad_w1 = torch.zeros(w1.shape, device=device, dtype=torch.float) + grad_b1 = torch.zeros(b1.shape, device=device, dtype=torch.float) + + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + x_chunk = x[st:ed] + + before_act = (torch.matmul(x_chunk, w1.t()) + b1) + before_act_2 = before_act**2 + tanh_out = torch.tanh(0.79788456 * before_act * (1 + 0.044715 * before_act_2)) + ff = 0.5 * before_act * ((1 - tanh_out * tanh_out) * + (0.79788456 + 0.1070322243 * before_act_2)) + 0.5 * (1 + tanh_out) + grad_w2.add_( + torch.matmul(grad_output[st:ed].reshape(-1, grad_output.shape[2]).t(), + (before_act * 0.5 * (1 + tanh_out)).reshape(-1, before_act.shape[2]))) + del before_act, before_act_2, tanh_out + + grad_inter = torch.matmul(grad_output[st:ed], w2) * ff + del ff + + grad_w1.add_(torch.matmul( + grad_inter.reshape(-1, grad_inter.shape[2]).t(), x_chunk.reshape(-1, x.shape[2]))) + grad_b1.add_(grad_inter.sum(0).sum(0)) + + x[st:ed].copy_(torch.matmul(grad_inter, w1)) + + del grad_inter + + if add_bias: + grad_b2.add_(grad_output[st:ed].sum(0).sum(0)) + + return x, grad_w1.to(dtype), grad_b1.to(dtype), grad_w2.to(dtype), grad_b2.to(dtype), None, None + + +class FPDT_LogitsLoss(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, lm_output, labels, logit_weights, rank, spg_size, spg, num_chunk): + labels = labels.t() + chunk_size = lm_output.shape[0] // num_chunk + assert chunk_size * num_chunk == lm_output.shape[0] + batch_size, local_seq_len = lm_output.shape[1], lm_output.shape[0] + loss = torch.empty((batch_size, local_seq_len), dtype=torch.float, device=lm_output.device) + + ctx.num_chunk = num_chunk + ctx.chunk_size = chunk_size + ctx.device = lm_output.device + ctx.dtype = lm_output.dtype + + ctx.rank = rank + ctx.local_seq_len = local_seq_len + with torch.no_grad(): + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + logits_chunk = torch.matmul(lm_output[st:ed], logit_weights.t()).float() + + vocab_size = logits_chunk.size(2) + # nll + softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) + loss_chunk = torch.nn.functional.nll_loss(softmax.log().reshape(-1, vocab_size).contiguous(), + labels[st:ed, :].reshape(-1).contiguous(), + reduction='none') + loss[:, st:ed] = loss_chunk.reshape(chunk_size, batch_size).t() + + del logits_chunk + ctx.save_for_backward(lm_output.to('cpu'), labels) + ctx.logit_weights = logit_weights + + seqlen = local_seq_len * spg_size + batch_size = loss.size(0) + loss = loss.t().contiguous() + loss_all = torch.empty(seqlen, batch_size, dtype=loss.dtype, device=loss.device).contiguous() + + dist.allgather_fn(loss_all, loss, group=spg) + + return loss_all + + @staticmethod + def backward(ctx, grad_output): + lm_output, labels = ctx.saved_tensors + logit_weights = ctx.logit_weights + device = ctx.device + dtype = ctx.dtype + num_chunk = ctx.num_chunk + chunk_size = ctx.chunk_size + + rank = ctx.rank + local_seq_len = ctx.local_seq_len + + grad_output = grad_output[rank * local_seq_len:(rank + 1) * local_seq_len] + grad_lm_output = [None for _ in range(num_chunk)] + grad_logit_weights = torch.zeros(logit_weights.shape, device=grad_output.device, dtype=torch.float) + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + lm_output_chunk = lm_output[st:ed].to(device) + logits_chunk = torch.matmul(lm_output_chunk, logit_weights.t()).float() + + # nll + softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) + vocab_size = logits_chunk.size(2) + + grad_input = softmax + grad_2d = grad_input.reshape(-1, vocab_size).contiguous() + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=device) + + grad_2d[arange_1d, labels[st:ed, :].reshape(-1).contiguous()] -= 1 + grad_input.mul_(grad_output[:chunk_size, :].unsqueeze(dim=-1)) + grad_input = grad_input.to(dtype) + + grad_output = grad_output[chunk_size:].contiguous() + + grad_lm_output_chunk = torch.matmul(grad_input, logit_weights) + grad_lm_output[i] = grad_lm_output_chunk + + grad_logit_weights.add_( + torch.matmul( + grad_input.reshape(-1, grad_input.shape[2]).t(), + lm_output_chunk.reshape(-1, lm_output_chunk.shape[2]))) + + return torch.cat(grad_lm_output, dim=0).to(dtype), None, grad_logit_weights.to(dtype), None, None, None, None diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/layer.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..ecbe0d94120e2312d82ce65bf6c12fcd5dbd48d9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/sequence/layer.py @@ -0,0 +1,440 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +import torch + +from typing import Any, Tuple +from torch import Tensor +from torch.nn import Module + +from einops import rearrange + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.tp_shard import get_shard_size_list, set_num_kv_heads, get_num_kv_heads +from deepspeed.utils import groups + + +def _generate_layout_params(scatter_idx, batch_dim_idx, seq_world_size, input): + """ + This function generates the parameters required for `permute` and `reshape` operations, + which are used to process data before and after `all2all` communication. + """ + if batch_dim_idx == 0: + if scatter_idx < 2: + bs, global_seq_len, num_local_head, head_dim = input.shape + pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim] + pre_all2all_permute_idx = (1, 0, 2, 3, 4) + + post_all2all_permute_idx = (1, 2, 0, 3, 4) + post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim] + else: + bs, local_seq_len, num_total_head, head_dim = input.shape + assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim] + pre_all2all_permute_idx = (2, 0, 1, 3, 4) + + post_all2all_permute_idx = (1, 0, 2, 3, 4) + post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim] + else: + if scatter_idx < 2: + global_seq_len, bs, num_local_head, head_dim = input.shape + pre_all2all_inp_shape = [seq_world_size, global_seq_len // seq_world_size, bs, num_local_head, head_dim] + pre_all2all_permute_idx = None + + post_all2all_permute_idx = (1, 2, 0, 3, 4) + post_all2all_res_shape = [bs, seq_world_size * global_seq_len, num_local_head // seq_world_size, head_dim] + else: + local_seq_len, bs, num_total_head, head_dim = input.shape + assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + pre_all2all_inp_shape = [local_seq_len, bs, seq_world_size, num_total_head // seq_world_size, head_dim] + pre_all2all_permute_idx = (2, 0, 1, 3, 4) + post_all2all_permute_idx = None + post_all2all_res_shape = [local_seq_len * seq_world_size, bs, num_total_head // seq_world_size, head_dim] + + return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape + + +def post_all2all(permute_idx, res_shape): + """ + Post-processing function for `all2all` communication. + """ + + def post_func(input): + if permute_idx is not None: + input = input.permute(permute_idx).contiguous() + output = input.reshape(res_shape).contiguous() + + return output + + return post_func + + +def pre_all2all_fun(permute_idx, inp_shape, input): + """ + Pre-processing function for `all2all` communication. + """ + input_t = input.reshape(inp_shape).contiguous() + if permute_idx is not None: + input_t = input_t.permute(permute_idx).contiguous() + return input_t + + +def _rotate_half(x): + """ + change sign so the last dimension becomes [-odd, +even] + """ + x = rearrange(x, '... (j d) -> ... j d', j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(t, freqs_cos, freqs_sin): + """ + input tensor t is of shape [seq_length, ..., dim] + rotary positional embeding tensor freqs is of shape [seq_length, ..., dim] + check https://kexue.fm/archives/8265 for detailed formulas + """ + rot_dim = freqs_cos.shape[-1] + # ideally t_pass is empty so rotary pos embedding is applied to all tensor t + t, t_pass = t[..., :rot_dim], t[..., rot_dim:] + + # first part is cosine component + # second part is sine component, need to change signs with _rotate_half method + t = (t * freqs_cos) + (_rotate_half(t) * freqs_sin) + + res = t if t_pass.shape[-1] == 0 else torch.cat((t, t_pass), dim=-1) + return res + + +def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group): + seq_world_size = dist.get_world_size(group) + inp_shape = list(input.shape) + assert batch_dim_idx in [0, 1], "batch_dim_idx must be either 0 or 1" + + if not (scatter_idx < 2): + input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size) + input = input.transpose(0, scatter_idx).contiguous() + local_heads = input_splits[groups._get_sequence_parallel_rank()] + output_splits = [local_heads] * seq_world_size + + output_buffer_shape = [seq_world_size * local_heads] + list(input.shape[1:]) + output = torch.empty(output_buffer_shape, device=input.device, dtype=input.dtype) + dist.all_to_all_single(output,input,output_split_sizes=output_splits,\ + input_split_sizes=input_splits,group=group) + ###[seq_ws*local_heads, ...] to [seq_ws, local_heads, ...] + output = output.view(seq_world_size, local_heads, *output.shape[1:]) + ###[seq_ws,local_heads,b,seq_len,...] to [seq_ws,seq_len,b,local_heads,...] + + ### batch_dim_idx=0 [seq_ws,local_heads,seq_len,b,...] to [b, seq_ws, seq_len, local_heads ...] + ### batch_dim_idx=1 [seq_ws,local_heads,b,seq_len,...] to [seq_ws,seq_len,b,local_heads,...] + if batch_dim_idx == 0: + order = [3, 0, 2, 1] + list(range(4, len(output.shape))) + output = output.permute(order).contiguous() + ###[b, seq_ws*local_seq_len, local_heads,...] + output = output.view(output.shape[0], inp_shape[gather_idx] * seq_world_size, + *output.shape[3:]).contiguous() + elif batch_dim_idx == 1: + output = output.transpose(1, 3).contiguous() + ###[seq_ws*local_seq_len, b, local_heads,...] + output = output.view(inp_shape[gather_idx] * seq_world_size, *output.shape[2:]).contiguous() + else: + # The compatibility handling of 4D and 3D tensors, standardizing to 3D. + input = input.reshape(input.shape[0], input.shape[1], -1) + + if batch_dim_idx == 0: #b,s,h + input = input.permute(1, 2, 0).contiguous() #s,h,b + elif batch_dim_idx == 1: #s,b,h + input = input.transpose(1, 2).contiguous() #s,h,b + seq_len, h, batch_size = input.shape + num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size) + local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()] + h_dim = h // local_heads + local_seq_len = seq_len // seq_world_size + + input = input.view(seq_len * h, batch_size) + local_seq_len_with_heads = int(input.shape[0] / seq_world_size) # dim size of local_seq_len*local_heads*hdim + input_splits = [local_seq_len_with_heads] * seq_world_size + coeff = local_seq_len_with_heads // local_heads #per head: dim size of local_seq_len*hdim + + #uneven seq_world_size coeff, total_heads/local_heads. + heads_scale_coeff = get_num_kv_heads() / local_heads + + output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list] + output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads) + total_h = int(inp_shape[gather_idx] * heads_scale_coeff) + output = torch.empty(output_buff_d1_size, input.shape[1], device=input.device, dtype=input.dtype) + dist.all_to_all_single(output,input,output_split_sizes=output_splits, \ + input_split_sizes=input_splits,group=group) + ################## + #suppose 7 heads divide into 4 ranks [2,2,2,1] + #chunk_num_heads_small=floor(7/4)=1 + #chunk_num_heads_large=ceil(7/4)=2 + #num_chunk_heads_large=len([2,2,2])=3, all2all_buffer_counts + #num_chunk_heads_small=len([1])=1, all2all_buffer_counts + #total_num_large_heads=sum([2,2,2])=7 + #total_num_small_heads=sum([1])=1 + + chunk_num_heads_small = get_num_kv_heads() // seq_world_size # even heads compatible + chunk_num_heads_large = chunk_num_heads_small + 1 + num_chunk_heads_large = get_num_kv_heads() % seq_world_size + num_chunk_heads_small = seq_world_size - num_chunk_heads_large + total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large + total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small + + heads_large_combine_size = coeff * total_num_large_heads + heads_small_combine_size = coeff * total_num_small_heads + heads_large_chunk, heads_small_chunk = output.split([heads_large_combine_size, heads_small_combine_size], + dim=0) + heads_large_chunk = heads_large_chunk.view(num_chunk_heads_large, local_seq_len, chunk_num_heads_large, h_dim, + batch_size) + heads_small_chunk = heads_small_chunk.view(num_chunk_heads_small, local_seq_len, chunk_num_heads_small, h_dim, + batch_size) + if batch_dim_idx == 0: + #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[batch,local_seq_len,all2all_buffer_counts*n_heads,dim] + order = [4, 1, 0, 2, 3] + heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(batch_size, local_seq_len, + total_num_large_heads, h_dim) + heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(batch_size, local_seq_len, + total_num_small_heads, h_dim) + elif batch_dim_idx == 1: + #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[local_seq_len,batch,all2all_buffer_counts*n_heads,dim] + order = [1, 4, 0, 2, 3] + heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(local_seq_len, batch_size, + total_num_large_heads, h_dim) + heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(local_seq_len, batch_size, + total_num_small_heads, h_dim) + + output = torch.cat([heads_large_chunk, heads_small_chunk], dim=2).contiguous() + + inp_shape[scatter_idx] = inp_shape[scatter_idx] // seq_world_size + output_shape= inp_shape[: gather_idx] + \ + [total_h,] + \ + inp_shape[gather_idx + 1:] + + output = output.view(output_shape) + + return output + + +def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, async_op=False, handle=None, type=None): + seq_world_size = dist.get_world_size(group) + # we only need num_heads once + num_heads = input.shape[2] + + if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2): + # Assuming here that the number of heads for q is consistent with kv + # If not, additional logic is required for cases like GQA + if get_num_kv_heads() is None: + assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})" + # set heads at first call by num_total_heads. + # then use ``get_num_kv_heads() is not None`` to re-entry uneven path. + set_num_kv_heads(num_heads) + assert async_op == False, "uneven head sp does not support async op" + return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group) + + pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = _generate_layout_params( + scatter_idx, batch_dim_idx, seq_world_size, input) + + input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input) + + post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape) + output = torch.empty_like(input_t) + work = dist.all_to_all_single(output, input_t, group=group, async_op=async_op) + + if async_op: + if type in ('dq', 'dk'): + handle[type + '_work'] = work + handle[type + '_grad'] = output + handle[type + '_post_all2all_func'] = post_all2all_fun + return output.view(post_all2all_res_shape) + + res = post_all2all_fun(output) + return res + + +class _DimZeroAllToAll(torch.autograd.Function): + """Differentiable All2All across dimension 0.""" + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor: + world_size = dist.get_world_size(group) + assert input.shape[0] == world_size, f"Dim 0 {input.shape[0]} is not world size" + + ctx.group = group + + output = torch.empty_like(input).contiguous() + # torch.distributed.nn.functional.all_to_all_single(output, input.contiguous(), group=group) + dist.all_to_all_single(output, input.contiguous(), group=group) + return output + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor]: + return (None, _DimZeroAllToAll.apply(ctx.group, *grad_output)) + + +class _SeqAllToAll(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, + group: dist.ProcessGroup, + input: Tensor, + scatter_idx: int, + gather_idx: int, + batch_dim_idx: int, + stream=None, + handle=None, + type=None, + is_fwd=True) -> Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + ctx.stream = stream + ctx.handle = handle + ctx.type = type + ctx.batch_dim_idx = batch_dim_idx + if ctx.handle is None: + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) + + else: + # overlap communication path + if not is_fwd and type == 'o': + assert ctx.stream != None + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) + get_accelerator().current_stream().wait_stream(ctx.stream) + # The computation of d o_weight can overlap with the communication of d o_input + + elif not is_fwd and type in ('q', 'k'): + # Achieve communication overlap by pipelining the matrix computation and communication of dq, dk, and dv + type = 'd' + type + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, True, handle, type) + + elif is_fwd and type in ('q', 'k'): + # Achieve communication overlap by pipelining the matrix computation and communication of q, k, and v + type = 'fwd_' + type + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False, handle, type) + + else: + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) + + return res + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: + + return (None, + _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx, ctx.batch_dim_idx, + ctx.stream, ctx.handle, ctx.type, False), None, None, None, None, None, None, None) + + +class DistributedAttention(torch.nn.Module): + """Initialization. + + Arguments: + local_attention (Module): local attention with q,k,v + sequence_process_group (ProcessGroup): sequence parallel process group + scatter_idx (int): scatter_idx for all2all comm + gather_idx (int): gather_idx for all2all comm + """ + + def __init__( + self, + local_attention: Module, + sequence_process_group: dist.ProcessGroup, + scatter_idx: int = 2, + gather_idx: int = 0, + sp_stream=None, + ) -> None: + + super(DistributedAttention, self).__init__() + self.local_attn = local_attention + self.spg = sequence_process_group + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + self.sp_overlap_comm = False + self.overlap_handles = None + self.sp_stream = sp_stream + if sp_stream is not None: + self.overlap_handles = {} + self.sp_overlap_comm = True + self.default_stream = get_accelerator().default_stream() + + def layer_sync(self, layer): + if self.sp_overlap_comm and hasattr(layer, 'done_event'): + self.default_stream.wait_event(layer.done_event) + + def forward(self, + query: Tensor, + key: Tensor, + value: Tensor, + batch_dim_idx: int, + rotary_pos_emb=None, + *args: Any, + **kwargs) -> Tensor: + """ forward + + Arguments: + query (Tensor): query input to the layer + key (Tensor): key input to the layer + value (Tensor): value input to the layer + batch_dim_idx (int): indicating which dim is batch + args: other args + + Returns: + * output (Tensor): context output + """ + + # TODO Merge three alltoall calls into one + # TODO (Reza): change the api on the megatron-deepspeed side so that we only receive all data (q,k, and v) together! + #in shape : e.g., [s/p:h:] + + def bwd_hook(layer_type): + + def pre_hook_fun(grad): + type = 'd' + layer_type + self.overlap_handles[type + '_work'].wait() + self.sp_stream.wait_stream(self.default_stream) + all2all_output = self.overlap_handles[type + '_grad'] + grad = list(grad) + grad[0] = self.overlap_handles[type + '_post_all2all_func'](all2all_output) + grad = tuple(grad) + + return pre_hook_fun + + self.layer_sync(query) + query_layer = _SeqAllToAll.apply(self.spg, query, self.scatter_idx, self.gather_idx, batch_dim_idx, None, + self.overlap_handles, 'q') + self.layer_sync(key) + key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx, batch_dim_idx, None, + self.overlap_handles, 'k') + if self.sp_overlap_comm: + self.default_stream.wait_stream(self.sp_stream) + + value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx, batch_dim_idx, None, + self.overlap_handles, 'v') + + if self.sp_overlap_comm: + # Register a hook to synchronize dq and dk after the all-to-all + # operation when the gradient data is used. + # Place this logic after the q, k, v all-to-all operation to + # improve interpreter speed to + # call and launch of the forward all-to-all communication. + grad_fn_q = query.grad_fn.next_functions[0][0] + grad_fn_q.register_prehook(bwd_hook(layer_type='q')) + grad_fn_k = key.grad_fn.next_functions[0][0] + grad_fn_k.register_prehook(bwd_hook(layer_type='k')) + + #out shape : e.g., [s:h/p:] + if rotary_pos_emb is not None: + pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) + query_layer = apply_rotary_pos_emb(query_layer, pos_emb_cos, pos_emb_sin) + key_layer = apply_rotary_pos_emb(key_layer, pos_emb_cos, pos_emb_sin) + + context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs) + + output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx, batch_dim_idx, + self.sp_stream, self.overlap_handles, 'o') + + #out e.g., [s/p::h] + return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d902a883f6eeee72b3c9ccbf5a11abf22de8e33 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .logging import logger, log_dist +from .comms_logging import get_caller_func +#from .distributed import init_distributed +from .init_on_device import OnDevice +from .groups import * +from .nvtx import instrument_w_nvtx +# TODO: Move tensor fragment and mixed precision to zero utils +from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address, get_full_hp_grad, map_to_flat_opt_states +from .tensor_fragment import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state +from .tensor_fragment import set_full_hp_param, set_full_hp_grad +from .tensor_fragment import safe_set_full_fp32_param, safe_set_full_optimizer_state, safe_set_full_grad +from .tensor_fragment import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state +from .tensor_fragment import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state +from .tensor_fragment import safe_update_full_grad_vectorized +from .z3_leaf_module import set_z3_leaf_modules, unset_z3_leaf_modules, get_z3_leaf_modules, z3_leaf_module, z3_leaf_parameter, set_z3_leaf_module +from .mixed_precision_linkage import link_hp_params, lazy_init_hp_params_optimizer_state +from deepspeed.runtime.dataloader import RepeatingLoader +from .numa import get_numactl_cmd diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e97c4b73da902815fce75bf1af723d1786f36968 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/bwc.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/bwc.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c0151367e892bb79f161a83cfdf2f6d5a4e9b00e Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/bwc.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a43e83619e6a7985cce48f87c38affb1a1a42ee3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54a36185f40095d6c55fa9d233a798ba225525cd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/debug.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/debug.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a436ccc447f65cea7fa32e1eefbfa99da08bd058 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/debug.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e599018a9a209d40c28ef7d6c3085cc83cc4bed4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/groups.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/groups.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39a06e80514b1f52db9446477eb72478aaaf88fd Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/groups.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0bceb451056c071a9e541771a745f40352f633e3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/logging.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..529d98f85d01f99c4f8efa5286dcc8316146e999 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/logging.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54b800910fe0721cdd2446789d49f3942960c5b9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/numa.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/numa.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf6b7f29152b1aafa5da25845daae72704ac0cb1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/numa.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5de2f1e3aedc1a6c5a7eff2cc868ad421af7673d Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35e3ef11987b0ca83954f66b9797cc1e7f2b0284 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/timer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/timer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f49ca5d1010f4aac90cf09db5dfb7042745a49a Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/timer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/torch.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/torch.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff93fdfd757b4260b69a598f715149ddbc558bbb Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/torch.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/types.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/types.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..150a5bb6412542819eded38ba7553c539245deb9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/types.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f98d73b3386d459cbba4536ad17210613fb235a9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2471b769c3180d410e98b9d6c45ba30509fb4016 Binary files /dev/null and b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py b/venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py new file mode 100644 index 0000000000000000000000000000000000000000..69fcc251a68429a768892b7cf7c6824de7cb5ba2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + + +def bwc_tensor_model_parallel_rank(mpu=None): + """Backwards-compatible way of querying the tensor model parallel rank from + an ``mpu`` object. + + *Tensor* model parallelism means that tensors are physically split across + processes. This contrasts with *pipeline* model parallelism, in which the + layers are partitioned but tensors left intact. + + The API for tensor model parallelism has changed across versions and this + helper provides a best-effort implementation across versions of ``mpu`` + objects. The preferred mechanism is + ``mpu.get_tensor_model_parallel_rank()``. + + This should "just work" with both Megatron-LM and DeepSpeed's pipeline + parallelism. + + Args: + mpu (model parallel unit, optional): The tensor model parallel rank. + If ``mpu=None``, returns 0. Defaults to ``None``. + + Returns: + int: the rank + """ + if mpu is None: + # No model parallelism in easy :) + return 0 + + if hasattr(mpu, 'get_tensor_model_parallel_rank'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + return mpu.get_tensor_model_parallel_rank() + elif hasattr(mpu, 'get_slice_parallel_rank'): + # Some DeepSpeed + pipeline parallelism versions + return mpu.get_slice_parallel_rank() + else: + # Deprecated Megatron and DeepSpeed convention + return mpu.get_model_parallel_rank() + + +def bwc_tensor_model_parallel_world_size(mpu=None): + """Backwards-compatible way of querying the tensor model parallel world size. + Similar to bwc_tensor_model_parallel_rank. + """ + if mpu is None: + return 1 + + if hasattr(mpu, 'get_tensor_model_parallel_world_size'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + return mpu.get_tensor_model_parallel_world_size() + elif hasattr(mpu, 'get_slice_parallel_world_size'): + # Some DeepSpeed + pipeline parallelism versions + return mpu.get_slice_parallel_world_size() + else: + # Deprecated Megatron and DeepSpeed convention + return mpu.get_model_parallel_world_size() + + +def bwc_tensor_model_parallel_group(mpu=None): + """Backwards-compatible way of querying the tensor model parallel group. + Similar to bwc_tensor_model_parallel_rank. + """ + if mpu is None: + return None + + if hasattr(mpu, 'get_tensor_model_parallel_group'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + return mpu.get_tensor_model_parallel_group() + elif hasattr(mpu, 'get_slice_parallel_group'): + # Some DeepSpeed + pipeline parallelism versions + return mpu.get_slice_parallel_group() + else: + # Deprecated Megatron and DeepSpeed convention + return mpu.get_model_parallel_group() + + +def bwc_pipeline_parallel_world_size(mpu=None): + """Backwards-compatible way of querying the pipeline parallel world size.""" + world_size = 1 + if mpu is not None: + if hasattr(mpu, 'get_pipeline_model_parallel_world_size'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + world_size = mpu.get_pipeline_model_parallel_world_size() + elif hasattr(mpu, 'get_pipe_parallel_world_size'): + # DeepSpeed Topology + world_size = mpu.get_pipe_parallel_world_size() + return world_size + + +def bwc_pipeline_parallel_group(mpu=None): + """Backwards-compatible way of querying the pipeline parallel group.""" + if mpu is None: + return None + if hasattr(mpu, 'get_pipeline_model_parallel_group'): + # Megatron + return mpu.get_pipeline_model_parallel_group() + elif hasattr(mpu, 'get_pipe_parallel_group'): + # DeepSpeed Topology + return mpu.get_pipe_parallel_group() + assert False, 'mpu does not support pipeline parallel group' diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/comms_logging.py b/venv/lib/python3.10/site-packages/deepspeed/utils/comms_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..abe0c9d47cd07e124d95d35736ae64f2bc2d85fa --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/comms_logging.py @@ -0,0 +1,181 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +from deepspeed.utils import log_dist + + +def get_caller_func(frame=3): + import sys + return sys._getframe(frame).f_code.co_name + + +def print_rank_0(message): + import deepspeed.comm as dist + if dist.get_rank() == 0: + print(message) + + +# Helper function to pretty-print message sizes +def convert_size(size_bytes): + if size_bytes == 0: + return "0B" + size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB") + i = int(math.floor(math.log(size_bytes, 1024))) + p = math.pow(1024, i) + s = round(size_bytes / p, 2) + return "%s %s" % (s, size_name[i]) + + +# Helper function to calculate algbw and busbw. +# See https://gist.github.com/jeffra/b5e80466b4c86be00ea3b6f130fb7a36 and https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md +def calc_bw_log(comm_op, size, duration): + import deepspeed.comm as dist + + n = dist.get_world_size() + tput = 0 + busbw = 0 + if comm_op == "all_to_all_single": + tput = (size / duration) + busbw = (size / duration) * ((n - 1) / n) + elif comm_op == "all_gather" or comm_op == "all_gather_into_tensor" or comm_op == "reduce_scatter" or comm_op == "reduce_scatter_tensor": + size *= n + tput = (size / duration) + busbw = (size / duration) * ((n - 1) / n) + elif comm_op == "all_reduce" or comm_op == "all_reduce_coalesced" or comm_op == "inference_all_reduce": + tput = (size * 2 / duration) + busbw = (size / duration) * (2 * (n - 1) / n) + elif comm_op == "send" or comm_op == "recv" or comm_op == "isend" or comm_op == "irecv" or comm_op == "broadcast" or comm_op == "reduce" or comm_op == "gather" or comm_op == "scatter" or comm_op == "barrier": + tput = (size / duration) + busbw = tput + else: + print_rank_0("wrong comm_op specified") # noqa: F821 + exit(0) + + # convert to Gbps + tput *= 8 + busbw *= 8 + + tput /= 1e6 + busbw /= 1e6 + + return tput, busbw + + +class CommsLogger: + + def __init__(self): + from deepspeed.comm.constants import COMMS_LOGGER_VERBOSE_DEFAULT, COMMS_LOGGER_DEBUG_DEFAULT, COMMS_LOGGER_PROF_OPS_DEFAULT, COMMS_LOGGER_PROF_ALL_DEFAULT, COMMS_LOGGER_ENABLED_DEFAULT + self.comms_dict = {} + self.verbose = COMMS_LOGGER_VERBOSE_DEFAULT + self.debug = COMMS_LOGGER_DEBUG_DEFAULT + self.prof_ops = COMMS_LOGGER_PROF_OPS_DEFAULT + self.prof_all = COMMS_LOGGER_PROF_ALL_DEFAULT + self.enabled = COMMS_LOGGER_ENABLED_DEFAULT + + def configure(self, comms_config): + self.enabled = comms_config.comms_logger_enabled + if self.enabled: + self.verbose = comms_config.comms_logger.verbose + self.debug = comms_config.comms_logger.debug + self.prof_ops = comms_config.comms_logger.prof_ops + self.prof_all = comms_config.comms_logger.prof_all + + # There are three settings for the op profiler: + # - Global profiling (profile all comms) + # - Op-type profiling (e.g. profile all all_reduce comms) + # - Op profiling (e.g. profile a specific all_reduce op) + def start_profiling_comms(self): + self.prof_all = True + + def stop_profiling_comms(self): + self.prof_all = True + + # E.g. start_profiling_op('all_reduce') + def start_profiling_op(self, op_name_list): + self.prof_ops = list(set(self.prof_ops) | set(op_name_list)) + + def stop_profiling_op(self, op_name_list): + self.prof_ops = [op for op in self.prof_ops if op not in op_name_list] + + # Add log entry + def append(self, raw_name, record_name, latency, msg_size): + algbw, busbw = calc_bw_log(raw_name, msg_size, latency) + if record_name in self.comms_dict.keys(): + # If this comm_op has already been logged with this message size, just add to existing record + if msg_size in self.comms_dict[record_name].keys(): + self.comms_dict[record_name][msg_size][0] += 1 + self.comms_dict[record_name][msg_size][1].append(latency) + self.comms_dict[record_name][msg_size][2].append(algbw) + self.comms_dict[record_name][msg_size][3].append(busbw) + # If this is a new message size for this comm_op, add new record under existing comm_op + else: + self.comms_dict[record_name][msg_size] = [1, [latency], [algbw], [busbw]] + else: + # Create entirely new record + self.comms_dict[record_name] = {msg_size: [1, [latency], [algbw], [busbw]]} + # If verbose, print every comm op + # TODO: Add to tensorboard + if self.verbose: + log_str = f"comm op: {record_name} | time (ms): {latency:.2f} | msg size: {convert_size(msg_size)} | algbw (Gbps): {algbw:.2f} | busbw (Gbps): {busbw:.2f}" + log_dist(log_str, [0]) + + # Print summary at end of iteration, epoch, or training + def log_all(self, print_log=True, show_straggler=False): + import torch + from deepspeed.utils.timer import trim_mean + import deepspeed.comm as dist + from deepspeed.comm.reduce_op import ReduceOp + from deepspeed.accelerator import get_accelerator + + if print_log: + print( + f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total Latency(ms)': <20}{'Avg Latency(ms)': <20}{'tput_avg (Gbps)': <20}{'busbw_avg (Gbps)': <20}" + ) + for record_name in self.comms_dict.keys(): + if print_log: + print(record_name) + for msg_size, vals in sorted(self.comms_dict[record_name].items()): + # vals[0] is the count for each msg size + count = vals[0] + # vals[1] is a list of latency records for each msg size + total_lat = sum(vals[1]) + # vals[2] and vals[3] are the lists of algbw and busbw, respectively + # Get rid of outliers when we print + avg_lat = trim_mean(vals[1], 0.1) + avg_algbw = trim_mean(vals[2], 0.1) + avg_busbw = trim_mean(vals[3], 0.1) + if print_log: + print( + f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{avg_lat: <20.2f}{avg_algbw: <20.2f}{avg_busbw: <20.2f}" + ) + + if show_straggler: + if print_log: + print("_______________________________") + print("Breakdown with straggler effect") + print("-------------------------------") + print( + f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total comm lat(ms)': <20}{'Total straggler(ms)': <20}{'Avg comm lat(ms)': <20}{'Avg straggler(ms)': <20}" + ) + device = get_accelerator().current_device_name() + for record_name in self.comms_dict.keys(): + if print_log: + print(record_name) + for msg_size, vals in sorted(self.comms_dict[record_name].items()): + # vals[0] is the count for each msg size + count = vals[0] + # vals[1] is a list of latency records for each msg size + lats = torch.tensor(vals[1], device=device) + min_lats = torch.tensor(vals[1], device=device) + dist.all_reduce(min_lats, op=ReduceOp.MIN) + total_lat = min_lats.sum().item() + total_straggler = (lats - min_lats).sum().item() + avg_lat = trim_mean(min_lats.tolist(), 0.1) + avg_straggler = trim_mean((lats - min_lats).tolist(), 0.1) + if print_log: + print( + f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{total_straggler: <20.2f}{avg_lat: <20.2f}{avg_straggler: <20.2f}" + ) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/config.py b/venv/lib/python3.10/site-packages/deepspeed/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..15f37ca7d8745cccee562d08b0eae3f10c5d9877 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/config.py @@ -0,0 +1,46 @@ +# 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): + """ Configure throughput timers """ + + enabled: bool = True + """ Turn on/off throughput timers """ + + synchronized: bool = True + """ Whether to synchronize a device when measuring the time. + Synchronizing a device is required to produce the most accurate timer measurements. + However, this comes at the expense of performance degradation. The CPU timer provides + sufficient accuracy in many cases. + """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/debug.py b/venv/lib/python3.10/site-packages/deepspeed/utils/debug.py new file mode 100644 index 0000000000000000000000000000000000000000..6174ba45137e1e70dd004593a69ff41015c0212c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/debug.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import deepspeed.comm as dist + +# For lazy import with printflock() +fcntl = None + +# for debug purposes map module and param objects to their fully qualified names +module_names = {} +param_names = {} + + +def debug_clear_module_and_param_names(): + global module_names + global param_names + module_names = {} + param_names = {} + + +def debug_extract_module_and_param_names(model): + # extract the fully qualified names as soon as the model is acquired + global module_names + global param_names + # XXX: can probably make a map of param2module and vice-versa + module_names = {module: name for name, module in model.named_modules()} + param_names = {param: name for name, param in model.named_parameters()} + + +def debug_module2name(module): + if module in module_names: + return module_names[module] + else: + return "unknown" + + +def debug_module2name_id(module): + return f"name={debug_module2name(module)}" + + +def debug_module2name_class(module): + return f"name={debug_module2name(module)} {module.__class__.__name__}" + + +def debug_param2name(param): + if param in param_names: + return param_names[param] + else: + return "unknown" + + +def debug_param2name_id(param): + return f"name={debug_param2name(param)} id={param.ds_id}" + + +def debug_param2name_id_shape(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape}" + + +def debug_param2name_id_shape_device(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape} device={param.device}" + + +def debug_param2name_id_numel(param): + return f"name={debug_param2name(param)} id={param.ds_id} numel={param.numel()}" + + +def debug_param2name_id_shape_status(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape} status={param.ds_status}" + + +def printflock(*msgs): + """ + + For printing messages for all concurrent gpus w/o getting interleaved text. + + This is useful when debugging issues where multi-gpus don't sync. + + 1. Enable the force debug in say partitioning and zero3 files + 2. Override the usual versions with :: + + def print_rank_0(message, debug=False, force=False): + rank = deepspeed.comm.get_rank() + printflock(f"[{rank}] {message}") + 3. run the program and you get both logs non-interleaved + + But this makes it very difficult to make sense of the output, so the ``log_rank_file`` helper + function might be more useful, as it's easier to send each log stream into a separate file and + then compare those. + + """ + global fcntl + if fcntl is None: + import fcntl + + with open(__file__, "r") as fh: + fcntl.flock(fh, fcntl.LOCK_EX) + try: + print(*msgs) + finally: + fcntl.flock(fh, fcntl.LOCK_UN) + + +fh = None + + +def log_rank_file(rank, *msgs): + """ + Print to a log file of the given rank + + This is useful for debugging hanging in sync processes. Here is a possible workflow: + + 1. Enable the force debug in say partitioning and zero3 files + 2. Override the usual versions of print_rank_0 in those files with :: + + def print_rank_0(message, debug=False, force=False): + rank = deepspeed.comm.get_rank() + log_rank_file(rank, message) + + 3. run the program + 4. fix up the expected differences, e.g. different cuda numbers :: + + perl -pi -e 's|cuda:1|cuda:0|' log_rank_* + + 5. now diff and see where names and ids diverge - you will find where the gpus don't do the same + work (e.g. when some layers get conditionally skipped on one gpu but not all) + + diff -u log_rank_0.txt log_rank_1.txt | less + + """ + global fh + if fh is None: + fh = open(f"log_rank_{rank}.txt", "w") + for m in msgs: + fh.write(f"{m}\n") + fh.flush() + + +def print_backward_tensors(tensor): + + def _print_bwd_tensors(grad_fn): + print(f"Backward tensors in {grad_fn}") + for funcs in grad_fn.next_functions: + if funcs[0]: + try: + tensor = getattr(funcs[0], 'variable') + print(funcs[0]) + print(f"Tensor - id: {id(tensor)}, shape: {tensor.shape}, data: {tensor}, grad: {tensor.grad}") + except AttributeError as e: + _print_bwd_tensors(funcs[0]) + + if hasattr(tensor, 'grad_fn'): + _print_bwd_tensors(tensor.grad_fn) + + +def print_rank(*msg, force=False): + """print something on all global ranks with [rank] prefix. + """ + if not force: + return + global_rank = dist.get_rank() + print(f"[{global_rank}]", *msg) + + +def print_rank0(*msg, force=False): + """print something only on rank 0""" + if not force: + return + global_rank = dist.get_rank() + if global_rank == 0: + print(f"[{global_rank}]", *msg) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/exceptions.py b/venv/lib/python3.10/site-packages/deepspeed/utils/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..9dba0933cc9e236ec150bfd5439f7ce8e157394a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/exceptions.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + + +class DeprecatedException(Exception): + pass diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/groups.py b/venv/lib/python3.10/site-packages/deepspeed/utils/groups.py new file mode 100644 index 0000000000000000000000000000000000000000..250aa06b1a7c0610c705313dd1a06c28672f7ef6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/groups.py @@ -0,0 +1,749 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# The file has been adapted from https://github.com/NVIDIA/Megatron-LM and retains the following license from the original file + +# Copyright (c) 2019, 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. +""" + Support different forms of parallelism in DeepSpeed using multiple process groups. + Given that there are multiple scenarios and use-cases, this file is going to be updated + frequently. For now, the group creation needed for the training scenario is being implemented. + For inference and other new scenarios, the code will be either reused or added to this file. +""" + +from deepspeed import comm as dist +from deepspeed.utils import log_dist +from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size, bwc_pipeline_parallel_world_size +from deepspeed.utils.exceptions import DeprecatedException +from deepspeed.accelerator import get_accelerator + +# Expert parallel group that the current rank belongs to. +_EXPERT_PARALLEL_GROUP = {} +# Mapping of expert parallel group to ranks +_EXPERT_PARALLEL_GROUP_RANKS = {} +# Expert data parallel group that the current rank belongs to. +_EXPERT_DATA_PARALLEL_GROUP = {} +# Mapping of expert data parallel group to ranks +_EXPERT_DATA_PARALLEL_GROUP_RANKS = {} +# dist world group needs to be cloned for some cases +_WORLD_GROUP = None +# ZeRO parameter partitioning group that the current rank belongs to. +_ZERO_PARAM_INTRA_PARALLEL_GROUP = None +# global object to maintain mpu object if passed by a Megatron client +mpu = None +# global object that stores tensor parallel world size for experts +expert_tensor_parallel_world_size = 1 +# All to All quantized graident communication groups +_ALL_TO_ALL_GROUP = {} + +mesh_device = None + + +# Deprecated groups initialize function. +def initialize(ep_size=1, mpu=None): + """ Deprecated function. Retained to inform the users.""" + raise DeprecatedException( + "Please do not use the groups.initialize() API as it is deprecated. Instead, pass the desired ep_size to deepspeed.moe.layer.MoE(..,ep_size,..)" + ) + + +def _ensure_divisibility(numerator, denominator): + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator) + + +# ======== Start: Tensor Parallel Group Attributes ======== + +# Intra-layer model parallel group that the current rank belongs to. +_TENSOR_MODEL_PARALLEL_GROUP = None + +# Model parallel group (both intra- and pipeline) that the current rank belongs to. +_MODEL_PARALLEL_GROUP = None +# Data parallel group that the current rank belongs to. +_DATA_PARALLEL_GROUP = None + +# These values enable us to change the mpu sizes on the fly. +_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None +_MPU_TENSOR_MODEL_PARALLEL_RANK = None + + +def _init_tp_mesh_device(tensor_model_parallel_size=1, data_parallel_size=None): + """Initialize model data parallel groups.""" + + global _DATA_PARALLEL_GROUP + global _MODEL_PARALLEL_GROUP + global _TENSOR_MODEL_PARALLEL_GROUP + + if _TENSOR_MODEL_PARALLEL_GROUP is not None: + return + + if data_parallel_size is None: + data_parallel_size = dist.get_world_size() // tensor_model_parallel_size + + mesh_device = dist.initialize_mesh_device((data_parallel_size, tensor_model_parallel_size), + ("data_parallel", "tensor_parallel")) + _TENSOR_MODEL_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="tensor_parallel") + _DATA_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="data_parallel") + + # They are always equal only in 2D (DP + TP) parallelism. + # _MODEL_PARALLEL_GROUP is assigned the same value as _TENSOR_MODEL_PARALLEL_GROUP + # to allow for future potential changes. + _MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP + + return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP + + +def get_tensor_model_parallel_group(): + """Get the tensor model parallel group the caller rank belongs to.""" + + assert _TENSOR_MODEL_PARALLEL_GROUP is not None, \ + 'intra_layer_model parallel group is not initialized' + return _TENSOR_MODEL_PARALLEL_GROUP + + +def get_model_parallel_group(): + """Get the model parallel group the caller rank belongs to.""" + + assert _MODEL_PARALLEL_GROUP is not None, \ + 'model parallel group is not initialized' + return _MODEL_PARALLEL_GROUP + + +def get_data_parallel_group(): + """Get the data parallel group the caller rank belongs to.""" + assert _DATA_PARALLEL_GROUP is not None, \ + 'data parallel group is not initialized' + return _DATA_PARALLEL_GROUP + + +def set_tensor_model_parallel_world_size(world_size): + """Set the tensor model parallel size""" + global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE + _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size + + +def get_tensor_model_parallel_world_size(): + """Return world size for the tensor model parallel group.""" + global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE + if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None: + return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE + return dist.get_world_size(group=get_tensor_model_parallel_group()) + + +def get_model_parallel_world_size(): + return get_tensor_model_parallel_world_size() + + +def set_tensor_model_parallel_rank(rank): + """Set tensor model parallel rank.""" + global _MPU_TENSOR_MODEL_PARALLEL_RANK + _MPU_TENSOR_MODEL_PARALLEL_RANK = rank + + +def get_tensor_model_parallel_rank(): + """Return my rank for the tensor model parallel group.""" + global _MPU_TENSOR_MODEL_PARALLEL_RANK + if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None: + return _MPU_TENSOR_MODEL_PARALLEL_RANK + return dist.get_rank(group=get_tensor_model_parallel_group()) + + +def get_model_parallel_rank(): + return get_tensor_model_parallel_rank() + + +def get_tensor_model_parallel_src_rank(): + """Calculate the global rank corresponding to the first local rank + in the tensor model parallel group.""" + global_rank = dist.get_rank() + local_world_size = get_tensor_model_parallel_world_size() + return (global_rank // local_world_size) * local_world_size + + +def get_data_parallel_world_size(): + """Return world size for the data parallel group.""" + return dist.get_world_size(group=get_data_parallel_group()) + + +def get_data_parallel_rank(): + """Return my rank for the data parallel group.""" + return dist.get_rank(group=get_data_parallel_group()) + + +# ======== End: Tensor Parallel Group Attributes ======== + + +# Not currently used. Helper function to create a model (tensor) parallel group. +def _create_model_parallel(model_parallel_size_): + """ + Initialize model data parallel groups. + + Arguments: + model_parallel_size: number of GPUs used to parallelize model. + + Returns: + Tuple of data parallel group and model parallel group + + Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we + use 2 GPUs to parallelize the model. The present function will + create 4 model parallel groups and 2 data parallel groups as: + 4 model parallel groups: + [g0, g1], [g2, g3], [g4, g5], [g6, g7] + 2 data parallel groups: + [g0, g2, g4, g6], [g1, g3, g5, g7] + Note that for efficiency, the caller should make sure adjacent ranks + are on the same DGX box. For example if we are using 2 DGX-1 boxes + with a total of 16 GPUs, rank 0 to 7 belong to the first box and + ranks 8 to 15 belong to the second box. + """ + log_dist(f'Creating model parallel group with size {model_parallel_size_}', ranks=[0]) + # Get world size and rank. Ensure some consistencies. + assert dist.is_initialized() + world_size = dist.get_world_size() + model_parallel_size = min(model_parallel_size_, world_size) + _ensure_divisibility(world_size, model_parallel_size) + rank = dist.get_rank() + + _DATA_PARALLEL_GROUP = None + _MODEL_PARALLEL_GROUP = None + # Build the data parallel groups. + for i in range(model_parallel_size): + ranks = range(i, world_size, model_parallel_size) + group = dist.new_group(ranks) + if i == (rank % model_parallel_size): + _DATA_PARALLEL_GROUP = group + + # Build the model parallel groups. + for i in range(world_size // model_parallel_size): + ranks = range(i * model_parallel_size, (i + 1) * model_parallel_size) + group = dist.new_group(ranks) + if i == (rank // model_parallel_size): + _MODEL_PARALLEL_GROUP = group + + return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP + + +def _create_expert_and_data_parallel(expert_parallel_size_, use_data_before_expert_parallel_=False): + """ + Create expert and data parallel groups. + + Note: Caller of this function is responsible to check if the groups already exist. + + Example - E + D parallel + world_size = 16 + expert_parallel_size = 2 # number of experts in same group + expert_data_parallel_group = [0,2,4,6,8,10,12,14], [1,3,5,7,9,11,13,15] - all reduce is only on MoE params + expert_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - no all reduce, but all to all + data_parallel_group = [0,1,...,15] - all reduce is only on non-MoE + use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology + """ + assert dist.is_initialized() + + log_dist(f'Creating expert and data parallel groups with size {expert_parallel_size_}', ranks=[0]) + world_size = dist.get_world_size() + pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu) + rank = dist.get_rank() + + pp_stride = world_size // pp_world_size + _ensure_divisibility(pp_stride, expert_parallel_size_) + + group_name = f"ep_size_{expert_parallel_size_}" + + # Build the expert data parallel groups. + global _EXPERT_DATA_PARALLEL_GROUP + global _EXPERT_DATA_PARALLEL_GROUP_RANKS + + ep_stride = pp_stride // expert_parallel_size_ + + # Only create group if it does not already exist + if group_name not in _EXPERT_DATA_PARALLEL_GROUP: + for pp_stage_start in range(0, world_size, pp_stride): + for i in range(expert_parallel_size_): + if use_data_before_expert_parallel_: + ranks = range(pp_stage_start + i * ep_stride, pp_stage_start + (i + 1) * ep_stride) + else: + ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, expert_parallel_size_) + group = dist.new_group(ranks) + log_dist(f'Creating expert data parallel process group named {group_name} with ranks: {list(ranks)}', + [0]) + if rank in ranks: + _EXPERT_DATA_PARALLEL_GROUP[group_name] = group + _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks + + # Build the expert parallel groups. + global _EXPERT_PARALLEL_GROUP + global _EXPERT_PARALLEL_GROUP_RANKS + + # Only create group if it does not already exist + if group_name not in _EXPERT_PARALLEL_GROUP: + if use_data_before_expert_parallel_: + for pp_stage_start in range(0, world_size, pp_stride): + for i in range(ep_stride): + ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, ep_stride) + group = dist.new_group(ranks) + log_dist( + f'creating expert parallel process group named {group_name} ' + f'with ranks: {list(ranks)}', [0]) + if rank in ranks: + _EXPERT_PARALLEL_GROUP[group_name] = group + _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks + else: + for i in range(world_size // expert_parallel_size_): + ranks = range(i * expert_parallel_size_, (i + 1) * expert_parallel_size_) + group = dist.new_group(ranks) + log_dist(f'creating expert parallel process group named {group_name} ' + f'with ranks: {list(ranks)}', [0]) + if rank in ranks: + _EXPERT_PARALLEL_GROUP[group_name] = group + _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks + + +def _get_expert_parallel_ranks(world_size, + tensor_parallel_size_, + expert_parallel_size_, + pipeline_parallel_size_=1, + use_data_before_expert_parallel_=False): + """Generate expert parallel and expert data parallel group ranks list. + + Example - E + M + D parallel + world_size = 16 + model_degree = 2 + expert_degree = 4 # number of experts in same group + mp_group = [0, 1], [2,3], [4,5] ... + data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] + expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] + expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] + + Args: + world_size (int): Distributed world size. + tensor_parallel_size_ (int): Tensor parallel group size. + expert_parallel_size_ (int): Expert parallel group size. + pipeline_parallel_size_ (int): Pipeline parallel group size + use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology + Returns: + Expert parallel group ranks and Expert data parallel group ranks list. + """ + _ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_) + dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_) + _ensure_divisibility(dp_world_size, expert_parallel_size_) + + # Generate data parallel groups + data_parallel_groups = [] + dp_group_size = tensor_parallel_size_ + pp_stride = world_size // pipeline_parallel_size_ + + if use_data_before_expert_parallel_: + dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_ + for pp_stage_start in range(0, world_size, pp_stride): + pp_stage_next = pp_stage_start + pp_stride + for i in range(dp_group_size): + data_parallel_groups.append(list()) + for ds in range(dp_stride): + # [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30] + # [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31] + data_parallel_groups[-1].extend( + list( + range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next, + dp_stride * tensor_parallel_size_))) + else: + for pp_stage_start in range(0, world_size, pp_stride): + pp_stage_next = pp_stage_start + pp_stride + for i in range(dp_group_size): + data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size))) + + expert_parallel_groups = [] + expert_data_parallel_groups = [] + for dp_ranks in data_parallel_groups: + # partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14] + part_ep_groups = [] + for i in range(0, dp_world_size, expert_parallel_size_): + part_ep_groups.append(dp_ranks[i:i + expert_parallel_size_]) + expert_parallel_groups.extend(part_ep_groups) + + # zip part_ep_groups get expert data parallel ranks, e.g [0,8],[2,10],[4,12],[6,14] + for expert_dp_ranks in zip(*part_ep_groups): + expert_data_parallel_groups.append(list(expert_dp_ranks)) + + return expert_parallel_groups, expert_data_parallel_groups + + +def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False): + """ + Create expert and data parallel groups based on MPU (model parallel) group. + + Note: Caller of this function is responsible to check if the groups already exist. + + Example - E + M + D parallel + world_size = 16 + model_degree = 2 + expert_degree = 4 # number of experts in same group + mp_group = [0, 1], [2,3], [4,5] ... + data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] + expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] + expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] + """ + assert dist.is_initialized(), "dist is not initialized" + tensor_parallel_size_ = bwc_tensor_model_parallel_world_size(mpu) + + global expert_tensor_parallel_world_size + expert_tensor_parallel_world_size = tensor_parallel_size_ + + world_size = dist.get_world_size() + rank = dist.get_rank() + dp_world_size = mpu.get_data_parallel_world_size() + pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu) + + _ensure_divisibility(world_size, tensor_parallel_size_) + _ensure_divisibility(dp_world_size, expert_parallel_size_) + + log_dist( + f"Creating deepspeed groups with model parallel size {tensor_parallel_size_}, " + f"pipeline parallel size {pp_world_size}, expert parallel size {expert_parallel_size_}, " + f"world size {world_size}, dp world size {dp_world_size}", [0]) + + global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP + global _EXPERT_PARALLEL_GROUP_RANKS, _EXPERT_DATA_PARALLEL_GROUP_RANKS + + group_name = f"ep_size_{expert_parallel_size_}" + + # Only create groups if they don't already exist + # Need to check conditions outside the group creation loop because of the way torch.dist group creation works + if group_name not in _EXPERT_DATA_PARALLEL_GROUP and group_name not in _EXPERT_PARALLEL_GROUP: + expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks( + world_size, tensor_parallel_size_, expert_parallel_size_, pp_world_size, use_data_before_expert_parallel_) + for ranks in expert_parallel_groups: + group = dist.new_group(ranks) + if rank in list(ranks): + _EXPERT_PARALLEL_GROUP[group_name] = group + _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks + + for ranks in expert_data_parallel_groups: + group = dist.new_group(ranks) + if rank in list(ranks): + _EXPERT_DATA_PARALLEL_GROUP[group_name] = group + _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks + + +def _get_max_expert_size(): + """Get the maximum ep_size from all the created groups.""" + assert _EXPERT_PARALLEL_GROUP is not None, "Warning! Process group not initialized" + keylist = [] + for key in _EXPERT_PARALLEL_GROUP.keys(): + # index 2 is ep_size in the group name: ep_size_ + index = 2 + keylist.append(int(key.split('_')[index])) + return max(keylist) if len(keylist) > 0 else None + + +def _get_max_expert_size_name(): + """Get the name of the group with max. ep_size""" + return f'ep_size_{_get_max_expert_size()}' + + +def _get_max_expert_parallel_group(): + """Get the max expert parallel size.""" + return _get_expert_parallel_group(_get_max_expert_size_name()) + + +def _get_expert_parallel_group(group_name): + """Get the expert parallel group the caller rank belongs to.""" + assert group_name in _EXPERT_PARALLEL_GROUP, \ + 'expert parallel group is not initialized' + return _EXPERT_PARALLEL_GROUP[group_name] + + +def _get_expert_parallel_group_ranks(group_name): + """Get the ranks of the expert parallel group the caller rank belongs to.""" + assert group_name in _EXPERT_PARALLEL_GROUP_RANKS, \ + 'expert parallel group is not initialized' + return _EXPERT_PARALLEL_GROUP_RANKS[group_name] + + +def _get_expert_parallel_group_dict(): + """Get the expert parallel group dict.""" + return _EXPERT_PARALLEL_GROUP + + +def _get_expert_data_parallel_group(group_name): + """Get the expert data parallel group the caller rank belongs to.""" + assert group_name in _EXPERT_DATA_PARALLEL_GROUP, \ + 'expert data parallel group is not initialized' + return _EXPERT_DATA_PARALLEL_GROUP[group_name] + + +def _get_expert_data_parallel_group_ranks(group_name): + """Get the ranks of the expert data parallel group the caller rank belongs to.""" + assert group_name in _EXPERT_DATA_PARALLEL_GROUP_RANKS, \ + 'expert data parallel group is not initialized' + return _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] + + +def _get_expert_data_parallel_group_dict(): + """Get the expert data parallel group dict.""" + return _EXPERT_DATA_PARALLEL_GROUP + + +def _clone_world_group(): + """Create a clone of the world group + Note: We need to clone the dist world group because we + use dist.get_global_rank() utility function in DeepSpeed at many places. + As that function does not work on dist.group.WORLD, we + need to keep a clone of it. + """ + assert dist.is_initialized(), "dist is not initialized" + global _WORLD_GROUP + if _WORLD_GROUP is None: + # If not cloned already, clone the world group + _WORLD_GROUP = dist.new_group(ranks=range(dist.get_world_size())) + return _WORLD_GROUP + + +def _get_local_all_to_all_group(): + assert dist.is_initialized(), 'dist is not initialized' + global _ALL_TO_ALL_GROUP + device_per_node = get_accelerator().device_count() + num_local = dist.get_world_size() // device_per_node + if num_local == 0 and dist.get_world_size() > 0: + assert dist.get_world_size() >= 1, 'num_gpus must >=1, cannot initialize All-To-All' + cur_rank = [] + for i in range(dist.get_world_size()): + cur_rank.append(i) + _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=cur_rank) + elif num_local == 1: + assert dist.get_world_size( + ) == device_per_node, 'num_gpus not equal to device per node, cannot initialize All-To-All' + _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=[i for i in range(device_per_node)]) + else: + assert dist.get_world_size() > device_per_node, 'num_nodes<2 cannot initialize All-To-All' + for i in range(num_local): + local_rank = [j + device_per_node * i for j in range(device_per_node)] + _ALL_TO_ALL_GROUP[f"local_{i}"] = dist.new_group(ranks=local_rank) + + for i in range(device_per_node): + cur_rank = [] + for j in range(num_local): + cur_rank.append(i + j * device_per_node) + _ALL_TO_ALL_GROUP[f"global_{i}"] = dist.new_group(ranks=cur_rank) + return _ALL_TO_ALL_GROUP + + +def _get_data_parallel_group(): + """Get the data parallel group the caller rank belongs to.""" + assert dist.is_initialized(), 'dist is not initialized' + global mpu + if mesh_device is not None: + return mesh_device.get_group(mesh_dim="data_parallel") + if mpu is not None: + if hasattr(mpu, 'initialize_sequence_parallel'): + return None + else: + return mpu.get_data_parallel_group() + + # Return the clone of dist world group + return _clone_world_group() + + +def _get_data_parallel_group_ranks(): + """Get the ranks of data parallel group the caller rank belongs to.""" + assert dist.is_initialized(), \ + 'dist is not initialized' + global mpu + if mpu is not None: + return mpu.get_data_parallel_group_ranks() + # Return all ranks + return range(dist.get_world_size()) + + +def _get_broadcast_src_rank(): + return dist.get_global_rank(_get_sequence_data_parallel_group(), 0) + + +def _get_expert_broadcast_src_rank(group_name): + return dist.get_global_rank(_get_expert_data_parallel_group(group_name), 0) + + +def _get_expert_parallel_world_size(group_name): + """Return world size for the expert parallel group.""" + return dist.get_world_size(group=_get_expert_parallel_group(group_name)) + + +def _get_expert_data_parallel_world_size(group_name): + """Return world size for the expert data parallel group.""" + return dist.get_world_size(group=_get_expert_data_parallel_group(group_name)) + + +def _get_expert_parallel_rank(group_name): + """Return my rank for the expert parallel group.""" + return dist.get_rank(group=_get_expert_parallel_group(group_name)) + + +def _get_expert_parallel_src_rank(group_name): + """Calculate the global rank corresponding to a local rank zero + in the expert parallel group.""" + global_rank = dist.get_rank() + local_world_size = _get_expert_parallel_world_size(group_name) + return (global_rank // local_world_size) * local_world_size + + +def _get_expert_data_parallel_rank(group_name): + """Return my rank for the expert data parallel group.""" + return dist.get_rank(group=_get_expert_data_parallel_group(group_name)) + + +def _get_data_parallel_world_size(): + """Return world size for the data parallel group.""" + if mesh_device is not None: + return dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) + global mpu + if mpu is not None: + if hasattr(mpu, 'initialize_sequence_parallel'): + return None + else: + return mpu.get_data_parallel_world_size() + return dist.get_world_size(group=_get_data_parallel_group()) + + +def _get_model_parallel_world_size(): + """Return world size for the model parallel group.""" + global mpu + if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): + return 1 + return mpu.get_model_parallel_world_size() + + +def _get_data_parallel_rank(): + """Return my rank for the data parallel group.""" + return dist.get_rank(group=_get_data_parallel_group()) + + +def _get_sequence_parallel_world_size(): + """Return world size for the sequence parallel group.""" + global mpu + if mesh_device is not None: + return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel")) + if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'): + return mpu.get_sequence_parallel_world_size() + return 1 + + +def _get_sequence_parallel_rank(): + """Return my rank for the sequence parallel group.""" + global mpu + if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'): + return mpu.get_sequence_parallel_rank() + if mesh_device is not None: + return dist.get_rank(mesh_device.get_group(mesh_dim="sequence_parallel")) + return 0 + + +def _get_sequence_parallel_group(): + global mpu + if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'): + if mesh_device is None: + raise KeyError("No sequence parallel group found") + return mesh_device.get_group(mesh_dim="sequence_parallel") + return mpu.get_sequence_parallel_group() + + +def _get_sequence_data_parallel_world_size(): + """Return world size for the model parallel group.""" + global mpu + if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'): + return mpu.get_sequence_data_parallel_world_size() + return _get_data_parallel_world_size() + + +def _get_sequence_data_parallel_rank(): + """Return my rank for the data parallel group.""" + global mpu + if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'): + return mpu.get_sequence_data_parallel_rank() + return _get_data_parallel_rank() + + +def _get_sequence_data_parallel_group(): + global mpu + # When sequence parallelism is enabled, the process group for zero sharding and + # gradient allreduce must be across both dimensions of data and sequence parallelism. + if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'): + return mpu.get_sequence_data_parallel_group() + return _get_data_parallel_group() + + +def _get_expert_model_parallel_world_size(): + global expert_tensor_parallel_world_size + return expert_tensor_parallel_world_size + + +def _create_zero_param_parallel_group(group_size): + """ + Create parameter partitioning group within ZeRO data parallel groups. + + Example - ZP + D parallel + world_size = 16 + zero_hpz_partition_size = 2 # number of ranks with replicated params (dual partitioning) + zero_param_intra_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - segmented (subgroup) with rep partition + data_parallel_group = [0,1,...,15] - all reduce is on ZeRO model + """ + assert dist.is_initialized() + global _ZERO_PARAM_INTRA_PARALLEL_GROUP + # Only create group if it does not already exist + assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is None, \ + 'ZeRO parameter intra parallel group is already initialized' + + world_size = dist.get_world_size() + rank = dist.get_rank() + + zero_param_parallel_size_ = min(group_size, world_size) + _ensure_divisibility(world_size, zero_param_parallel_size_) + + # Build the ZeRO param intra parallel groups. + for i in range(world_size // zero_param_parallel_size_): + ranks = range(i * zero_param_parallel_size_, (i + 1) * zero_param_parallel_size_) + group = dist.new_group(ranks) + if i == (rank // zero_param_parallel_size_): + _ZERO_PARAM_INTRA_PARALLEL_GROUP = group + + +def _get_zero_param_intra_parallel_group(): + """Get the ZeRO parameter partitioning intra parallel group the caller rank belongs to.""" + #assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is not None, \ + # 'ZeRO parameter partitioning group is not initialized' + #TODO: Add warning + return _ZERO_PARAM_INTRA_PARALLEL_GROUP + + +def _zero_param_parallel_is_initialized(): + """Check if ZeRO data parallel with parameter partititioning groups are initialized.""" + ###TODO: assert that MPU is not set + if _ZERO_PARAM_INTRA_PARALLEL_GROUP is None and _DATA_PARALLEL_GROUP is None: + return False + + +def _get_zero_param_intra_parallel_rank_in_mygroup(): + """Return my rank for the ZeRO parameter inter parallel group.""" + return dist.get_rank(group=_get_zero_param_intra_parallel_group()) + + +def _get_zero_param_intra_parallel_group_world_size(): + """Return world size for the ZeRO parameter parallel group.""" + return dist.get_world_size(group=_get_zero_param_intra_parallel_group()) + + +def _get_zero_param_intra_parallel_group_ranks(): + """Return all ranks for the ZeRO parameter intra parallel group.""" + return dist.get_all_ranks_from_group(group=_get_zero_param_intra_parallel_group()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/init_on_device.py b/venv/lib/python3.10/site-packages/deepspeed/utils/init_on_device.py new file mode 100644 index 0000000000000000000000000000000000000000..52dbf71d9562b73d868f71ee7ec7894a187e00f0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/init_on_device.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from typing import Callable +from torch import Tensor +from packaging import version as pkg_version + + +class OnDevice(object): + """ + Create modules/tensors w. specific devices and dtypes. Examples: + + Create MyModule which consists of many different sub-modules and parameters. In this case we can create + MyModule as a collection of 'meta' tensors by passing `device='meta'` or we can create the module _directly_ + on a CUDA device by passing `device=f'cuda:{local_rank}'` (where `local_rank` is the local GPU id. + + with OnDevice(dtype=torch.float16, device='meta'): + model = MyModel() + + with OnDevice(dtype=torch.float16, device=f'cuda:{local_rank}'): + model = MyModel() + + """ + + _orig_torch_empty = torch.empty + _orig_torch_zeros = torch.zeros + _orig_torch_ones = torch.ones + _orig_torch_full = torch.full + + def __init__(self, dtype, device="meta", enabled=True): + self.dtype = dtype + self.enabled = enabled + self.device = device + + if device == "meta": + if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): + raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") + + def fp_tensor_constructor(self, fn: Callable, target_fp_dtype: torch.dtype) -> Callable: + + def wrapped_fn(*args, **kwargs) -> Tensor: + if kwargs.get("device", None) is None: + kwargs['device'] = self.device + tensor: Tensor = fn(*args, **kwargs) + if tensor.is_floating_point(): + tensor = tensor.to(target_fp_dtype) + return tensor + + return wrapped_fn + + def get_new_tensor_fn_for_dtype(self, dtype: torch.dtype) -> Callable: + + def new_tensor(cls, *args) -> Tensor: + tensor = OnDevice._orig_torch_empty(0, device=self.device).new_empty(*args) + if tensor.is_floating_point(): + tensor = tensor.to(dtype) + return tensor + + return new_tensor + + def __enter__(self): + if not self.enabled: + return + torch.Tensor.__old_new__ = torch.Tensor.__new__ + torch.Tensor.__new__ = self.get_new_tensor_fn_for_dtype(self.dtype) + torch.empty = self.fp_tensor_constructor(self._orig_torch_empty, self.dtype) + torch.zeros = self.fp_tensor_constructor(self._orig_torch_zeros, self.dtype) + torch.ones = self.fp_tensor_constructor(self._orig_torch_ones, self.dtype) + torch.full = self.fp_tensor_constructor(self._orig_torch_full, self.dtype) + + def __exit__(self, exc_type, exc_value, traceback): + if not self.enabled: + return + torch.Tensor.__new__ = torch.Tensor.__old_new__ + torch.empty = self._orig_torch_empty + torch.zeros = self._orig_torch_zeros + torch.ones = self._orig_torch_ones + torch.full = self._orig_torch_full diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/logging.py b/venv/lib/python3.10/site-packages/deepspeed/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..ed691e3985e1f1ca7360fe8dccb30f886013f411 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/logging.py @@ -0,0 +1,163 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import functools +import logging +import sys +import os +import torch +from deepspeed.utils.torch import required_torch_version + +log_levels = { + "debug": logging.DEBUG, + "info": logging.INFO, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, +} + + +class LoggerFactory: + + @staticmethod + def create_logger(name=None, level=logging.INFO): + """create a logger + + Args: + name (str): name of the logger + level: level of logger + + Raises: + ValueError is name is None + """ + + if name is None: + raise ValueError("name for logger cannot be None") + + formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] " + "[%(filename)s:%(lineno)d:%(funcName)s] %(message)s") + + logger_ = logging.getLogger(name) + logger_.setLevel(level) + logger_.propagate = False + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(level) + ch.setFormatter(formatter) + logger_.addHandler(ch) + if required_torch_version(min_version=2.6) and os.getenv("DISABLE_LOGS_WHILE_COMPILING", "0") == "1": + excluded_set = { + item.strip() + for item in os.getenv("LOGGER_METHODS_TO_EXCLUDE_FROM_DISABLE", "").split(",") + } + ignore_set = {'info', 'debug', 'error', 'warning', 'critical', 'exception', 'isEnabledFor'} - excluded_set + for method in ignore_set: + original_logger = getattr(logger_, method) + torch._dynamo.config.ignore_logger_methods.add(original_logger) + return logger_ + + +logger = LoggerFactory.create_logger(name="DeepSpeed", level=logging.INFO) + + +@functools.lru_cache(None) +def warning_once(*args, **kwargs): + """ + This method is identical to `logger.warning()`, but will emit the warning with the same message only once + + Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache. + The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to + another type of cache that includes the caller frame information in the hashing function. + """ + logger.warning(*args, **kwargs) + + +logger.warning_once = warning_once + + +def print_configuration(args, name): + logger.info("{}:".format(name)) + for arg in sorted(vars(args)): + dots = "." * (29 - len(arg)) + logger.info(" {} {} {}".format(arg, dots, getattr(args, arg))) + + +def log_dist(message, ranks=None, level=logging.INFO): + from deepspeed import comm as dist + """Log message when one of following condition meets + + + not dist.is_initialized() + + dist.get_rank() in ranks if ranks is not None or ranks = [-1] + + Args: + message (str) + ranks (list) + level (int) + + """ + should_log = not dist.is_initialized() + ranks = ranks or [] + my_rank = dist.get_rank() if dist.is_initialized() else -1 + if ranks and not should_log: + should_log = ranks[0] == -1 + should_log = should_log or (my_rank in set(ranks)) + if should_log: + final_message = "[Rank {}] {}".format(my_rank, message) + logger.log(level, final_message) + + +def print_json_dist(message, ranks=None, path=None): + from deepspeed import comm as dist + """Print message when one of following condition meets + + + not dist.is_initialized() + + dist.get_rank() in ranks if ranks is not None or ranks = [-1] + + Args: + message (str) + ranks (list) + path (str) + + """ + should_log = not dist.is_initialized() + ranks = ranks or [] + my_rank = dist.get_rank() if dist.is_initialized() else -1 + if ranks and not should_log: + should_log = ranks[0] == -1 + should_log = should_log or (my_rank in set(ranks)) + if should_log: + message['rank'] = my_rank + import json + with open(path, 'w') as outfile: + json.dump(message, outfile) + os.fsync(outfile) + + +def get_current_level(): + """ + Return logger's current log level + """ + return logger.getEffectiveLevel() + + +def should_log_le(max_log_level_str): + """ + Args: + max_log_level_str: maximum log level as a string + + Returns ``True`` if the current log_level is less or equal to the specified log level. Otherwise ``False``. + + Example: + + ``should_log_le("info")`` will return ``True`` if the current log level is either ``logging.INFO`` or ``logging.DEBUG`` + """ + + if not isinstance(max_log_level_str, str): + raise ValueError(f"{max_log_level_str} is not a string") + + max_log_level_str = max_log_level_str.lower() + if max_log_level_str not in log_levels: + raise ValueError(f"{max_log_level_str} is not one of the logging levels") + + return get_current_level() <= log_levels[max_log_level_str] diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/mixed_precision_linkage.py b/venv/lib/python3.10/site-packages/deepspeed/utils/mixed_precision_linkage.py new file mode 100644 index 0000000000000000000000000000000000000000..c97515ca8fef7477e115914e6410bc0e778e3b8a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/mixed_precision_linkage.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +from deepspeed.utils import get_full_hp_param, get_full_hp_grad, get_hp_fragment_mapping +from deepspeed.utils import set_full_hp_param, set_full_hp_grad + + +def link_hp_params(lp_param_list, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, + param_group_index, partition_start, partition_size, dp_group): + local_lp_param_and_offset = _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group) + + for lp_param, lp_start in local_lp_param_and_offset: + lp_param._hp_mapping = get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, + offload_gradient_dict, use_offload, param_group_index, + partition_start, partition_size) + + +def lazy_init_hp_params_optimizer_state(lp_param_list, flat_hp_partition, optimizer_state): + for lp in lp_param_list: + if lp._hp_mapping is not None: + lp._hp_mapping.set_optim_state_fragment(flat_hp_partition, optimizer_state[flat_hp_partition]) + + +def _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group): + current_offset = 0 + param_and_offset_list = [] + partition_end = partition_start + partition_size + index_in_param_group = 0 + for i, lp_param in enumerate(lp_param_list): + lp_param._hp_mapping = None + lp_param._dp_group = dp_group + lp_param.get_full_hp_param = types.MethodType(get_full_hp_param, lp_param) + lp_param.get_full_hp_grad = types.MethodType(get_full_hp_grad, lp_param) + lp_param.set_full_hp_param = types.MethodType(set_full_hp_param, lp_param) + lp_param.set_full_hp_grad = types.MethodType(set_full_hp_grad, lp_param) + + # lp_param overlaps with partition if both are true + # 1) current_offset < partition_end, + # 2) current_offset + lp_param.numel() >= partition_start + lp_param_end = current_offset + lp_param.numel() + if current_offset < partition_end and lp_param_end > partition_start: + param_and_offset_list.append((lp_param, current_offset)) + lp_param._index_in_param_group = index_in_param_group + # Indices for params in this partition/GPU + index_in_param_group += 1 + current_offset += lp_param.numel() + + return param_and_offset_list diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/numa.py b/venv/lib/python3.10/site-packages/deepspeed/utils/numa.py new file mode 100644 index 0000000000000000000000000000000000000000..75f0442a04e53042b9593cb72a7b9aae836ba25c --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/numa.py @@ -0,0 +1,205 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +# return a list of list for cores to numa mapping +# [ +# [ cores for numa 0 ] +# [ cores belong to numa 1 ] +# ... +# ] + +import os +import psutil +import shutil +import subprocess + + +# return a list of list for cores to numa mapping +# [ +# [ cores for numa 0 ] +# [ cores belong to numa 1 ] +# ... +# ] +def get_numa_cores(): + ret = [] + try: + output = subprocess.check_output(['numactl', '--hardware']).decode("utf-8") + except: + return [] + lines = output.split('\n') + for line in lines: + if line.startswith('available:'): + num_numas = int(line.split(' ')[1]) + break + for numa in range(num_numas): + for line in lines: + if line.startswith(f'node {numa} cpus:'): + cores = line.split(' ')[3:] + ret.append([int(core) for core in cores]) + return ret + + +def check_for_numactl_pkg(): + libs = dict( + dpkg=["-l", "numactl", "apt"], + pacman=["-Q", "numactl", "pacman"], + rpm=["-q", "numactl", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + print(f"please install the {lib} package with {tool}") + break + return found + + +def parse_range(rng): + try: + value = int(rng) + return range(value, value + 1) + except ValueError: + # value is not a single number + parts = rng.split('-') + if len(parts) != 2: + raise ValueError("Bad range: '%s', range must be either a number or two number separated by dash" % + (rng, )) + start = int(parts[0]) + end = int(parts[1]) + if start > end: + raise ValueError("Bad range: '%s', range end must larger than or equal to start" % (rng, )) + return range(start, end + 1) + + +# parse comma and dash separated range list into list +# i.e. "0,2-4,6" --> [0, 2, 3, 4, 6] +# rules: +# 1. Range list number be comma separated, each item are either a single number, +# or a range marked by two numbers (both number are included in the range) +# 2. Sub ranges must be in ascend order and not overlap with each other +# 3. No space in the range expression +def parse_range_list(range_str): + number_list = [] + last = -1 + range_list = range_str.split(',') + for sub_range in range_list: + sub_number_list = parse_range(sub_range) + if sub_number_list[0] <= last: + raise ValueError( + "Bad range: '%s', sub ranges must not overlap with each other and should be in ascend order" % + (range_str, )) + last = sub_number_list[-1] + number_list.extend(sub_number_list) + return number_list + + +def get_numactl_cmd(bind_core_list, num_local_procs, local_rank): + numactl_cmd = [] + check_for_numactl_pkg() + if 'KMP_AFFINITY' in os.environ.keys(): + raise ValueError("Environment variable KMP_AFFINITY conflicts with numactl " + "because it interfere with how many CPU cores numactl can set. " + "Unset KMP_AFFINITY before launching deepspeed.\n\n" + "\t$ unset KMP_AFFINITY\n" + "\t$ deepspeed ") + if bind_core_list is not None: + core_list = parse_range_list(bind_core_list) + total_cores = len(core_list) + else: + total_cores = psutil.cpu_count(logical=False) + core_list = range(total_cores) + cores_per_rank = total_cores // num_local_procs + assert cores_per_rank >= 1, "At least one core needs to be assigned to each rank" + core_list_for_rank = core_list[cores_per_rank * local_rank:cores_per_rank * (local_rank + 1)] + numactl_cmd.append("numactl") + + # check if all cores belong to same numa, if true, bind process to that numa domain with -m parameter + numa_cores = get_numa_cores() + num_numas = len(numa_cores) + + numa_mode = "normal" + + non_empty_numa_list = [] + empty_numa_list = [] + previous_numa_cores = [] + numa_node_list = [] + numa_node_list_list = [] + for i in range(num_numas): + # look for empty numa which is HBM numa + if numa_cores[i] == []: + empty_numa_list.append(i) + else: + non_empty_numa_list.append(i) + + # check for fakenuma + if numa_cores[i] == previous_numa_cores: + if numa_node_list == []: + #first duplication, add previous node into list + numa_node_list.append(i - 1) + numa_node_list.append(i) + else: + if numa_node_list != []: + numa_node_list_list.append(numa_node_list) + numa_node_list = [] + previous_numa_cores = numa_cores[i] + if numa_node_list != []: + numa_node_list_list.append(numa_node_list) + + if empty_numa_list != [] and len(empty_numa_list) == len(non_empty_numa_list): + numa_mode = "flat_hbm" + numa_dict = dict(zip(non_empty_numa_list, empty_numa_list)) + elif numa_node_list_list != []: + numa_mode = "fake" + + if numa_mode == "normal": + for i in range(num_numas): + if set(core_list_for_rank) <= set(numa_cores[i]): + numactl_cmd.append("-m") + numactl_cmd.append(f"{i}") + break + elif numa_mode == "flat_hbm": + for i in range(num_numas): + if set(core_list_for_rank) <= set(numa_cores[i]): + numactl_cmd.append("-p") + numactl_cmd.append(f"{numa_dict[i]}") + break + elif numa_mode == "fake": + for i in range(num_numas): + if set(core_list_for_rank) <= set(numa_cores[i]): + for nodes in numa_node_list_list: + if i in nodes: + numactl_cmd.append("-m") + numactl_cmd.append(f"{','.join(map(str, nodes))}") + break + # the following construct break the outer loop if inner loop breaks + else: + continue + break + + numactl_cmd.append("-C") + last_core = core_list_for_rank[0] + first_core = last_core + core_list_str = f"{last_core}" + for core_id in core_list_for_rank[1:]: + if core_id == last_core + 1: + last_core = core_id + continue + else: + if first_core == last_core: + core_list_str = f"{core_list_str},{core_id}" + else: + core_list_str = f"{core_list_str}-{last_core},{core_id}" + first_core = core_id + last_core = core_id + if first_core != last_core: + core_list_str = f"{core_list_str}-{last_core}" + numactl_cmd.append(f"{core_list_str}") + return cores_per_rank, numactl_cmd diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py b/venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py new file mode 100644 index 0000000000000000000000000000000000000000..72d7c863a33f5628fd7c327a6b11652ccba17a01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py @@ -0,0 +1,25 @@ +# 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 diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/tensor_fragment.py b/venv/lib/python3.10/site-packages/deepspeed/utils/tensor_fragment.py new file mode 100644 index 0000000000000000000000000000000000000000..305f8d56fd978e714a666a7e9617a23324e15ef5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/tensor_fragment.py @@ -0,0 +1,481 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from dataclasses import dataclass +from deepspeed import comm as dist +from typing import Dict, List, Callable + + +@dataclass +class fragment_address: + numel: int + start: int + + +@dataclass +class tensor_fragment: + lp_fragment: torch.Tensor + lp_fragment_address: fragment_address + hp_fragment: torch.Tensor + hp_fragment_address: fragment_address + gradient_dict: Dict + offload_gradient_dict: Dict + use_offload: bool + param_group_index: int + optim_fragment: Dict = None + + def update_hp(self): + self.hp_fragment.data.copy_(self.lp_fragment.data) + + def update_lp(self): + self.lp_fragment.data.copy_(self.hp_fragment.data) + + def get_optim_state_fragment(self, key): + if key in self.optim_fragment: + return self.optim_fragment[key] + else: + raise ValueError(f'{key} not found in optimizer state fragment') + + def set_optim_state_fragment(self, flat_hp_partition, optim_fragment): + self.optim_fragment = { + key: value.narrow(0, self.hp_fragment_address.start, self.hp_fragment_address.numel) + for key, value in optim_fragment.items() + if torch.is_tensor(value) and value.shape == flat_hp_partition.shape + } + + def get_hp_fragment_address(self): + return self.hp_fragment_address + + def get_optim_state_keys(self): + return list(self.optim_fragment.keys()) + + def get_hp_fragment(self, optim_state_key=None): + if optim_state_key is None: + return self.hp_fragment + return self.get_optim_state_fragment(optim_state_key) + + def get_lp_grad_fragment(self, index_in_param_group): + if self.use_offload: + gradient_dict = self.offload_gradient_dict + else: + gradient_dict = self.gradient_dict + + if self.param_group_index not in gradient_dict or gradient_dict[self.param_group_index] is None: + raise ValueError("Gradients are only available immediately after backward and before engine step") + + return gradient_dict[self.param_group_index][index_in_param_group] + + +def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys): + for key in opt_keys: + hp_param = flat_hp_tensor + buffer = torch.zeros_like(hp_param) + + for lp in lp_tensors: + if lp._hp_mapping is not None: + hp_fragment_address = lp._hp_mapping.get_hp_fragment_address() + hp_fragment = buffer.narrow(0, hp_fragment_address.start, hp_fragment_address.numel) + hp_fragment.data.copy_(lp._hp_mapping.get_hp_fragment(optim_state_key=key).data) + lp._hp_mapping.hp_fragment = hp_fragment + + optim_state[hp_param][key] = buffer + + +def get_full_hp_param(self, optim_state_key=None): + reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten() + if self._hp_mapping is not None: + lp_frag_address = self._hp_mapping.lp_fragment_address + reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel) + hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key) + reduce_fragment.data.copy_(hp_fragment.data) + dist.all_reduce(reduce_buffer, group=self._dp_group) + return reduce_buffer.reshape_as(self) + + +def set_full_hp_param(self, value, optim_state_key=None): + if self._hp_mapping is not None: + lp_frag_address = self._hp_mapping.lp_fragment_address + value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel) + hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key) + hp_fragment.data.copy_(value_fragment.data) + + +def get_full_hp_grad(self): + reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten() + if self._hp_mapping is not None: + lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group) + hp_grad_fragment = lp_grad_fragment.to(torch.float32).flatten() + + lp_frag_address = self._hp_mapping.lp_fragment_address + reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel) + + if self.view(-1).shape == hp_grad_fragment.shape: + reduce_buffer.data.copy_(hp_grad_fragment.data) + else: + reduce_fragment.data.copy_(hp_grad_fragment.data) + + dist.all_reduce(reduce_buffer, group=self._dp_group) + return reduce_buffer.reshape_as(self) + + +def set_full_hp_grad(self, value): + if self._hp_mapping is not None: + lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group) + lp_frag_address = self._hp_mapping.lp_fragment_address + value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel) + lp_grad_fragment.data.copy_(value_fragment.data.reshape_as(lp_grad_fragment.data)) + if hasattr(self, '_zero_optimizer'): + self._zero_optimizer.update_offload_overflow_tracker(value) + + +def safe_get_full_fp32_param(param): + """Assemble and return the fp32 parameter of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + return param._z3_optimizer.get_full_hp_param(param) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + return param.get_full_hp_param() + return None + + +def safe_set_full_fp32_param(param, value): + """Update the partitioned fp32 parameter of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + value (``torch.Tensor``): New value + """ + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + param._z3_optimizer.set_full_hp_param(value, param) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + param.set_full_hp_param(value) + + +def safe_get_full_optimizer_state(param, optim_state_key): + """Assemble and return the fp32 optimizer state of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device +""" + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + return param._z3_optimizer.get_full_hp_param(param, optim_state_key) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + return param.get_full_hp_param(optim_state_key) + return None + + +def safe_set_full_optimizer_state(param, value, optim_state_key): + """Update the partitioned fp32 optimizer state of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + value (``torch.Tensor``): New value + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) + """ + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + param._z3_optimizer.set_full_hp_param(value, param, optim_state_key) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + param.set_full_hp_param(value, optim_state_key) + + +# TODO: Figure out the correct return dtype +def safe_get_full_grad(param): + """ + Assemble and return the fp32 gradient of a low-precision (e.g., fp16) parameter. + The return data type is that used for gradient accumulation. This is usually the param data type, + but could also be different (e.g., bf16 param training with fp32 gradient accumulation). + + Args: + param (``torch.nn.Parameter``): A model parameter + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + if param.grad is not None: + return param.grad + + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + return param._z3_optimizer.get_fp32_grad_for_param(param) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + return param.get_full_hp_grad() + + return None + + +def safe_set_full_grad(param, value): + """ + Update the partitioned gradient of a low-precision (e.g., fp16) parameter. + To avoid precision issues, the update value should have the data type of + gradient accumulation. + + Args: + param (``torch.nn.Parameter``): A model parameter + value (``torch.Tensor``): The un-partitioned new gradient value. + """ + if param.grad is not None: + param.grad.copy_(value) + elif hasattr(param, 'ds_id'): + # ZeRO stage 3 param + param._z3_optimizer.set_fp32_grad_for_param(value, param) + elif hasattr(param, '_hp_mapping'): + # ZeRO stage 1, 2, and bf16_optimizer params + param.set_full_hp_grad(value) + + +### Local API START ### +def safe_get_local_grad(param): + """ + Get the local gradient partition of a ZeRO-3 partitioned parameter. + The return data type is that used for gradient accumulation. This is usually the param data type, + but could also be different (e.g., bf16 param training with fp32 gradient accumulation). + + Args: + param (``torch.nn.Parameter``): A model parameter + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + return param._z3_optimizer.get_local_fp32_grad_for_param(param) + + +def safe_set_local_grad(param, value): + """ + Update the local gradient partition of a ZeRO-3 partitioned parameter. + To avoid precision issues, the update value should have the data type of + gradient accumulation. + + Args: + param (``torch.nn.Parameter``): A model parameter. + value (``torch.Tensor``): New value of local gradient partition. + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + param._z3_optimizer.set_local_grad_for_param(value, param) + + +def safe_get_local_fp32_param(param): + """Get the local partition of a ZeRO-3 partitioned parameter in fp32 precision. + + Args: + param (``torch.nn.Parameter``): A model parameter. + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + return param._z3_optimizer.get_local_fp32_param(param) + + +def safe_get_local_optimizer_state(param, optim_state_key): + """Get the local optimizer state partition of ZeRO-3 partitioned parameter in fp32 precision. + + Args: + param (``torch.nn.Parameter``): A model parameter + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + return param._z3_optimizer.get_local_fp32_param(param, optim_state_key) + + +def safe_set_local_optimizer_state(param, value, optim_state_key): + """Update the local optimizer state partition of a ZeRO-3 partitioned parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter. + value (``torch.Tensor``): New value of local optimizer state partition. + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer). + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + param._z3_optimizer.set_local_hp_param(value, param, optim_state_key) + + +def safe_set_local_fp32_param(param, value): + """Update the local partition of ZeRO-3 partitioned parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter. + value (``torch.Tensor``): New value of local parameter partition. + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + param._z3_optimizer.set_local_hp_param(value, param) + + +### Local API END ### + + +### VECTORIZED API BEGIN ### +def safe_update_full_grad_vectorized(param_list: List[torch.nn.Parameter], update_func: Callable): + """ + Vectorized update of the partitioned gradients of a list of low-precision (e.g., fp16) parameters. + To avoid precision issues, the update value should have the data type of + gradient accumulation. + + Args: + param_list (``List[torch.nn.Parameter]``): List of model parameters + update_func (``torch.Tensor``): A function that takes current full gradient value and returns new one. + """ + partitioned_grad_params = [] + for p in param_list: + if p.grad is not None: + p.grad.copy_(update_func(p.grad, p)) + elif p.requires_grad: + partitioned_grad_params.append(p) + + if not partitioned_grad_params: + return + + if hasattr(partitioned_grad_params[0], 'ds_id'): + # ZeRO stage 3 param + partitioned_grad_params[0]._z3_optimizer.update_fp32_grad_for_param_vectorized( + update_func, partitioned_grad_params) + elif hasattr(partitioned_grad_params[0], '_hp_mapping'): + # ZeRO stage 1, 2, and bf16_optimizer params + for p in partitioned_grad_params: + old_grad = safe_get_full_grad(p) + new_grad = update_func(old_grad, p) + p.set_full_hp_grad(new_grad) + + +### VECTORIZED API END ### + + +def get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, + param_group_index, partition_start, partition_size): + lp_end = lp_param.numel() + lp_start + hp_start = partition_start + hp_end = partition_start + partition_size + + fragment_start = max(lp_start, hp_start) + fragment_end = min(lp_end, hp_end) + assert fragment_start < fragment_end, \ + f'fragment start {fragment_start} should be < fragment_end {fragment_end}' + + fragment_numel = fragment_end - fragment_start + hp_frag_address = fragment_address(start=fragment_start - hp_start, numel=fragment_numel) + hp_fragment_tensor = flat_hp_partition.narrow(0, hp_frag_address.start, hp_frag_address.numel) + + lp_frag_address = fragment_address(start=fragment_start - lp_start, numel=fragment_numel) + lp_fragment_tensor = lp_param.flatten().narrow(0, lp_frag_address.start, lp_frag_address.numel) + + return tensor_fragment(lp_fragment=lp_fragment_tensor, + lp_fragment_address=lp_frag_address, + hp_fragment=hp_fragment_tensor, + hp_fragment_address=hp_frag_address, + gradient_dict=gradient_dict, + offload_gradient_dict=offload_gradient_dict, + use_offload=use_offload, + param_group_index=param_group_index) + + +''' +Logic for lp_param to hp_param mapping + +lp lp0 lp1 lp2 lp3 lp4 <------- indices/names +lp [ ][ ][ ][ ][ ] <-------- tensors +flat_lp [ ] <-------- flat lp params +flat_hp [ ] <------------------ flat hp partition on current rank +full_hp [ ] <------- full flat hp params + + +lp2 + full numel = 16 + lp_frag + numel = 12 + frag_start = 3 + frag_end = 15 + hp_frag + numel = 12 + frag_start = 0 + frag_end = 11 + + hp_frag.copy_(lp_frag) + + +lp3: + full numel = 4 + lp_frag + numel = 4 + start = 0 + end = 3 + hp_frag + numel = 4 + start = 12 + end = 15 + + +lp4: + full numel = 12 + lp_frag + numel = 4 + start = 0 + end = 3 + hp_frag + numel = 4 + start = 16 + end = 19 + + + +Visual depiction of above +lp { } +flat_lp [ ] +flat_hp ( ) + + +flat_lp [ { ( } ) ] + lx hx ly hy + ly-hx + + +lp { } +flat_lp [ ] +flat_hp ( ) + + +flat_lp [ ( { ) } ] + hx lx hy ly + hy-lx + +lp { } +flat_lp [ ] +flat_hp ( ) + + +flat_lp [ ( { } ) ] + hx lx ly hy + ly-lx + +lp -> (lx, hy) +flat_hp -> (hx, hy) +''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/timer.py b/venv/lib/python3.10/site-packages/deepspeed/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..64ae8ac0e5b4f73d3bf36e8fdddc9dd3ad69d760 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/timer.py @@ -0,0 +1,313 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import time +from numpy import mean +from deepspeed.utils.logging import log_dist +from deepspeed.accelerator import get_accelerator + +FORWARD_MICRO_TIMER = 'fwd_microstep' +FORWARD_GLOBAL_TIMER = 'fwd' +BACKWARD_MICRO_TIMER = 'bwd_microstep' +BACKWARD_GLOBAL_TIMER = 'bwd' +BACKWARD_INNER_MICRO_TIMER = 'bwd_inner_microstep' +BACKWARD_INNER_GLOBAL_TIMER = 'bwd_inner' +BACKWARD_REDUCE_MICRO_TIMER = 'bwd_allreduce_microstep' +BACKWARD_REDUCE_GLOBAL_TIMER = 'bwd_allreduce' +STEP_MICRO_TIMER = 'step_microstep' +STEP_GLOBAL_TIMER = 'step' +TIME_EPSILON = 1e-6 + +try: + import psutil + + PSUTILS_INSTALLED = True +except ImportError: + PSUTILS_INSTALLED = False + pass + + +class CudaEventTimer(object): + + def __init__(self, start_event: get_accelerator().Event, end_event: get_accelerator().Event): + self.start_event = start_event + self.end_event = end_event + + def get_elapsed_msec(self): + get_accelerator().current_stream().wait_event(self.end_event) + self.end_event.synchronize() + return self.start_event.elapsed_time(self.end_event) + + +class SynchronizedWallClockTimer: + """Group of timers. Borrowed from Nvidia Megatron code""" + + class Timer: + """Timer.""" + + def __init__(self, name): + self.name_ = name + self.started_ = False + self.event_timers = [] + self.use_host_timer = get_accelerator().use_host_timers() + self.start_event = None + self.elapsed_records = None + self.start_time = 0.0 + self.end_time = 0.0 + + def start(self): + """Start the timer.""" + assert not self.started_, f"{self.name_} timer has already been started" + if self.use_host_timer: + self.start_time = time.time() + else: + event_class = get_accelerator().Event + self.start_event = event_class(enable_timing=True) + self.start_event.record() + self.started_ = True + + def stop(self, reset=False, record=False): + """Stop the timer.""" + assert self.started_, "timer is not started" + event_class = get_accelerator().Event + if self.use_host_timer: + self.end_time = time.time() + self.event_timers.append(self.end_time - self.start_time) + else: + event_class = get_accelerator().Event + end_event = event_class(enable_timing=True) + end_event.record() + self.event_timers.append(CudaEventTimer(self.start_event, end_event)) + self.start_event = None + self.started_ = False + + def _get_elapsed_msec(self): + if self.use_host_timer: + self.elapsed_records = [et * 1000.0 for et in self.event_timers] + else: + self.elapsed_records = [et.get_elapsed_msec() for et in self.event_timers] + self.event_timers.clear() + return sum(self.elapsed_records) + + def reset(self): + """Reset timer.""" + self.started_ = False + self.start_event = None + self.elapsed_records = None + self.event_timers.clear() + + def elapsed(self, reset=True): + """Calculate the elapsed time.""" + started_ = self.started_ + # If the timing in progress, end it first. + if self.started_: + self.stop() + # Get the elapsed time. + elapsed_ = self._get_elapsed_msec() + # Reset the elapsed time + if reset: + self.reset() + # If timing was in progress, set it back. + if started_: + self.start() + return elapsed_ + + def mean(self): + self.elapsed(reset=False) + return trim_mean(self.elapsed_records, 0.1) + + def __init__(self): + self.timers = {} + + def get_timers(self): + return self.timers + + def __call__(self, name): + if name not in self.timers: + self.timers[name] = self.Timer(name) + return self.timers[name] + + @staticmethod + def memory_usage(): + alloc = "mem_allocated: {:.4f} GB".format(get_accelerator().memory_allocated() / (1024 * 1024 * 1024)) + max_alloc = "max_mem_allocated: {:.4f} GB".format(get_accelerator().max_memory_allocated() / + (1024 * 1024 * 1024)) + cache = "cache_allocated: {:.4f} GB".format(get_accelerator().memory_cached() / (1024 * 1024 * 1024)) + max_cache = "max_cache_allocated: {:.4f} GB".format(get_accelerator().max_memory_cached() / + (1024 * 1024 * 1024)) + return " | {} | {} | {} | {}".format(alloc, max_alloc, cache, max_cache) + + def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None): + """Log a group of timers.""" + assert normalizer > 0.0 + string = f"time (ms)" + for name in names: + if name in self.timers: + elapsed_time = (self.timers[name].elapsed(reset=reset) / normalizer) + string += " | {}: {:.2f}".format(name, elapsed_time) + + log_dist(string, ranks=ranks or [0]) + + def get_mean(self, names, normalizer=1.0, reset=True): + """Get the mean of a group of timers.""" + assert normalizer > 0.0 + means = {} + for name in names: + if name in self.timers: + elapsed_time = (self.timers[name].mean() * 1000.0 / normalizer) + means[name] = elapsed_time + return means + + +class NoopTimer: + + class Timer: + + def start(self): + ... + + def reset(self): + ... + + def stop(self, **kwargs): + ... + + def elapsed(self, **kwargs): + return 0 + + def mean(self): + return 0 + + def __init__(self): + self.timer = self.Timer() + + def __call__(self, name): + return self.timer + + def get_timers(self): + return {} + + def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None): + ... + + def get_mean(self, names, normalizer=1.0, reset=True): + ... + + +class ThroughputTimer: + + def __init__(self, config, batch_size, start_step=2, steps_per_output=None, monitor_memory=False, logging_fn=None): + from deepspeed.utils import logger + self.config = config + self.start_time = 0 + self.end_time = 0 + self.started = False + self.batch_size = 1 if batch_size is None else batch_size + self.start_step = start_step + self.epoch_count = 0 + self.micro_step_count = 0 + self.global_step_count = 0 + self.total_elapsed_time = 0 + self.step_elapsed_time = 0 + self.steps_per_output = steps_per_output + self.monitor_memory = monitor_memory + self.logging = logging_fn + if self.logging is None: + self.logging = logger.info + self.initialized = False + + if self.monitor_memory and not PSUTILS_INSTALLED: + raise ImportError("Unable to import 'psutils', please install package") + + def update_epoch_count(self): + self.epoch_count += 1 + self.micro_step_count = 0 + + def _init_timer(self): + self.initialized = True + + def start(self): + if not self.config.enabled: + return + self._init_timer() + self.started = True + if self.global_step_count >= self.start_step: + if self.config.synchronized: + get_accelerator().synchronize() + self.start_time = time.time() + + def _is_report_boundary(self): + if self.steps_per_output is None: + return False + return self.global_step_count % self.steps_per_output == 0 + + def stop(self, global_step=False, report_speed=True): + if not self.config.enabled or not self.started: + return + self.started = False + self.micro_step_count += 1 + if global_step: + self.global_step_count += 1 + + if self.start_time > 0: + if self.config.synchronized: + get_accelerator().synchronize() + self.end_time = time.time() + duration = self.end_time - self.start_time + self.total_elapsed_time += duration + self.step_elapsed_time += duration + + if global_step: + if report_speed and self._is_report_boundary(): + self.logging( + "epoch={}/micro_step={}/global_step={}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, " + "MemAllocated={}GB, MaxMemAllocated={}GB".format( + self.epoch_count, + self.micro_step_count, + self.global_step_count, + self.avg_samples_per_sec(), + self.batch_size / (self.step_elapsed_time + TIME_EPSILON), + round(get_accelerator().memory_allocated() / 1024**3, 2), + round(get_accelerator().max_memory_allocated() / 1024**3, 2), + )) + if self.monitor_memory: + virt_mem = psutil.virtual_memory() + swap = psutil.swap_memory() + self.logging("epoch={}/micro_step={}/global_step={}, vm %: {}, swap %: {}".format( + self.epoch_count, + self.micro_step_count, + self.global_step_count, + virt_mem.percent, + swap.percent, + )) + self.step_elapsed_time = 0 + + def avg_samples_per_sec(self): + if self.global_step_count > 0: + total_step_offset = self.global_step_count - self.start_step + avg_time_per_step = self.total_elapsed_time / total_step_offset + # training samples per second + return self.batch_size / avg_time_per_step + return float("-inf") + + +def trim_mean(data, trim_percent): + """Compute the trimmed mean of a list of numbers. + + Args: + data (list): List of numbers. + trim_percent (float): Percentage of data to trim. + + Returns: + float: Trimmed mean. + """ + assert 0.0 <= trim_percent <= 1.0 + n = len(data) + # Account for edge case of empty list + if len(data) == 0: + return 0 + data.sort() + k = int(round(n * (trim_percent))) + return mean(data[k:n - k]) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/torch.py b/venv/lib/python3.10/site-packages/deepspeed/utils/torch.py new file mode 100644 index 0000000000000000000000000000000000000000..1d32775fe64a8e9bf1bff5df6aaa111c2974d53d --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/torch.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from packaging import version as pkg_version + +import torch + + +def required_torch_version(min_version=None, max_version=None): + assert min_version or max_version, "Must provide a min_version or max_version argument" + + torch_version = pkg_version.parse(torch.__version__) + + if min_version and pkg_version.parse(str(min_version)) > torch_version: + return False + + if max_version and pkg_version.parse(str(max_version)) < torch_version: + return False + + return True + + +def register_grad_hook(param, hook): + if required_torch_version(min_version=2.1): + return param.register_post_accumulate_grad_hook(hook) + else: + param_tmp = param.expand_as(param) + grad_acc = param_tmp.grad_fn.next_functions[0][0] + return grad_acc.register_hook(hook) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/types.py b/venv/lib/python3.10/site-packages/deepspeed/utils/types.py new file mode 100644 index 0000000000000000000000000000000000000000..96b5df625965401f52f6fd9b60396821943fe177 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/types.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import IntEnum + + +class ActivationFuncType(IntEnum): + UNKNOWN = 0 + GELU = 1 + ReLU = 2 + GATED_GELU = 3 + GATED_SILU = 4 + + +GATED_ACTIVATION_TYPES = [ + ActivationFuncType.GATED_GELU, + ActivationFuncType.GATED_SILU, +] + + +class NormType(IntEnum): + UNKNOWN = 0 + LayerNorm = 1 + GroupNorm = 2 + RMSNorm = 3 diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/z3_leaf_module.py b/venv/lib/python3.10/site-packages/deepspeed/utils/z3_leaf_module.py new file mode 100644 index 0000000000000000000000000000000000000000..14e8ae2d28235c46016b9030b3492a2eb4e7122a --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/z3_leaf_module.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from typing import List, Type, Union + + +def z3_leaf_module(model: torch.nn.Module) -> bool: + """Returns whether a module in `model` has been flagged as a 'leaf' module. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + Returns: + bool: Whether the module has been flagged as a 'leaf' module. + """ + return hasattr(model, '_z3_leaf') and model._z3_leaf + + +def z3_leaf_parameter(model: torch.nn.Parameter) -> bool: + """Returns whether a parameter belongs to a leaf module. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Parameter): The parameter to which the leaf module flag will be applied. + Returns: + bool: Whether the parameter belongs to a leaf module. + """ + return hasattr(model, 'ds_z3_leaf_module') + + +def get_z3_leaf_modules(model: torch.nn.Module) -> List[torch.nn.Module]: + """Returns a list of modules in `model` that have been flagged as 'leaf' modules. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + Returns: + List[torch.nn.Module]: A list of modules that have been flagged as 'leaf' modules. + """ + return [module for module in model.modules() if z3_leaf_module(module)] + + +def set_z3_leaf_module(model: torch.nn.Module, flag: bool): + model._z3_leaf = flag + + +def _do_set_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: Union[List[Type], List[str]], + flag: bool) -> List[torch.nn.Module]: + assert all(isinstance(module_class, (type, str) ) for module_class in leaf_module_classes), \ + f'leaf_module_classes must be a list of types or names, got {leaf_module_classes}' + + leaf_modules = [] + + def _set_z3_leaf_flag(model: torch.nn.Module): + nonlocal leaf_modules + for module in leaf_module_classes: + if (isinstance(module, type) and model.__class__ == module) or \ + (isinstance(module, str) and model.__class__.__name__ == module): + model._z3_leaf = flag + leaf_modules.append(model) + + model.apply(_set_z3_leaf_flag) + + if len(leaf_modules) == 0: + raise ValueError(f'No modules of type {leaf_module_classes} found in model {model}') + + return leaf_modules + + +def set_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: Union[List[Type], + List[str]]) -> List[torch.nn.Module]: + """Sets a flag within a module in `model` to instruct ZeRO3 to stop setting hooks recursively when it encounters a module class listed in `leaf_module_classes`. + This is particularly useful in the context of Mixture of Experts (MoE) models. In MoE models, the computation order of experts varies across forward passes. This variability can disrupt ZeRO3's functionality, as ZeRO3 relies on tracking the computation order of modules to prefetch parameters efficiently. By designating a module as a 'leaf' node, ZeRO3 will prefetch parameters for all child modules upon entering the module. + Another scenario where this functionality is beneficial is in models with excessively fine-grained nested modules, where it helps to avoid the overhead associated with hooks. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules. + Returns: + List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`. + """ + return _do_set_z3_leaf_modules(model, leaf_module_classes, True) + + +def unset_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: List[Type]) -> List[torch.nn.Module]: + """Unsets a flag within a module in `model` to instruct ZeRO3 to resume setting hooks recursively when it encounters a module class listed in `leaf_module_classes`. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules. + Returns: + List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`. + """ + return _do_set_z3_leaf_modules(model, leaf_module_classes, False) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/zero_to_fp32.py b/venv/lib/python3.10/site-packages/deepspeed/utils/zero_to_fp32.py new file mode 100644 index 0000000000000000000000000000000000000000..0e759146cadd92ddfefab3680146c2bd6a2b5c04 --- /dev/null +++ b/venv/lib/python3.10/site-packages/deepspeed/utils/zero_to_fp32.py @@ -0,0 +1,760 @@ +#!/usr/bin/env python + +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets +# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in +# the future. Once extracted, the weights don't require DeepSpeed and can be used in any +# application. +# +# example: +# python zero_to_fp32.py . output_dir/ +# or +# python zero_to_fp32.py . output_dir/ --safe_serialization + +import argparse +import torch +import glob +import math +import os +import re +import gc +import json +import numpy as np +from tqdm import tqdm +from collections import OrderedDict +from dataclasses import dataclass + +# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with +# DeepSpeed data structures it has to be available in the current python environment. +from deepspeed.utils import logger +from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS, + FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES, + FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS) + + +@dataclass +class zero_model_state: + buffers: dict() + param_shapes: dict() + shared_params: list + ds_version: int + frozen_param_shapes: dict() + frozen_param_fragments: dict() + + +debug = 0 + +# load to cpu +device = torch.device('cpu') + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split(r'(\d+)', text)] + + +def get_model_state_file(checkpoint_dir, zero_stage): + if not os.path.isdir(checkpoint_dir): + raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") + + # there should be only one file + if zero_stage <= 2: + file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt") + elif zero_stage == 3: + file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") + + if not os.path.exists(file): + raise FileNotFoundError(f"can't find model states file at '{file}'") + + return file + + +def get_checkpoint_files(checkpoint_dir, glob_pattern): + # XXX: need to test that this simple glob rule works for multi-node setup too + ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys) + + if len(ckpt_files) == 0: + raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'") + + return ckpt_files + + +def get_optim_files(checkpoint_dir): + return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt") + + +def get_model_state_files(checkpoint_dir): + return get_checkpoint_files(checkpoint_dir, "*_model_states.pt") + + +def parse_model_states(files): + zero_model_states = [] + for file in files: + state_dict = torch.load(file, map_location=device, weights_only=False) + + if BUFFER_NAMES not in state_dict: + raise ValueError(f"{file} is not a model state checkpoint") + buffer_names = state_dict[BUFFER_NAMES] + if debug: + print("Found buffers:", buffer_names) + + # recover just the buffers while restoring them to fp32 if they were saved in fp16 + buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names} + param_shapes = state_dict[PARAM_SHAPES] + + # collect parameters that are included in param_shapes + param_names = [] + for s in param_shapes: + for name in s.keys(): + param_names.append(name) + + # update with frozen parameters + frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None) + if frozen_param_shapes is not None: + if debug: + print(f"Found frozen_param_shapes: {frozen_param_shapes}") + param_names += list(frozen_param_shapes.keys()) + + # handle shared params + shared_params = [[k, v] for k, v in state_dict["shared_params"].items()] + + ds_version = state_dict.get(DS_VERSION, None) + + frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None) + + z_model_state = zero_model_state(buffers=buffers, + param_shapes=param_shapes, + shared_params=shared_params, + ds_version=ds_version, + frozen_param_shapes=frozen_param_shapes, + frozen_param_fragments=frozen_param_fragments) + zero_model_states.append(z_model_state) + + return zero_model_states + + +def parse_optim_states(files, ds_checkpoint_dir): + total_files = len(files) + state_dicts = [] + for f in tqdm(files, desc='Loading checkpoint shards'): + state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False) + # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights + # and also handle the case where it was already removed by another helper script + state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) + state_dicts.append(state_dict) + + if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]: + raise ValueError(f"{files[0]} is not a zero checkpoint") + zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE] + world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT] + + # For ZeRO-2 each param group can have different partition_count as data parallelism for expert + # parameters can be different from data parallelism for non-expert parameters. So we can just + # use the max of the partition_count to get the dp world_size. + + if type(world_size) is list: + world_size = max(world_size) + + if world_size != total_files: + raise ValueError( + f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. " + "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes." + ) + + # the groups are named differently in each stage + if zero_stage <= 2: + fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS + elif zero_stage == 3: + fp32_groups_key = FP32_FLAT_GROUPS + else: + raise ValueError(f"unknown zero stage {zero_stage}") + + fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))] + return zero_stage, world_size, fp32_flat_groups + + +def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters): + """ + Returns fp32 state_dict reconstructed from ds checkpoint + + Args: + - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are) + + """ + print(f"Processing zero checkpoint '{ds_checkpoint_dir}'") + + optim_files = get_optim_files(ds_checkpoint_dir) + zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir) + print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}") + + model_files = get_model_state_files(ds_checkpoint_dir) + + zero_model_states = parse_model_states(model_files) + print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}') + + if zero_stage <= 2: + return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters) + elif zero_stage == 3: + return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters) + + +def _zero2_merge_frozen_params(state_dict, zero_model_states): + if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: + return + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + frozen_param_fragments = zero_model_states[0].frozen_param_fragments + + if debug: + num_elem = sum(s.numel() for s in frozen_param_shapes.values()) + print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') + + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = sum([p.numel() for p in frozen_param_fragments.values()]) + print(f'Frozen params: Have {avail_numel} numels to process.') + print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') + + total_params = 0 + total_numel = 0 + for name, shape in frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + state_dict[name] = frozen_param_fragments[name] + + if debug: + print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") + + print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +def _has_callable(obj, fn): + attr = getattr(obj, fn, None) + return callable(attr) + + +def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): + param_shapes = zero_model_states[0].param_shapes + + # Reconstruction protocol: + # + # XXX: document this + + if debug: + for i in range(world_size): + for j in range(len(fp32_flat_groups[0])): + print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}") + + # XXX: memory usage doubles here (zero2) + num_param_groups = len(fp32_flat_groups[0]) + merged_single_partition_of_fp32_groups = [] + for i in range(num_param_groups): + merged_partitions = [sd[i] for sd in fp32_flat_groups] + full_single_fp32_vector = torch.cat(merged_partitions, 0) + merged_single_partition_of_fp32_groups.append(full_single_fp32_vector) + avail_numel = sum( + [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups]) + + if debug: + wanted_params = sum([len(shapes) for shapes in param_shapes]) + wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes]) + # not asserting if there is a mismatch due to possible padding + print(f"Have {avail_numel} numels to process.") + print(f"Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + total_numel = 0 + total_params = 0 + for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups): + offset = 0 + avail_numel = full_single_fp32_vector.numel() + for name, shape in shapes.items(): + + unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape) + total_numel += unpartitioned_numel + total_params += 1 + + if debug: + print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") + state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape) + offset += unpartitioned_numel + + # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and + # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex + # paddings performed in the code it's almost impossible to predict the exact numbers w/o the + # live optimizer object, so we are checking that the numbers are within the right range + align_to = 2 * world_size + + def zero2_align(x): + return align_to * math.ceil(x / align_to) + + if debug: + print(f"original offset={offset}, avail_numel={avail_numel}") + + offset = zero2_align(offset) + avail_numel = zero2_align(avail_numel) + + if debug: + print(f"aligned offset={offset}, avail_numel={avail_numel}") + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements") + + +def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters): + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + if not exclude_frozen_parameters: + _zero2_merge_frozen_params(state_dict, zero_model_states) + + _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def zero3_partitioned_param_info(unpartitioned_numel, world_size): + remainder = unpartitioned_numel % world_size + padding_numel = (world_size - remainder) if remainder else 0 + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + return partitioned_numel, padding_numel + + +def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states): + if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: + return + + if debug: + for i in range(world_size): + num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values()) + print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size + print(f'Frozen params: Have {avail_numel} numels to process.') + print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') + + total_params = 0 + total_numel = 0 + for name, shape in zero_model_states[0].frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states) + state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape) + + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) + + if debug: + print( + f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +class GatheredTensor: + """ + A pseudo tensor that collects partitioned weights. + It is more memory efficient when there are multiple groups. + """ + + def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape): + self.flat_groups = flat_groups + self.flat_groups_offset = flat_groups_offset + self.offset = offset + self.partitioned_numel = partitioned_numel + self.shape = shape + self.dtype = self.flat_groups[0][0].dtype + + def contiguous(self): + """ + Merge partitioned weights from flat_groups into a single tensor. + """ + end_idx = self.offset + self.partitioned_numel + world_size = len(self.flat_groups) + pad_flat_param_chunks = [] + + for rank_i in range(world_size): + # for each rank, we need to collect weights from related group/groups + flat_groups_at_rank_i = self.flat_groups[rank_i] + start_group_id = None + end_group_id = None + for group_id in range(len(self.flat_groups_offset)): + if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]: + start_group_id = group_id + if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]: + end_group_id = group_id + break + # collect weights from related group/groups + for group_id in range(start_group_id, end_group_id + 1): + flat_tensor = flat_groups_at_rank_i[group_id] + start_offset = self.offset - self.flat_groups_offset[group_id] + end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id] + pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset]) + + # collect weights from all ranks + pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0) + param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous() + return param + + +def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): + param_shapes = zero_model_states[0].param_shapes + avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size + + # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each + # param, re-consolidating each param, while dealing with padding if any + + # merge list of dicts, preserving order + param_shapes = {k: v for d in param_shapes for k, v in d.items()} + + if debug: + for i in range(world_size): + print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}") + + wanted_params = len(param_shapes) + wanted_numel = sum(shape.numel() for shape in param_shapes.values()) + # not asserting if there is a mismatch due to possible padding + avail_numel = fp32_flat_groups[0].numel() * world_size + print(f"Trainable params: Have {avail_numel} numels to process.") + print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + offset = 0 + total_numel = 0 + total_params = 0 + flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]])) + for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'): + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + total_params += 1 + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) + + if debug: + print( + f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + # memory efficient tensor + tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape) + state_dict[name] = tensor + offset += partitioned_numel + + offset *= world_size + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements") + + +def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters): + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + if not exclude_frozen_parameters: + _zero3_merge_frozen_params(state_dict, world_size, zero_model_states) + + _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def to_torch_tensor(state_dict, return_empty_tensor=False): + """ + Convert state_dict of GatheredTensor to torch tensor + """ + torch_state_dict = {} + converted_tensors = {} + for name, tensor in state_dict.items(): + tensor_id = id(tensor) + if tensor_id in converted_tensors: # shared tensors + shared_tensor = torch_state_dict[converted_tensors[tensor_id]] + torch_state_dict[name] = shared_tensor + else: + converted_tensors[tensor_id] = name + if return_empty_tensor: + torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype) + else: + torch_state_dict[name] = tensor.contiguous() + return torch_state_dict + + +def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, + tag=None, + exclude_frozen_parameters=False, + lazy_mode=False): + """ + Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with + ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example + via a model hub. + + Args: + - ``checkpoint_dir``: path to the desired checkpoint folder + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14`` + - ``exclude_frozen_parameters``: exclude frozen parameters + - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient. + Convert the pesduo tensor to torch tensor by ``.contiguous()`` + + Returns: + - pytorch ``state_dict`` + + A typical usage might be :: + + from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + # do the training and checkpoint saving + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu + model = model.cpu() # move to cpu + model.load_state_dict(state_dict) + # submit to model hub or save the model to share with others + + In this example the ``model`` will no longer be usable in the deepspeed context of the same + application. i.e. you will need to re-initialize the deepspeed engine, since + ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. + + If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead. + + Note: the above usage may not work if your application doesn't have sufficient free CPU memory. + You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with + the checkpoint. Or you can load state_dict in lazy mode :: + + from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu + for name, lazy_tensor in state_dict.item(): + tensor = lazy_tensor.contiguous() # to cpu + print(name, tensor) + # del tensor to release memory if it no longer in use + """ + if tag is None: + latest_path = os.path.join(checkpoint_dir, 'latest') + if os.path.isfile(latest_path): + with open(latest_path, 'r') as fd: + tag = fd.read().strip() + else: + raise ValueError(f"Unable to find 'latest' file at {latest_path}") + + ds_checkpoint_dir = os.path.join(checkpoint_dir, tag) + + if not os.path.isdir(ds_checkpoint_dir): + raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist") + + state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters) + if lazy_mode: + return state_dict + else: + return to_torch_tensor(state_dict) + + +def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, + output_dir, + max_shard_size="5GB", + safe_serialization=False, + tag=None, + exclude_frozen_parameters=False): + """ + Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be + loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed. + + Args: + - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) + - ``output_dir``: directory to the pytorch fp32 state_dict output files + - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB + - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` + - ``exclude_frozen_parameters``: exclude frozen parameters + """ + + # Dependency pre-check + if safe_serialization: + try: + from safetensors.torch import save_file + except ImportError: + print('If you want to use `safe_serialization`, please `pip install safetensors`') + raise + if max_shard_size is not None: + try: + from huggingface_hub import split_torch_state_dict_into_shards + except ImportError: + print('If you want to use `max_shard_size`, please `pip install huggingface_hub`') + raise + + # Convert zero checkpoint to state_dict + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, + tag, + exclude_frozen_parameters, + lazy_mode=True) + + # Shard the model if it is too big. + weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin" + if max_shard_size is not None: + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") + # an memory-efficient approach for sharding + empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True) + state_dict_split = split_torch_state_dict_into_shards(empty_state_dict, + filename_pattern=filename_pattern, + max_shard_size=max_shard_size) + else: + from collections import namedtuple + StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"]) + state_dict_split = StateDictSplit(is_sharded=False, + filename_to_tensors={weights_name: list(state_dict.keys())}) + + # Save the model by shard + os.makedirs(output_dir, exist_ok=True) + filename_to_tensors = state_dict_split.filename_to_tensors.items() + for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"): + shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors} + shard_state_dict = to_torch_tensor(shard_state_dict) + output_path = os.path.join(output_dir, shard_file) + if safe_serialization: + save_file(shard_state_dict, output_path, metadata={"format": "pt"}) + else: + torch.save(shard_state_dict, output_path) + # release the memory of current shard + for tensor_name in list(shard_state_dict.keys()): + del state_dict[tensor_name] + del shard_state_dict[tensor_name] + del shard_state_dict + gc.collect() + + # Save index if sharded + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json" + save_index_file = os.path.join(output_dir, save_index_file) + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + + +def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None): + """ + 1. Put the provided model to cpu + 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` + 3. Load it into the provided model + + Args: + - ``model``: the model object to update + - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` + + Returns: + - ``model`: modified model + + Make sure you have plenty of CPU memory available before you call this function. If you don't + have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it + conveniently placed for you in the checkpoint folder. + + A typical usage might be :: + + from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint + model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir) + # submit to model hub or save the model to share with others + + Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context + of the same application. i.e. you will need to re-initialize the deepspeed engine, since + ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. + + """ + logger.info(f"Extracting fp32 weights") + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag) + + logger.info(f"Overwriting model with fp32 weights") + model = model.cpu() + model.load_state_dict(state_dict, strict=False) + + return model + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint_dir", + type=str, + help="path to the desired checkpoint folder, e.g., path/checkpoint-12") + parser.add_argument("output_dir", + type=str, + help="directory to the pytorch fp32 state_dict output files" + "(e.g. path/checkpoint-12-output/)") + parser.add_argument( + "--max_shard_size", + type=str, + default="5GB", + help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size" + "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`" + "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances" + "without CPU OOM issues.") + parser.add_argument( + "--safe_serialization", + default=False, + action='store_true', + help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).") + parser.add_argument("-t", + "--tag", + type=str, + default=None, + help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1") + parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters") + parser.add_argument("-d", "--debug", action='store_true', help="enable debug") + args = parser.parse_args() + + debug = args.debug + + convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, + args.output_dir, + max_shard_size=args.max_shard_size, + safe_serialization=args.safe_serialization, + tag=args.tag, + exclude_frozen_parameters=args.exclude_frozen_parameters) diff --git a/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/METADATA b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..f6ddf4f0e77a8a767d533125c4c7f37fba08c896 --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/METADATA @@ -0,0 +1,64 @@ +Metadata-Version: 2.4 +Name: hf_transfer +Version: 0.1.9 +Classifier: Programming Language :: Rust +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +License-File: LICENSE +Summary: Speed up file transfers with the Hugging Face Hub. +Requires-Python: >=3.7 +Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM +Project-URL: Issues, https://github.com/huggingface/hf_transfer/issues +Project-URL: Repository, https://github.com/huggingface/hf_transfer.git + +# HF Transfer + +Speed up file transfers with the Hub. + +# DISCLAIMER + +This library is a power user tool, to go beyond `~500MB/s` on very high bandwidth +network, where Python cannot cap out the available bandwidth. + +This is *not* meant to be a general usability tool. +It purposefully lacks progressbars and comes generally as-is. + +Please file issues *only* if there's an issue on the underlying downloaded file. + +## Contributing + +```sh +python3 -m venv ~/.venv/hf_transfer +source ~/.venv/hf_transfer/bin/activate +pip install maturin +maturin develop +``` + +### `huggingface_hub` + +If you are working on changes with `huggingface_hub` + +```sh +git clone git@github.com:huggingface/huggingface_hub.git +# git clone https://github.com/huggingface/huggingface_hub.git + +cd huggingface_hub +python3 -m pip install -e ".[quality]" +``` + +You can use the following test script: + +```py +import os + +# os.environ["HF_ENDPOINT"] = "http://localhost:5564" +os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" + +from huggingface_hub import HfApi, logging + +logging.set_verbosity_debug() +hf = HfApi() +hf.upload_file(path_or_fileobj="/path/to/my/repo/some_file", path_in_repo="some_file", repo_id="my/repo", repo_type="model") +``` + + diff --git a/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/RECORD b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..41ad2b97e7ef815a8c5670da224011be269aaa7f --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/RECORD @@ -0,0 +1,9 @@ +hf_transfer-0.1.9.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +hf_transfer-0.1.9.dist-info/METADATA,sha256=cISfuZ8iKrpsZCWkxI-IFRcC3htDBOCMXXT-liJ9FJA,1747 +hf_transfer-0.1.9.dist-info/RECORD,, +hf_transfer-0.1.9.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +hf_transfer-0.1.9.dist-info/WHEEL,sha256=473T-iBDPc6DHc-9q66Gpnn0HSHmHhv_27uSTUVy86w,127 +hf_transfer-0.1.9.dist-info/licenses/LICENSE,sha256=8dFfk5jBdKElwqtcWTyhfvBqN7VD_-J4bwXLJC-rszk,11346 +hf_transfer/__init__.py,sha256=0OjYn5WBWp_eK3hJKyYtyMWxAStjTCTRjSh_eXoOldE,127 +hf_transfer/__pycache__/__init__.cpython-310.pyc,, +hf_transfer/hf_transfer.abi3.so,sha256=ZQGGGzyeqivKo3-kIKYQQ7m-wW0WILRHX8AW-Q3x1bU,9653960 diff --git a/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/WHEEL b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..4a2f9abc544203bc4e582b0d692e0ddbf94eed91 --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: maturin (1.8.1) +Root-Is-Purelib: false +Tag: cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64 diff --git a/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..98f5731a66fcac01ff798880d5c396621400b3f7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer-0.1.9.dist-info/licenses/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2023 HuggingFace Inc. + + 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. diff --git a/venv/lib/python3.10/site-packages/hf_transfer/__init__.py b/venv/lib/python3.10/site-packages/hf_transfer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ec78d8378ffa59f297d250095503af8f11a6f0db --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer/__init__.py @@ -0,0 +1,5 @@ +from .hf_transfer import * + +__doc__ = hf_transfer.__doc__ +if hasattr(hf_transfer, "__all__"): + __all__ = hf_transfer.__all__ \ No newline at end of file diff --git a/venv/lib/python3.10/site-packages/hf_transfer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/hf_transfer/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..053f7ac25980f162be6bcbc9dd17f1f9a47afee4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/hf_transfer/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/hf_transfer/hf_transfer.abi3.so b/venv/lib/python3.10/site-packages/hf_transfer/hf_transfer.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..210f0bf9408780d2e376d254f8c874e3f45b1971 --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_transfer/hf_transfer.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6501861b3c9eaa2bcaa37fa420a61043b9bec16d1620b4475fc016f90df1d5b5 +size 9653960