Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Semi Auto] Revise spmd rule static mode API #57269

Merged
merged 23 commits into from
Sep 19, 2023
Merged

[Semi Auto] Revise spmd rule static mode API #57269

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
ba1e4bf
adapt general spmd rule
chenwhql Aug 30, 2023
5917bb2
polish details
chenwhql Aug 31, 2023
e9a4455
add new rules
JZ-LIANG Aug 31, 2023
8a58468
bugfix
JZ-LIANG Aug 31, 2023
459d633
register
JZ-LIANG Aug 31, 2023
4f21444
bugfix for set_partial
JZ-LIANG Aug 31, 2023
5526841
Merge remote-tracking branch 'cwh/ap/add_general_spmd_rule_utils' int…
JZ-LIANG Sep 4, 2023
3c744a7
bugfix
JZ-LIANG Sep 5, 2023
69fce04
unitest
JZ-LIANG Sep 5, 2023
78d6d90
unitest
JZ-LIANG Sep 5, 2023
92aadb5
static done
JZ-LIANG Sep 6, 2023
c6571f9
adapt for dy graph
JZ-LIANG Sep 6, 2023
44605fb
clean
JZ-LIANG Sep 6, 2023
946c991
adapt for argument for tensor and vector of tensor
JZ-LIANG Sep 6, 2023
4a328d4
Merge remote-tracking branch 'upstream/develop' into semi-auto/defaul…
JZ-LIANG Sep 6, 2023
5e2a99b
pass as reference
JZ-LIANG Sep 7, 2023
c690246
bugfix
JZ-LIANG Sep 11, 2023
1c316a5
parse vector
JZ-LIANG Sep 11, 2023
1fde8d4
Merge remote-tracking branch 'upstream/develop' into semi-auto/spmd-r…
JZ-LIANG Sep 12, 2023
acb2821
revise spmd rule static api
JZ-LIANG Sep 13, 2023
11469ab
polish comment
JZ-LIANG Sep 13, 2023
d19df7c
bugfix in win
JZ-LIANG Sep 15, 2023
06cfbd4
bugfix for include static function
JZ-LIANG Sep 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
264 changes: 152 additions & 112 deletions paddle/fluid/pybind/auto_parallel_py.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

