Files
2026-07-13 12:40:42 +08:00

295 lines
9.9 KiB
Python

# Copyright (c) 2021 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.
import inspect
import os
import pickle
import shlex
import shutil
import sys
import unittest
from collections import OrderedDict
import numpy as np
import paddle
from paddle.distributed.fleet.launch_utils import run_with_coverage
from paddle.distributed.passes.pass_base import PassBase, PassManager
def prepare_python_path_and_return_module(path):
dirname, filename = os.path.split(path)
py_suffix = ".py"
assert filename.endswith(py_suffix), filename
env_name = 'PYTHONPATH'
python_path = os.environ.get(env_name, '')
if python_path:
paths = [p for p in python_path.split(":") if p]
if dirname not in paths:
paths.append(dirname)
python_path = ":".join(paths)
else:
python_path = dirname
os.environ[env_name] = python_path
print('GLOG_v=', os.environ.get('GLOG_v', None), flush=1)
return filename[: -len(py_suffix)]
def remove_path_if_exists(path):
if not os.path.exists(path):
return
if os.path.isfile(path):
os.remove(path)
else:
shutil.rmtree(path)
# NOTE: only support GPU now
class DistPassTestBase(unittest.TestCase):
def setUp(self):
paddle.enable_static()
if paddle.is_compiled_with_cuda():
paddle.set_flags({'FLAGS_cudnn_deterministic': 1})
seed = int(os.environ.get('SEED', -1))
if seed <= 0:
seed = np.random.randint(low=1, high=1000000, size=[1])[0]
os.environ['SEED'] = str(seed)
self.seed = seed
paddle.seed(self.seed)
self.rtol = 1e-5
self.atol = 1e-8
self.equal_nan = False
self.init()
def init(self):
pass
def get_model(self, place, **kwargs):
raise NotImplementedError
def apply_passes(self, main_prog, startup_prog):
raise NotImplementedError
def check_main(self, model=None, gpus=None, **kwargs):
pass_rets = self._distributed_launch(
model=model, apply_pass=True, gpus=gpus, **kwargs
)
no_pass_rets = self._distributed_launch(
model=model, apply_pass=False, gpus=gpus, **kwargs
)
self.check_results(no_pass_rets, pass_rets)
def check_results(self, no_pass_rets, pass_rets):
self.assertEqual(len(no_pass_rets), len(pass_rets))
for no_pass_ret, pass_ret in zip(no_pass_rets, pass_rets):
self.assertEqual(len(no_pass_ret), len(pass_ret))
for i, (out_var_no_pass, out_var_pass) in enumerate(
zip(no_pass_ret, pass_ret)
):
if out_var_no_pass is None:
self.assertIsNone(out_var_pass)
else:
self.assertEqual(len(out_var_pass), len(out_var_no_pass))
for i in range(0, len(out_var_pass)):
np.testing.assert_allclose(
out_var_no_pass[i],
out_var_pass[i],
rtol=self.rtol,
atol=self.atol,
equal_nan=self.equal_nan,
)
@classmethod
def _to_var_names(cls, names_or_vars):
if not isinstance(names_or_vars, (list, tuple)):
names_or_vars = [names_or_vars]
ret_var_names = []
for name_or_var in names_or_vars:
if isinstance(name_or_var, str):
ret_var_names.append(name_or_var)
else:
ret_var_names.append(name_or_var.name)
return ret_var_names
def _run_gpu_main(self, model, apply_pass, dump_file, **kwargs):
gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
place = paddle.CUDAPlace(gpu_id)
scope = paddle.static.Scope()
if model is None:
model = self.get_model
with (
paddle.static.program_guard(
paddle.static.Program(), paddle.static.Program()
),
paddle.static.scope_guard(scope),
paddle.base.unique_name.guard(),
):
main_prog, startup_prog, inputs, outputs, reader = model(
place, **kwargs
)
inputs = self._to_var_names(inputs)
outputs = self._to_var_names(outputs)
if apply_pass:
self.apply_passes(main_prog, startup_prog)
all_fetch_values = []
exe = paddle.static.Executor(place)
with paddle.static.scope_guard(scope):
exe.run(startup_prog)
for batch_id, input_data in enumerate(reader()):
assert len(input_data) == len(inputs), (
f"{len(input_data)} vs {len(inputs)}"
)
feed = dict(zip(inputs, input_data))
fetch_values = exe.run(main_prog, feed=feed, fetch_list=outputs)
if paddle.distributed.get_rank() == 0:
output_dict = OrderedDict(zip(outputs, fetch_values))
print(f'batch {batch_id}, outputs {output_dict}')
all_fetch_values.append(fetch_values)
with open(dump_file, "wb") as f:
pickle.dump(all_fetch_values, f)
@classmethod
def _get_default_gpu_lists(cls):
visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
if visible_devices is None:
visible_devices = os.getenv("FLAGS_selected_gpus")
if visible_devices is None:
num_gpus = paddle.device.cuda.device_count()
return list(range(num_gpus))
else:
return [
int(s.strip()) for s in visible_devices.split(",") if s.strip()
]
def _distributed_launch(self, model, apply_pass, gpus=None, **kwargs):
if gpus is None:
gpus = self._get_default_gpu_lists()
num_gpus = len(gpus)
gpus = ','.join([str(gpu_id) for gpu_id in gpus])
pid = os.getpid()
if apply_pass:
output_dir = f"test_with_pass_{pid}"
else:
output_dir = f"test_without_pass_{pid}"
remove_path_if_exists(output_dir)
os.makedirs(output_dir, mode=0o777)
input_dump_file = os.path.join(output_dir, 'inputs.bin')
model_dump_file = os.path.join(output_dir, 'model.bin')
if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
run_with_coverage(True)
coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
else:
coverage_args = []
file_dir = os.path.dirname(os.path.abspath(__file__))
try:
with open(input_dump_file, 'wb') as f:
pickle.dump(kwargs, f)
if model is not None:
with open(model_dump_file, 'wb') as f:
pickle.dump(model, f)
cmd = [
sys.executable,
"-u",
*coverage_args,
"-m",
"launch",
"--log_dir",
output_dir,
"--gpus",
gpus,
os.path.join(file_dir, "pass_run_main.py"),
"--file_path",
inspect.getfile(type(self)),
"--class_name",
type(self).__name__,
"--input_file",
input_dump_file,
"--output_dir",
output_dir,
]
if apply_pass:
cmd += ["--apply_pass"]
if model is not None:
cmd += ["--model_file", model_dump_file]
cmd = [shlex.quote(c) for c in cmd]
prepare_python_path_and_return_module(__file__)
exitcode = os.system(' '.join(cmd))
self.assertEqual(
exitcode,
0,
f"Pass test failed with apply_pass = {apply_pass}, please view log in {output_dir}",
)
results = []
for i in range(num_gpus):
dump_file = f'{output_dir}/{i}.bin'
self.assertTrue(
os.path.exists(dump_file),
f"Pass test failed with apply_pass = {apply_pass}, please view log in {output_dir}",
)
with open(dump_file, "rb") as f:
results.append(pickle.load(f))
return results
finally:
if int(os.environ.get("DEBUG", 0)) == 0:
remove_path_if_exists(output_dir)
class PassConflictChecker(DistPassTestBase):
def setUp(self):
os.environ['DEBUG'] = '0'
super().setUp()
def pass_config(self):
raise NotImplementedError
def apply_passes(self, main_prog, startup_prog):
passes = self.pass_config()
if not isinstance(passes, (list, tuple)):
passes = [passes]
for p in passes:
self.assertTrue(isinstance(p, PassBase))
auto_pass_manager = PassManager(passes, auto_solve_conflict=True)
new_passes = auto_pass_manager.passes
self.assertEqual(
len(passes),
len(new_passes),
f"After solving conflicts, the left passes are: {auto_pass_manager.names}",
)
for i, (p1, p2) in enumerate(zip(passes, new_passes)):
self.assertEqual(
id(p1),
id(p2),
f"After solving conflicts, the {i}-th pass is different: {p1.name} vs {p2.name}",
)
auto_pass_manager.apply([main_prog], [startup_prog])