wehub-resource-sync
142 lines
5.2 KiB
Python
142 lines
5.2 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import deepspeed.compile.inductor as inductor
|
|
|
|
|
|
def _compiler(name):
|
|
|
|
def compiler(*args, **kwargs):
|
|
return name, args, kwargs
|
|
|
|
compiler.__name__ = name
|
|
return compiler
|
|
|
|
|
|
def _install_patch_spies(monkeypatch):
|
|
compiler_calls = []
|
|
partition_calls = []
|
|
|
|
def fake_patch_compiler(original_compiler, dc_compiler, z3_partition, graph_id, graph_param_manager, bwd):
|
|
wrapped = {
|
|
"original_compiler": original_compiler,
|
|
"dc_compiler": dc_compiler,
|
|
"z3_partition": z3_partition,
|
|
"graph_id": graph_id,
|
|
"graph_param_manager": graph_param_manager,
|
|
"bwd": bwd,
|
|
}
|
|
compiler_calls.append(wrapped)
|
|
return wrapped
|
|
|
|
def fake_wrap_partition_fn(z3_partition, partition_fn, real_inputs, param_indices, frame_id, frames_partitioned):
|
|
wrapped = {
|
|
"z3_partition": z3_partition,
|
|
"partition_fn": partition_fn,
|
|
"real_inputs": real_inputs,
|
|
"param_indices": param_indices,
|
|
"frame_id": frame_id,
|
|
"frames_partitioned": frames_partitioned,
|
|
}
|
|
partition_calls.append(wrapped)
|
|
return wrapped
|
|
|
|
monkeypatch.setattr(inductor, "patch_compiler", fake_patch_compiler)
|
|
monkeypatch.setattr(inductor, "wrap_partition_fn", fake_wrap_partition_fn)
|
|
return compiler_calls, partition_calls
|
|
|
|
|
|
def _patch_kwargs(kwargs, monkeypatch):
|
|
compiler_calls, partition_calls = _install_patch_spies(monkeypatch)
|
|
make_fw_graph = object()
|
|
make_bw_graph = object()
|
|
real_inputs = object()
|
|
param_indices = object()
|
|
param_manager = object()
|
|
frames_partitioned = set()
|
|
|
|
applied = inductor._patch_deepcompile_aot_kwargs(kwargs,
|
|
graph_id=7,
|
|
z3_partition=True,
|
|
make_fw_graph=make_fw_graph,
|
|
make_bw_graph=make_bw_graph,
|
|
real_inputs=real_inputs,
|
|
param_indices=param_indices,
|
|
param_manager=param_manager,
|
|
frame_id=11,
|
|
frames_partitioned=frames_partitioned)
|
|
|
|
return {
|
|
"applied": applied,
|
|
"compiler_calls": compiler_calls,
|
|
"partition_calls": partition_calls,
|
|
"make_fw_graph": make_fw_graph,
|
|
"make_bw_graph": make_bw_graph,
|
|
"real_inputs": real_inputs,
|
|
"param_indices": param_indices,
|
|
"param_manager": param_manager,
|
|
"frames_partitioned": frames_partitioned,
|
|
}
|
|
|
|
|
|
def test_legacy_inductor_shape_wraps_explicit_bw_compiler(monkeypatch):
|
|
fw_compiler = _compiler("fw")
|
|
bw_compiler = _compiler("bw")
|
|
inference_compiler = _compiler("inference")
|
|
partition_fn = _compiler("partition")
|
|
kwargs = {
|
|
"fw_compiler": fw_compiler,
|
|
"bw_compiler": bw_compiler,
|
|
"inference_compiler": inference_compiler,
|
|
"partition_fn": partition_fn,
|
|
}
|
|
|
|
result = _patch_kwargs(kwargs, monkeypatch)
|
|
|
|
assert result["applied"] is True
|
|
assert len(result["compiler_calls"]) == 2
|
|
assert result["compiler_calls"][0]["original_compiler"] is fw_compiler
|
|
assert result["compiler_calls"][0]["dc_compiler"] is result["make_fw_graph"]
|
|
assert result["compiler_calls"][0]["bwd"] is False
|
|
assert result["compiler_calls"][1]["original_compiler"] is bw_compiler
|
|
assert result["compiler_calls"][1]["dc_compiler"] is result["make_bw_graph"]
|
|
assert result["compiler_calls"][1]["bwd"] is True
|
|
assert kwargs["fw_compiler"] is result["compiler_calls"][0]
|
|
assert kwargs["bw_compiler"] is result["compiler_calls"][1]
|
|
assert kwargs["inference_compiler"] is result["compiler_calls"][0]
|
|
assert kwargs["partition_fn"] is result["partition_calls"][0]
|
|
assert result["partition_calls"][0]["partition_fn"] is partition_fn
|
|
|
|
|
|
def test_missing_bw_compiler_uses_original_fw_compiler_for_backward(monkeypatch):
|
|
fw_compiler = _compiler("fw")
|
|
partition_fn = _compiler("partition")
|
|
kwargs = {
|
|
"fw_compiler": fw_compiler,
|
|
"partition_fn": partition_fn,
|
|
}
|
|
|
|
result = _patch_kwargs(kwargs, monkeypatch)
|
|
|
|
assert result["applied"] is True
|
|
assert result["compiler_calls"][0]["original_compiler"] is fw_compiler
|
|
assert result["compiler_calls"][0]["bwd"] is False
|
|
assert result["compiler_calls"][1]["original_compiler"] is fw_compiler
|
|
assert result["compiler_calls"][1]["dc_compiler"] is result["make_bw_graph"]
|
|
assert result["compiler_calls"][1]["bwd"] is True
|
|
assert kwargs["bw_compiler"] is result["compiler_calls"][1]
|
|
|
|
|
|
def test_torchxla_openxla_shape_passes_through_unchanged(monkeypatch):
|
|
kwargs = {"fw_compiler": _compiler("openxla_eval_boxed")}
|
|
original_kwargs = dict(kwargs)
|
|
|
|
result = _patch_kwargs(kwargs, monkeypatch)
|
|
|
|
assert result["applied"] is False
|
|
assert result["compiler_calls"] == []
|
|
assert result["partition_calls"] == []
|
|
assert kwargs == original_kwargs
|