# Copyright (c) 2023 PaddlePaddle Authors. 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 __future__ import annotations import os import pathlib import pickle import subprocess import sys import tempfile import uuid from collections import defaultdict from typing import cast import numpy as np from prim_op_test import OpTestUtils, _as_list, convert_uint16_to_float, flatten from utils import dygraph_guard import paddle import paddle.distributed as dist IMPORT_PACKAGE_TEMPLATE = """ import pathlib import pickle import sys """ IMPORT_FORWARD_TEST_CLASS_TEMPLATE = """ sys.path.append( str(pathlib.Path(__file__).resolve().parents[0] / 'test/legacy_test') ) from auto_parallel_op_test import AutoParallelForwardChecker, convert_input_dims_map_to_placements """ IMPORT_GRAD_TEST_CLASS_TEMPLATE = """ sys.path.append( str(pathlib.Path(__file__).resolve().parents[0] / 'test/legacy_test') ) from auto_parallel_op_test import AutoParallelGradChecker, convert_input_dims_map_to_placements """ LOAD_TEST_INFO_TEMPLATE = """ def load_test_info(test_info_path): with open(test_info_path, "rb") as f: test_info = pickle.load(f) return test_info """ FORWARD_TEST_FUNCTION_TEMPLATE = """ def run_forward_check(test_info): auto_parallel_forward_checker = AutoParallelForwardChecker( test_info["op_type"], python_api, test_info["dtype"], convert_input_dims_map_to_placements(test_info["dims_map"], test_info["inputs"], 1), test_info["inputs"], test_info["attrs"], test_info["outputs"], test_info["place"], test_info["eager_auto_parallel_threshold"], test_info["python_out_sig"], ) auto_parallel_forward_checker.check() """ GRAD_TEST_FUNCTION_TEMPLATE = """ def run_grad_check(test_info): auto_parallel_forward_checker = AutoParallelGradChecker( test_info["op_type"], python_api, test_info["dtype"], convert_input_dims_map_to_placements(test_info["dims_map"], test_info["inputs"], 1), test_info["inputs"], test_info["attrs"], test_info["outputs"], test_info["place"], test_info["inputs_to_check"], test_info["output_names"], test_info["no_grad_set"], test_info["user_defined_grad_outputs"], test_info["eager_auto_parallel_threshold"], test_info["python_out_sig"], ) auto_parallel_forward_checker.check() """ LOAD_PYTHON_API_TEMPLATE = """ from {module} import {function} python_api = {function} """ TEST_BODY_TEMPLATE = """ if __name__ == "__main__": test_info = load_test_info(r'{test_info_path}') {load_python_api} {run_test} """ def is_ban_auto_parallel_test(place): if ( isinstance(place, paddle.base.libpaddle.CUDAPlace) and paddle.device.cuda.device_count() < 2 or not paddle.is_compiled_with_distribute() or ( os.environ.get("WITH_COVERAGE") == "ON" and os.environ.get("FLAGS_COVERAGE_RUN_AUTO_PARALLEL_IN_OP_TEST") != "1" ) ): return True else: return False def gen_import_packages(check_grad): import_code = '' import_code += IMPORT_PACKAGE_TEMPLATE import_code += ( IMPORT_FORWARD_TEST_CLASS_TEMPLATE if not check_grad else IMPORT_GRAD_TEST_CLASS_TEMPLATE ) return import_code def gen_auto_parallel_test_file( check_grad, test_info_path, test_file_path, python_api_info ): test_code = '' test_code += gen_import_packages(check_grad) test_code += LOAD_TEST_INFO_TEMPLATE.format(test_info_path=test_info_path) test_code += ( GRAD_TEST_FUNCTION_TEMPLATE if check_grad else FORWARD_TEST_FUNCTION_TEMPLATE ) run_test_str = ( "run_grad_check(test_info)" if check_grad else "run_forward_check(test_info)" ) load_python_api_str = LOAD_PYTHON_API_TEMPLATE.format( module=python_api_info["api_module"], function=python_api_info["api_name"], ) test_code += TEST_BODY_TEMPLATE.format( test_info_path=test_info_path, load_python_api=load_python_api_str, run_test=run_test_str, ) with open(test_file_path, "w") as f: f.write(test_code) def get_test_info_and_generated_test_path( test_class_name, op_type, backward=False ): suffixes = str(uuid.uuid4()) current_path = pathlib.Path(__file__).resolve().parents[0] forward_or_backward = "forward" if not backward else "backward" test_info_path = ( current_path / f"{test_class_name}_{op_type}_{forward_or_backward}_info_{suffixes}.pkl" ) generated_test_path = ( current_path / f"{test_class_name}_{op_type}_{forward_or_backward}_test_{suffixes}.py" ) return str(test_info_path), str(generated_test_path) def check_auto_parallel_info(op_test): assert hasattr(op_test, 'python_api'), ( "If you want to check auto parallel, please set python_api in setUp function." ) assert hasattr(op_test, 'placements'), ( "If you want to check auto parallel, please set placements in setUp function." ) def dump_test_info( op_test, place, test_info_path, backward=False, backward_extra_test_info=None, ): check_auto_parallel_info(op_test) test_info = {} with open(test_info_path, "wb") as f: test_info["op_type"] = op_test.op_type test_info["dtype"] = op_test.dtype test_info["dims_map"] = convert_input_placements_to_dims_map( op_test.placements, op_test.inputs ) test_info["inputs"] = op_test.inputs test_info["attrs"] = op_test.attrs if hasattr(op_test, "attrs") else {} test_info["outputs"] = op_test.outputs if isinstance(place, paddle.base.libpaddle.CPUPlace): test_info["place"] = "cpu" if isinstance(place, paddle.base.libpaddle.CUDAPlace): test_info["place"] = "gpu" eager_auto_parallel_threshold = { "atol": ( op_test.eager_auto_parallel_atol if hasattr(op_test, "eager_auto_parallel_atol") else None ), "rtol": ( op_test.eager_auto_parallel_atol if hasattr(op_test, "eager_auto_parallel_atol") else None ), } test_info["eager_auto_parallel_threshold"] = ( eager_auto_parallel_threshold ) test_info["python_out_sig"] = ( op_test.python_out_sig if hasattr(op_test, "python_out_sig") else None ) if backward: test_info["inputs_to_check"] = backward_extra_test_info[ "inputs_to_check" ] test_info["output_names"] = backward_extra_test_info["output_names"] test_info["no_grad_set"] = backward_extra_test_info["no_grad_set"] test_info["user_defined_grad_outputs"] = backward_extra_test_info[ "user_defined_grad_outputs" ] try: pickle.dump(test_info, f) except Exception as e: raise Exception( "Dump test info failed, please check your test info." ) def get_subprocess_runtime_envs(place): runtime_envs = os.environ if ( "CUDA_VISIBLE_DEVICES" not in runtime_envs or len(runtime_envs["CUDA_VISIBLE_DEVICES"].split(",")) < 2 ): runtime_envs.update({"CUDA_VISIBLE_DEVICES": "0,1"}) if isinstance(place, paddle.base.libpaddle.CPUPlace): runtime_envs.update({"backend": "cpu"}) if isinstance(place, paddle.base.libpaddle.CUDAPlace): runtime_envs.update({"backend": "gpu"}) return runtime_envs def get_subprocess_command(devices, test_file_path, log_dir=None): if log_dir: if os.path.isabs(log_dir): abs_log_dir = log_dir else: abs_log_dir = os.path.abspath(log_dir) else: abs_log_dir = tempfile.TemporaryDirectory().name start_command = f"{sys.executable} -m paddle.distributed.launch --devices {devices} --log_dir {abs_log_dir} {test_file_path}" return start_command def run_subprocess(start_command, env, timeout): start_command_list = start_command.strip().split() try: _launcher = subprocess.run( start_command_list, env=env, timeout=timeout, check=True, ) except subprocess.TimeoutExpired as err: raise TimeoutError( f"Timeout while running command {err.cmd}, try to set a longer period, {err.timeout} is not enough." ) except subprocess.CalledProcessError as err: raise RuntimeError( f"Error occurs when running this test case. The return code of command {err.cmd} is {err.returncode}" ) def convert_input_placements_to_dims_map(placements: dict, inputs: dict): all_dims_map = {} for name, item in inputs.items(): if name not in placements: continue # such as inputs = {"X": [("x0", x0), ("x1", x1), ("x2", x2)]} # placements = {"X": [("x0", [Shard(0)]), ("x1", [Shard(0)]), ("x2", [Shard(0)])]} if isinstance(item, list): all_dims_map[name] = [] for i in range(len(item)): dims_map = placements_to_dims_map( placements[name][i][1], inputs[name][i][1].ndim ) all_dims_map[name].append((item[i][0], dims_map)) # inputs like this : inputs = {'X': x} # placements = {"X": [Shard(0)]} else: dims_map = placements_to_dims_map( placements[name], inputs[name].ndim ) all_dims_map[name] = dims_map return all_dims_map def convert_input_dims_map_to_placements( dims_map: dict, inputs: dict, mesh_ndim: int ): placements_map = {} for name, item in inputs.items(): if name not in dims_map: continue # such as inputs = {"X": [("x0", x0), ("x1", x1), ("x2", x2)]} # dims_map = {"X": [("x0", [-1, 0]), ("x1", [-1, 0]), ("x2", [-1, 0]} if isinstance(item, list): placements_map[name] = [] for i in range(len(item)): placements = dims_map_to_placements( dims_map[name][i][1], mesh_ndim ) placements_map[name].append((item[i][0], placements)) # inputs like this : inputs = {'X': x} # placements = {"X": [Shard(0)]} else: placements = dims_map_to_placements(dims_map[name], mesh_ndim) placements_map[name] = placements return placements_map # TODO: This method has been implemented in # paddle/phi/core/distributed/auto_parallel/placement_types.h, bind it # python and it's logic. def placements_to_dims_map(placements: list, tensor_ndim: int) -> tuple[int]: r = [-1] * tensor_ndim for i, placement in enumerate(placements): if placement.is_shard(): shard_dim = cast("dist.Shard", placement).get_dim() if r[shard_dim] > -1: raise ValueError( f"Tensor dim {shard_dim} is already sharded on mesh dim {r[shard_dim]}," " DTensor operator implementation does not support things like hybrid" " sharding strategies yet (i.e. [Shard(0), Shard(0)])" ) r[shard_dim] = i return r # TODO: Add this method to # paddle/phi/core/distributed/auto_parallel/placement_types.h, and bind it to # python def dims_map_to_placements( dim_map: tuple[int], mesh_ndim: int, sums: tuple[int] = () ) -> tuple[dist.Placement]: """ Construct a placements from dim_map list and pending sum. Args: dim_map (tuple[int]): a list of integer that represents sharding on each tensor dimension, see `dim_map` property doc for details mesh_ndim (int): the ndim of Process mesh. sums (Tuple[int]): a list of integer that represents the dist tensor have pending sum on which device mesh dimension. Return: a placement sequence. """ # by default replicate on device mesh dims placements: list[dist.Placement] = [ dist.Replicate() for _ in range(mesh_ndim) ] # find all mesh dims that need pending reductions for s in sums: placements[s] = dist.Partial() for i, m in enumerate(dim_map): if m >= 0: placement = placements[m] if placement.is_shard(): placement = cast("dist.Shard", placement) raise RuntimeError( f"DeviceMesh dimension can't be mapped to two dimension of the same tensor: {i} and {placement.dim}" ) elif placement.is_partial(): raise RuntimeError( f"DeviceMesh dimension {m} cannot be both shard and partial!" ) placements[m] = dist.Shard(i) return tuple(placements) TOLERANCE = { np.dtype('float64'): {"rtol": 1e-15, "atol": 0}, np.dtype('float32'): {"rtol": 1e-6, "atol": 0}, np.dtype('float16'): {"rtol": 1e-3, "atol": 0}, np.dtype('uint16'): {"rtol": 1e-2, "atol": 0}, np.dtype('int32'): {"rtol": 0, "atol": 0}, } class AutoParallelForwardChecker: def __init__( self, op_type, python_api, dtype, placements_map, inputs, attrs, outputs, place, eager_auto_parallel_threshold, python_out_sig=None, ): self.checker_name = "AutoParallelForwardChecker" self.init_checker( op_type, python_api, dtype, placements_map, inputs, attrs, outputs, place, eager_auto_parallel_threshold, python_out_sig, ) def init_checker( self, op_type, python_api, dtype, placements_map, inputs, attrs, outputs, place, eager_auto_parallel_threshold, python_out_sig=None, ): self.op_type = op_type self.public_python_api = python_api self.dtype = np.dtype(dtype) self.placements_map = placements_map self.inputs = inputs self.attrs = attrs self.outputs = outputs self.place = place if self.place == "cpu": paddle.device.set_device("cpu") if self.place == "gpu": paddle.device.set_device("gpu:" + str(dist.get_rank())) self.python_out_sig = python_out_sig self.attrs = attrs self.outputs = outputs self.init_checker_threshold( eager_auto_parallel_threshold["atol"], eager_auto_parallel_threshold["rtol"], ) self.kernel_sig = self.get_kernel_sig() self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"]) def init_checker_threshold(self, atol=None, rtol=None): self.atol = atol if atol else TOLERANCE[self.dtype]["atol"] self.rtol = rtol if rtol else TOLERANCE[self.dtype]["rtol"] def check(self): self.eager_forward_desire = self.get_eager_desire() self.check_eager_auto_parallel() def check_eager_auto_parallel(self): with dygraph_guard(): actual_ret = self.get_eager_desire(dist_mode=True) # check eager auto parallel forward if len(actual_ret) != len(self.eager_forward_desire): msg = ( f"The eager auto parallel out tensor nums is different with eager out tensor nums on {self.place}." f'eager auto parallel out tensor nums = {len(actual_ret)}, eager out tensor nums = {len(self.eager_forward_desire)}. \n' ) raise RuntimeError(msg) for i in range(len(actual_ret)): np.testing.assert_allclose( actual_ret[i], self.eager_forward_desire[i], rtol=self.atol, atol=self.rtol, err_msg=( 'Check eager auto parallel failed. Mismatch between eager auto parallel outputs ' f'and eager outputs on {self.place!s}, the eager forward output tensor\'s index is : {i} \n' f'eager auto parallel output tensor:\n{actual_ret[i]}\n eager output tensor:\n{self.eager_forward_desire[i]}\n' ), ) def get_kernel_sig(self): with dygraph_guard(): ( eager_tensor_inputs, attrs_outputs, _, ) = self.get_eager_input_attr_and_inputdict(stop_gradient=True) eager_tensor_outputs = self.get_eager_empty_output( stop_gradient=True ) kernel_sig = OpTestUtils._get_kernel_signature( self.op_type, eager_tensor_inputs, eager_tensor_outputs, attrs_outputs, ) return kernel_sig def get_eager_desire(self, dist_mode=False): with dygraph_guard(): if dist_mode: ( eager_tensor_inputs, attrs_outputs, _, ) = self.get_eager_input_attr_and_inputdict( stop_gradient=True, dist_mode=True ) else: ( eager_tensor_inputs, attrs_outputs, _, ) = self.get_eager_input_attr_and_inputdict( stop_gradient=True, dist_mode=False ) args = OpTestUtils.prepare_python_api_arguments( self.public_python_api, eager_tensor_inputs, attrs_outputs, self.kernel_sig, target_dtype=paddle.core.VarDesc.VarType, ) inputs_sig, _, _ = self.kernel_sig args = OpTestUtils.assumption_assert_and_transform( args, len(inputs_sig) ) ret = flatten(_as_list(self.public_python_api(*args))) ret = paddle.utils.map_structure(lambda x: x.numpy(), ret) if OpTestUtils.is_bfloat16_type(self.dtype): ret = paddle.utils.map_structure( lambda x: convert_uint16_to_float(x), ret ) return ret def get_eager_input_attr_and_inputdict( self, stop_gradient, dist_mode=False ): attrs_outputs = {} for attrs_name in self.attrs: if self.attrs[attrs_name] is not None: attrs_outputs[attrs_name] = self.attrs[attrs_name] input_dict = {} eager_inputs = defaultdict(list) for name, item in self.inputs.items(): # such as inputs = {"X": [("x0", x0), ("x1", x1), ("x2", x2)]} # placements = {"X": [("x0", [Shard(0)]), ("x1", [Shard(0)]), ("x2", [Shard(0)])]} if isinstance(item, list): for i in range(len(item)): dtype = ( "bfloat16" if OpTestUtils.is_bfloat16_type(item[i][1].dtype) else item[i][1].dtype ) x = paddle.to_tensor( data=item[i][1], stop_gradient=stop_gradient, dtype=dtype, ) if not dist_mode or name not in self.placements_map: eager_inputs[name].append(x) input_dict.update({str(item[i][0]): x}) else: dist_x = dist.shard_tensor( x, self._mesh, self.placements_map[name][i][1] ) dist_x.stop_gradient = stop_gradient eager_inputs[name].append(dist_x) input_dict.update({str(item[i][0]): dist_x}) # inputs like this : inputs = {'X': x} # placements = {"X": [Shard(0)]} else: dtype = ( "bfloat16" if OpTestUtils.is_bfloat16_type(item.dtype) else item.dtype ) x = paddle.to_tensor( data=item, stop_gradient=stop_gradient, dtype=dtype, ) if not dist_mode or name not in self.placements_map: eager_inputs[name].append(x) input_dict.update({name: x}) else: dist_x = dist.shard_tensor( x, self._mesh, self.placements_map[name] ) dist_x.stop_gradient = stop_gradient eager_inputs[name].append(dist_x) input_dict.update({name: dist_x}) return eager_inputs, attrs_outputs, input_dict def get_eager_empty_output(self, stop_gradient): eager_outputs = defaultdict(list) for name, item in self.outputs.items(): if isinstance(item, list): for tup in item: dtype = ( "bfloat16" if OpTestUtils.is_bfloat16_type(tup[1].dtype) else tup[1].dtype ) x = paddle.to_tensor( data=[], stop_gradient=stop_gradient, dtype=dtype, ) eager_outputs[name].append(x) else: dtype = ( "bfloat16" if OpTestUtils.is_bfloat16_type(item.dtype) else item.dtype ) x = paddle.to_tensor( data=[], stop_gradient=stop_gradient, dtype=dtype, ) eager_outputs[name].append(x) return eager_outputs class AutoParallelGradChecker(AutoParallelForwardChecker): def __init__( self, op_type, python_api, dtype, placements_map, inputs, attrs, outputs, place, inputs_to_check, output_names, no_grad_set, grad_outputs, eager_auto_parallel_threshold, python_out_sig=None, ): super().__init__( op_type, python_api, dtype, placements_map, inputs, attrs, outputs, place, eager_auto_parallel_threshold, python_out_sig, ) self.checker_name = "AutoParallelGradChecker" self.inputs_to_check = inputs_to_check self.output_names = output_names self.no_grad_set = no_grad_set self.grad_outputs = grad_outputs def check(self): ( self.eager_forward_desire, self.eager_grad_desire, ) = self.get_eager_desire() self.check_eager_auto_parallel() def check_eager_auto_parallel(self): with dygraph_guard(): actual_forward_res, actual_grad_res = self.get_eager_desire( dist_mode=True ) # check eager auto parallel forward if len(actual_forward_res) != len(self.eager_forward_desire): msg = ( f"The eager auto parallel out tensor nums is different with eager out tensor nums on {self.place}." f'eager auto parallel out tensor nums = {len(actual_forward_res)}, eager out tensor nums = {len(self.eager_forward_desire)}. \n' ) raise RuntimeError(msg) for i in range(len(actual_forward_res)): np.testing.assert_allclose( actual_forward_res[i], self.eager_forward_desire[i], rtol=self.atol, atol=self.rtol, err_msg=( 'Check eager auto parallel failed. Mismatch between eager auto parallel outputs ' f'and eager outputs on {self.place!s}, the eager forward output tensor\'s index is : {i} \n' f'eager auto parallel output tensor:\n{actual_forward_res[i]}\n eager output tensor:\n{self.eager_forward_desire[i]}\n' ), ) # check eager auto parallel grad if len(actual_grad_res) != len(self.eager_grad_desire): msg = ( f"The eager auto parallel grad out tensor nums is different with eager grad out tensor nums on {self.place}." f'eager auto parallel grad out tensor nums = {len(actual_grad_res)}, eager grad out tensor nums = {len(self.eager_grad_desire)}. \n' ) raise RuntimeError(msg) for i in range(len(actual_grad_res)): np.testing.assert_allclose( actual_grad_res[i], self.eager_grad_desire[i], rtol=self.atol, atol=self.rtol, err_msg=( 'Check eager auto parallel backward failed. Mismatch between eager auto parallel grad outputs ' f'and eager grad outputs on {self.place!s}, the eager grad output tensor\'s index is : {i} \n' f'eager auto parallel grad output tensor:\n{actual_grad_res[i]}\n eager grad output tensor:\n{self.eager_grad_desire[i]}\n' ), ) def gen_eager_grad_outputs(self): if self.grad_outputs is None: return None eager_vs = [] for np_v in self.grad_outputs: eager_vs.append( paddle.to_tensor( data=np_v, place=self.place, dtype=( "bfloat16" if OpTestUtils.is_bfloat16_type(np_v.dtype) else np_v.dtype ), ) ) return eager_vs def get_output_dict(self, np_outputs, api_outputs, outputs_sig): assert len(api_outputs) <= len(outputs_sig), ( f"forward api outputs length must be the less than or equal to KernelSignature outputs,but receive {len(api_outputs)} and {len(outputs_sig)}" ) output_dict = {} for i in range(len(api_outputs)): output_name = outputs_sig[i] if output_name in np_outputs and isinstance( np_outputs[output_name], list ): for j, tup in enumerate(np_outputs[output_name]): output_dict.update({tup[0]: api_outputs[i][j]}) else: output_dict.update({output_name: api_outputs[i]}) return output_dict def gen_no_grad_set(self, var_dict): if self.no_grad_set is None: return None no_grad_set = set() for name in self.no_grad_set: if name in var_dict: no_grad_set.add(var_dict[name]) return no_grad_set def get_eager_desire(self, dist_mode=False): with dygraph_guard(): if dist_mode: ( eager_tensor_inputs, attrs_outputs, inputs_dict, ) = self.get_eager_input_attr_and_inputdict( stop_gradient=False, dist_mode=True ) else: ( eager_tensor_inputs, attrs_outputs, inputs_dict, ) = self.get_eager_input_attr_and_inputdict( stop_gradient=False, dist_mode=False ) args = OpTestUtils.prepare_python_api_arguments( self.public_python_api, eager_tensor_inputs, attrs_outputs, self.kernel_sig, target_dtype=paddle.core.VarDesc.VarType, ) inputs_sig, _, outputs_sig = self.kernel_sig if self.python_out_sig is not None: outputs_sig = self.python_out_sig args = OpTestUtils.assumption_assert_and_transform( args, len(inputs_sig) ) forward_res = _as_list(self.public_python_api(*args)) outputs_dict = self.get_output_dict( self.outputs, forward_res, outputs_sig ) ys = [] if isinstance(self.output_names, list): for output_name in self.output_names: ys.append(outputs_dict[output_name]) else: ys.append(outputs_dict[self.output_names]) xs = [] if isinstance(self.inputs_to_check, list): for input_name in self.inputs_to_check: xs.append(inputs_dict[input_name]) else: xs.append(inputs_dict[self.inputs_to_check]) vs = self.gen_eager_grad_outputs() no_grad_vars = self.gen_no_grad_set( var_dict=inputs_dict | outputs_dict ) grad_res = paddle.grad( ys, xs, vs, allow_unused=True, no_grad_vars=no_grad_vars ) forward_res = paddle.utils.map_structure( lambda x: x.numpy(), forward_res ) grad_res = paddle.utils.map_structure(lambda x: x.numpy(), grad_res) if OpTestUtils.is_bfloat16_type(self.dtype): forward_res = paddle.utils.map_structure( lambda x: convert_uint16_to_float(x), forward_res ) grad_res = paddle.utils.map_structure( lambda x: convert_uint16_to_float(x), grad_res ) return forward_res, grad_res