#include <Python.h>
#include <pybind11/operators.h>
#include <pybind11/stl.h>
#include <utility>
Expand All @@ -20,6 +21,8 @@
#include "paddle/fluid/framework/op_desc.h"
#include "paddle/fluid/framework/var_desc.h"
#include "paddle/fluid/pybind/auto_parallel_py.h"
#include "paddle/fluid/pybind/eager_utils.h"
#include "paddle/fluid/pybind/op_function_common.h"
#include "paddle/fluid/pybind/pybind_variant_caster.h"
#include "paddle/phi/core/device_context.h"
#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
Expand Down Expand Up @@ -47,6 +50,14 @@ namespace py = pybind11;
namespace paddle {
namespace pybind {

static bool PyCheckInteger(PyObject *obj) {
#if PY_VERSION_HEX < 0x03000000
return (PyLong_Check(obj) || PyInt_Check(obj)) && !PyBool_Check(obj);
#else
return PyLong_Check(obj) && !PyBool_Check(obj);
#endif
}

using paddle::distributed::auto_parallel::DistTensorSpec;
using paddle::distributed::auto_parallel::kDefault;
using paddle::distributed::auto_parallel::OperatorDistAttr;
Expand All @@ -66,6 +77,8 @@ using phi::distributed::auto_parallel::LinkCapability;
using phi::distributed::auto_parallel::Machine;

PyTypeObject *g_tensor_dist_attr_pytype = nullptr;
PyTypeObject *g_dist_tensor_spec_pytype = nullptr;
constexpr const char *infer_spmd_string = "infer_spmd";

static inline const ProcessMesh *get_tensor_process_mesh(
const TensorDistAttr &self) {
Expand Down Expand Up @@ -122,6 +135,11 @@ static inline void reset_operator_dist_attr(OperatorDistAttr *dist_attr) {
dist_attr->clear_annotated();
}

static std::pair<std::vector<TensorDistAttr>, std::vector<TensorDistAttr>>
infer_forward(const phi::distributed::SpmdRule &self, const py::args &args);
static std::pair<std::vector<TensorDistAttr>, std::vector<TensorDistAttr>>
infer_backward(const phi::distributed::SpmdRule &self, const py::args &args);

void BindAutoParallel(py::module *m) {
auto ReshardFunction =
py::class_<phi::distributed::ReshardFunction>(*m, "ReshardFunction")
Expand Down Expand Up @@ -368,118 +386,14 @@ void BindAutoParallel(py::module *m) {
// .def("infer_backward", &SPMDRuleBase::InferBackward) [revert in future]

py::class_<phi::distributed::SpmdRule>(*m, "SpmdRule")
.def("infer_forward",
[](const phi::distributed::SpmdRule &self,
const std::vector<DistTensorSpec> &input_specs,
const std::vector<phi::Attribute> &attrs) {
phi::distributed::InferSpmdContext ctx;
for (auto &spec : input_specs) {
ctx.EmplaceBackInput(phi::distributed::DistMetaTensor(
phi::make_ddim(spec.shape()), spec.dist_attr()));
}
for (auto &attr : attrs) {
ctx.EmplaceBackAttr(attr);
}
return self.InferForward(ctx);
})
.def("infer_forward", // for op that have vector argument
[](const phi::distributed::SpmdRule &self,
const std::vector<std::pair<int, int>> &input_ranges,
const std::vector<DistTensorSpec> &input_specs,
const std::vector<phi::Attribute> &attrs) {
/*
to distingish between single tensor argument and vector argument of
one tensor: start - end == 0: single tensor start - end == 1:
vector containing one tensor input_ranges: [(0, 0), (1, 3), (3, 4)]
+ input_specs: [t0, t1, t2, t3] --> t0, [t1, t2], [t3]
*/
phi::distributed::InferSpmdContext ctx;
paddle::small_vector<phi::distributed::DistMetaTensor,
phi::kInputSmallVectorSize>
ins;
for (auto &range : input_ranges) {
if (range.second - range.first == 0) {
auto &in = input_specs.at(range.first);
ctx.EmplaceBackInput(phi::distributed::DistMetaTensor(
phi::make_ddim(in.shape()), in.dist_attr()));
} else {
int start = range.first;
int end = range.second;
ins.reserve(end - start);
for (int i = start; i < end; ++i) {
auto &in = input_specs.at(i);
ins.emplace_back(phi::distributed::DistMetaTensor(
phi::make_ddim(in.shape()), in.dist_attr()));
}
ctx.EmplaceBackInputs(ins);
ins.clear();
}
}
for (auto &attr : attrs) {
ctx.EmplaceBackAttr(attr);
}
return self.InferForward(ctx);
})
.def("infer_backward",
[](const phi::distributed::SpmdRule &self,
const std::vector<DistTensorSpec> &input_specs,
const std::vector<DistTensorSpec> &output_specs,
const std::vector<phi::Attribute> &attrs) {
phi::distributed::InferSpmdContext ctx;
for (auto &spec : input_specs) {
ctx.EmplaceBackInput(phi::distributed::DistMetaTensor(
phi::make_ddim(spec.shape()), spec.dist_attr()));
}
for (auto &spec : output_specs) {
ctx.EmplaceBackInput(phi::distributed::DistMetaTensor(
phi::make_ddim(spec.shape()), spec.dist_attr()));
}
for (auto &attr : attrs) {
ctx.EmplaceBackAttr(attr);
}
return self.InferBackward(ctx);
})
.def("infer_backward", // for op that have vector argument
[](const phi::distributed::SpmdRule &self,
const std::vector<std::pair<int, int>> &input_ranges,
const std::vector<DistTensorSpec> &input_specs,
const std::vector<phi::Attribute> &attrs) {
/*
to distingish between single tensor argument and vector argument of
one tensor: start - end == 0: single tensor start - end == 1:
vector containing one tensor input_ranges: [(0, 0), (1, 3), (3, 4)]
+ input_specs: [t0, t1, t2, t3] --> t0, [t1, t2], [t3]
*/
phi::distributed::InferSpmdContext ctx;
paddle::small_vector<phi::distributed::DistMetaTensor,
phi::kInputSmallVectorSize>
ins;
for (auto &range : input_ranges) {
if (range.second - range.first == 0) {
auto &in = input_specs.at(range.first);
ctx.EmplaceBackInput(phi::distributed::DistMetaTensor(
phi::make_ddim(in.shape()), in.dist_attr()));
} else {
int start = range.first;
int end = range.second;
ins.reserve(end - start);
for (int i = start; i < end; ++i) {
auto &in = input_specs.at(i);
ins.emplace_back(phi::distributed::DistMetaTensor(
phi::make_ddim(in.shape()), in.dist_attr()));
}
ctx.EmplaceBackInputs(ins);
ins.clear();
}
}
for (auto &attr : attrs) {
ctx.EmplaceBackAttr(attr);
}
return self.InferBackward(ctx);
});

py::class_<DistTensorSpec>(*m, "DistTensorSpec")
.def(py::init<>())
.def("infer_forward", &infer_forward)
.def("infer_backward", &infer_backward);

py::class_<DistTensorSpec> py_dist_tensor_spec(
*m, "DistTensorSpec"); // TODO(ljz) remove and unify to DistTensor
g_dist_tensor_spec_pytype =
reinterpret_cast<PyTypeObject *>(py_dist_tensor_spec.ptr());
py_dist_tensor_spec.def(py::init<>())
.def(py::init<const DistTensorSpec &>())
.def(py::init<const std::vector<int64_t> &, const TensorDistAttr &>())
.def("dims_mapping", &DistTensorSpec::dims_mapping)
Expand Down Expand Up @@ -634,5 +548,131 @@ void BindAutoParallel(py::module *m) {
});
}

static void parse_tensors(PyObject *obj,
phi::distributed::InferSpmdContext *ctx,
const size_t arg_pos) {
Py_ssize_t len = PyList_Size(obj);
VLOG(6) << "args indx: [" << arg_pos << "] input vector of ["
<< static_cast<size_t>(len) << "] tensors.";
paddle::small_vector<phi::distributed::DistMetaTensor,
phi::kInputSmallVectorSize>
ins;
ins.reserve(static_cast<size_t>(len));
for (Py_ssize_t i = 0; i < len; i++) {
DistTensorSpec in = py::cast<DistTensorSpec>(PyList_GetItem(obj, i));
VLOG(6) << "Vector emplace_back DistTensorSpec: " << in.to_string();
ins.emplace_back(phi::distributed::DistMetaTensor(
phi::make_ddim(in.shape()), in.dist_attr()));
}
ctx->EmplaceBackInputs(ins);
}

static void parse_tensor(PyObject *obj,
phi::distributed::InferSpmdContext *ctx,
const size_t arg_pos) {
VLOG(6) << "args indx: [" << arg_pos << "] input one tensor.";
DistTensorSpec in = py::cast<DistTensorSpec>(obj);
VLOG(6) << "DistTensorSpec: " << in.to_string();
ctx->EmplaceBackInput(phi::distributed::DistMetaTensor(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can consider unify DistMetaTensor and DistTensorSpec in the future.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, that is on the list

phi::make_ddim(in.shape()), in.dist_attr()));
}

// TODO(ljz) support other types
static void parse_attrs(PyObject *obj,
PyObject *first_item,
phi::distributed::InferSpmdContext *ctx,
const size_t arg_pos) {
if (PyBool_Check(first_item)) {
auto attrs = CastPyArg2Booleans(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attrs);
} else if (PyCheckInteger(first_item)) {
auto attrs = CastPyArg2Ints(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attrs);
} else if (PyLong_Check(first_item)) {
auto attrs = CastPyArg2Longs(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attrs);
} else if (PyFloat_Check(first_item)) {
auto attrs = CastPyArg2Floats(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attrs);
} else {
PADDLE_THROW(platform::errors::InvalidArgument(
"%s(): argument (position %d) must be "
"list of int, float, bool or Tensor, but got %s",
infer_spmd_string,
arg_pos,
((PyTypeObject *)first_item->ob_type)->tp_name)); // NOLINT
}
}

// TODO(ljz) support other types
static void parse_attr(PyObject *obj,
phi::distributed::InferSpmdContext *ctx,
const size_t arg_pos) {
if (PyBool_Check(obj)) {
auto attr = CastPyArg2Boolean(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attr);
} else if (PyCheckInteger(obj)) {
auto attr = CastPyArg2Int(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attr);
} else if (PyLong_Check(obj)) {
auto attr = CastPyArg2Long(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attr);
} else if (PyFloat_Check(obj)) {
auto attr = CastPyArg2Float(obj, infer_spmd_string, arg_pos);
ctx->EmplaceBackAttr(attr);
} else { // TODO(ljz) support other types
PADDLE_THROW(platform::errors::InvalidArgument(
"%s(): argument (position %d) must be "
"int, float, bool or Tensor, but got %s",
infer_spmd_string,
arg_pos,
((PyTypeObject *)obj->ob_type)->tp_name)); // NOLINT
}
}

static void parse_single_pyobject(PyObject *obj,
phi::distributed::InferSpmdContext *ctx,
const size_t arg_pos) {
if (PyList_Check(obj)) { // list inputs, spmd not allow tuple inputs
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suggest we can support tuple and list both.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there is no tuple argument in Phi API yaml right now, we would support it once there is tuple input defined by Phi.

PyObject *first_item = PyList_GetItem(obj, 0);
if (PyObject_TypeCheck(first_item, g_dist_tensor_spec_pytype)) {
parse_tensors(obj, ctx, arg_pos);
} else {
parse_attrs(obj, first_item, ctx, arg_pos);
}
} else {
if (PyObject_TypeCheck(obj, g_dist_tensor_spec_pytype)) {
parse_tensor(obj, ctx, arg_pos);
} else {
parse_attr(obj, ctx, arg_pos);
}
}
}

static void prepare_ctx(phi::distributed::InferSpmdContext *ctx,
const py::args &args) {
VLOG(6) << "prepare_ctx ";
size_t inputs_size = args.size();
for (size_t i = 0; i < inputs_size; ++i) {
PyObject *obj = args[i].ptr();
parse_single_pyobject(obj, ctx, i);
}
}
static std::pair<std::vector<TensorDistAttr>, std::vector<TensorDistAttr>>
infer_forward(const phi::distributed::SpmdRule &self, const py::args &args) {
VLOG(6) << "infer_forward ";
phi::distributed::InferSpmdContext ctx;
prepare_ctx(&ctx, args);
return self.InferForward(ctx);
}

static std::pair<std::vector<TensorDistAttr>, std::vector<TensorDistAttr>>
infer_backward(const phi::distributed::SpmdRule &self, const py::args &args) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

py::args& can be const?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it work

VLOG(6) << "infer_backward ";
phi::distributed::InferSpmdContext ctx;
prepare_ctx(&ctx, args);
return self.InferBackward(ctx);
}

} // namespace pybind
} // namespace paddle
37 changes: 12 additions & 25 deletions test/auto_parallel/spmd_rules/test_default_data_parallel_rule.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,11 +61,8 @@ def test_default_dp_infer_forward(self):
# 2 inputs 2 outputs, sharded batch axis
in_vec = [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
out_vec = [self.out1_dist_tensor_spec, self.out2_dist_tensor_spec]
result_dist_attrs = self.rule.infer_forward(
[(0, len(in_vec)), (len(in_vec), len(in_vec) + len(out_vec))],
in_vec + out_vec,
[],
)
result_dist_attrs = self.rule.infer_forward(in_vec, out_vec)

self.assertEqual(len(result_dist_attrs), 2)
self.assertEqual(len(result_dist_attrs[0]), 2)
self.assertEqual(len(result_dist_attrs[1]), 2)
Expand All @@ -82,11 +79,9 @@ def test_default_dp_infer_forward(self):
self.out1_dist_tensor_spec,
self.out2_dist_tensor_spec,
]
result_dist_attrs = self.rule.infer_forward(
[(0, len(in_vec)), (len(in_vec), len(in_vec) + len(out_vec))],
in_vec + out_vec,
[],
)

result_dist_attrs = self.rule.infer_forward(in_vec, out_vec)

self.assertEqual(len(result_dist_attrs), 2)
self.assertEqual(len(result_dist_attrs[0]), 1)
self.assertEqual(len(result_dist_attrs[1]), 3)
Expand All @@ -103,11 +98,7 @@ def test_default_dp_infer_forward(self):
in_vec = [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
out_vec = [self.out1_dist_tensor_spec, self.out2_dist_tensor_spec]
with self.assertRaises(NotImplementedError):
result_dist_attrs = self.rule.infer_forward(
[(0, len(in_vec)), (len(in_vec), len(in_vec) + len(out_vec))],
in_vec + out_vec,
[],
)
result_dist_attrs = self.rule.infer_forward(in_vec, out_vec)

def test_default_dp_infer_backward(self):
# replicated out1 from [-1, 0, 1, -1] to [0, -1, -1, -1]
Expand All @@ -117,11 +108,9 @@ def test_default_dp_infer_backward(self):

in_vec = [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
out_vec = [self.out1_dist_tensor_spec, self.out2_dist_tensor_spec]
result_dist_attrs = self.rule.infer_backward(
[(0, len(in_vec)), (len(in_vec), len(in_vec) + len(out_vec))],
in_vec + out_vec,
[],
)

result_dist_attrs = self.rule.infer_backward(in_vec, out_vec)

self.assertEqual(len(result_dist_attrs), 2)
self.assertEqual(len(result_dist_attrs[0]), 2)
self.assertEqual(len(result_dist_attrs[1]), 2)
Expand All @@ -137,11 +126,9 @@ def test_default_dp_infer_backward(self):

in_vec = [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
out_vec = [self.out1_dist_tensor_spec, self.out2_dist_tensor_spec]
result_dist_attrs = self.rule.infer_backward(
[(0, len(in_vec)), (len(in_vec), len(in_vec) + len(out_vec))],
in_vec + out_vec,
[],
)

result_dist_attrs = self.rule.infer_backward(in_vec, out_vec)

self.assertEqual(len(result_dist_attrs), 2)
self.assertEqual(len(result_dist_attrs[0]), 2)
self.assertEqual(len(result_dist_attrs[1]), 2)
Expand Down
Loading