Skip to content

Commit 03e4e07

Browse files
lockshawelliottslaughterseemamirch
committed
Add support for Replicate op in distributed training
This is primarily (though not entirely) changes in task-spec, as realm-execution currently just uses multiple point-to-point copies to perform the Replicate. The key task-spec changes are: - copy_insertion does not insert copies around `Replicate`s, since copies are already done as part of the replicate - Since eventually `Replicate`s should be executed as a single collective call, shard_expansion no longer fully shard-expands `Replicate`s, limiting the expansion to the granularity of (what will eventually be) the collective call. To implement this, slots in the dynamic graph now need to optionally have a shard specified, since these partially shard-expanded task still need to output fully shard-expanded tensors/values. In realm-execution: - There's a new redop registry for reduction ops - There's an abstration for broadcast and reduction collectives (currently the actual implementation is just naive point-to-point operations, but the interface/abstraction allows this implementation to be easily swapped out for something more performant once it's added) - Machine information caching in RealmContext is simplified Additional changes made as part of this PR: - Commutativity tests for task-spec passes were added, and corresponding fixes made to the passes to satisfy the new tests - Most unordered_ containers are replaced by their ordered counterparts, for more predictable/deterministic behavior, particularly when distributed - A number of additional tests for task-spec passes were added to prevent regressions and check edge cases - proj was updated to print more datastructures in json by default, for easier debugging Co-authored-by: Elliott Slaughter <slaughter@cs.stanford.edu> Co-authored-by: Seema Mirchandaney <seemah@stanford.edu>
1 parent dfab204 commit 03e4e07

1,290 files changed

Lines changed: 18406 additions & 8693 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.proj.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ testsuite_macro = "FF_TEST_SUITE"
33
namespace_name = "FlexFlow"
44
header_extension = ".h"
55
doxygen = true
6+
build_tool = "ninja"
67
cuda_launch_cmd = [
78
"nixGL",
89
"--",

‎bin/run-model/src/run-model/main.cc‎

Lines changed: 41 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -71,47 +71,47 @@ int main(int argc, char **argv) {
7171
char **realm_argv = realm_args.data();
7272
RealmManager manager(&realm_argc, &realm_argv);
7373

74-
ControllerTaskResult result = manager.start_controller([&](RealmContext
75-
&ctx) {
76-
MappedParallelComputationGraph mpcg = [&]() {
77-
std::ifstream f(mapped_pcg_json);
78-
nlohmann::json mpcg_json = nlohmann::json::parse(f);
79-
return from_v1(mpcg_json.get<V1MappedParallelComputationGraph>());
80-
}();
81-
82-
// instantiate computation graph
83-
OptimizerAttrs optimizer_attrs =
84-
OptimizerAttrs{SGDOptimizerAttrs{/*lr=*/0.001,
85-
/*momentum=*/0.9,
86-
/*nesterov=*/false,
87-
/*weight_decay=*/0.001}};
88-
89-
std::unordered_map<DynamicValueAttrs, DynamicTensorAccessor> input_tensors;
90-
91-
DistributedFfHandle device_handle =
92-
create_distributed_ff_handle(ctx,
93-
/*workSpaceSize=*/1024 * 1024,
94-
/*allowTensorOpMathConversion=*/true);
95-
96-
PCGInstance pcg_instance = create_pcg_instance(
97-
/*ctx=*/ctx,
98-
/*mpcg=*/mpcg,
99-
/*optimizer=*/optimizer_attrs,
100-
/*loss=*/std::nullopt,
101-
/*input_tensors=*/input_tensors,
102-
/*profiling_settings=*/ProfilingSettings{0, 0},
103-
/*device_handle=*/device_handle,
104-
/*device_type=*/DeviceType::GPU);
105-
106-
// begin training loop
107-
int num_epochs = 5;
108-
for (int i = 0; i < num_epochs; i++) {
109-
perform_all_passes_for_pcg_instance(
110-
/*instance=*/pcg_instance,
111-
/*profiling_settings=*/ProfilingSettings{0, 1},
112-
/*device_handle=*/device_handle);
113-
}
114-
});
74+
ControllerTaskResult result =
75+
manager.start_controller([&](RealmContext &ctx) {
76+
MappedParallelComputationGraph mpcg = [&]() {
77+
std::ifstream f(mapped_pcg_json);
78+
nlohmann::json mpcg_json = nlohmann::json::parse(f);
79+
return from_v1(mpcg_json.get<V1MappedParallelComputationGraph>());
80+
}();
81+
82+
// instantiate computation graph
83+
OptimizerAttrs optimizer_attrs =
84+
OptimizerAttrs{SGDOptimizerAttrs{/*lr=*/0.001,
85+
/*momentum=*/0.9,
86+
/*nesterov=*/false,
87+
/*weight_decay=*/0.001}};
88+
89+
std::map<DynamicValueAttrs, DynamicTensorAccessor> input_tensors;
90+
91+
DistributedFfHandle device_handle =
92+
create_distributed_ff_handle(ctx,
93+
/*workSpaceSize=*/1024 * 1024,
94+
/*allowTensorOpMathConversion=*/true);
95+
96+
PCGInstance pcg_instance = create_pcg_instance(
97+
/*ctx=*/ctx,
98+
/*mpcg=*/mpcg,
99+
/*optimizer=*/optimizer_attrs,
100+
/*loss=*/std::nullopt,
101+
/*input_tensors=*/input_tensors,
102+
/*profiling_settings=*/ProfilingSettings{0, 0},
103+
/*device_handle=*/device_handle,
104+
/*device_type=*/DeviceType::GPU);
105+
106+
// begin training loop
107+
int num_epochs = 5;
108+
for (int i = 0; i < num_epochs; i++) {
109+
perform_all_passes_for_pcg_instance(
110+
/*instance=*/pcg_instance,
111+
/*profiling_settings=*/ProfilingSettings{0, 1},
112+
/*device_handle=*/device_handle);
113+
}
114+
});
115115
result.wait();
116116

117117
return 0;

‎bin/sp-ization-benchmarking/include/sp-ization-benchmarking/distributions.h‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
#define _FLEXFLOW_BIN_SP_IZATION_BENCHMARKING_INCLUDE_SP_IZATION_BENCHMARKING_DISTRIBUTIONS_H
33

44
#include "utils/graph/node/node.dtg.h"
5+
#include <map>
56
#include <random>
6-
#include <unordered_map>
7-
#include <unordered_set>
7+
#include <set>
88

99
namespace FlexFlow {
1010

@@ -55,21 +55,19 @@ struct GaussianNoise {
5555
};
5656

5757
template <typename Dist>
58-
std::unordered_map<Node, float>
59-
make_cost_map(std::unordered_set<Node> const &nodes,
60-
Dist const &distribution) {
61-
std::unordered_map<Node, float> cost_map;
58+
std::map<Node, float> make_cost_map(std::set<Node> const &nodes,
59+
Dist const &distribution) {
60+
std::map<Node, float> cost_map;
6261
for (Node const &node : nodes) {
6362
cost_map[node] = distribution();
6463
}
6564
return cost_map;
6665
}
6766

6867
template <typename Noise>
69-
std::unordered_map<Node, float>
70-
add_noise_to_cost_map(std::unordered_map<Node, float> cost_map,
71-
Noise const &noise) {
72-
std::unordered_map<Node, float> noisy_cost_map;
68+
std::map<Node, float> add_noise_to_cost_map(std::map<Node, float> cost_map,
69+
Noise const &noise) {
70+
std::map<Node, float> noisy_cost_map;
7371
for (auto const &[node, cost] : cost_map) {
7472
noisy_cost_map[node] = noise() * cost;
7573
}

‎bin/sp-ization-benchmarking/src/sp-ization-benchmarking/sample_graphs.cc‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ DiGraph make_full_taso_nasnet(size_t num_reduction_cells, size_t N) {
108108
(i % (N + 1) == N) ? make_reduction_taso_nasnet_cell()
109109
: make_normal_taso_nasnet_cell();
110110
Node cell_output = get_only(get_terminal_nodes(s));
111-
std::unordered_map<Node, Node> node_map = parallel_extend(g, s);
111+
std::map<Node, Node> node_map = parallel_extend(g, s);
112112
later_input = node_map.at(later_input);
113113
earlier_input = node_map.at(earlier_input);
114114
cell_output = node_map.at(cell_output);
@@ -331,8 +331,8 @@ DiGraph make_2_terminal_random_dag(size_t num_nodes, float p, size_t step) {
331331
}
332332
}
333333
}
334-
std::unordered_set<Node> sinks = get_terminal_nodes(g);
335-
std::unordered_set<Node> sources = get_initial_nodes(g);
334+
std::set<Node> sinks = get_terminal_nodes(g);
335+
std::set<Node> sources = get_initial_nodes(g);
336336
Node sink = g.add_node();
337337
Node source = g.add_node();
338338
for (Node s : sources) {

‎flake.lock‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎flake.nix‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@
155155
expect
156156
universal-ctags
157157
ninja
158+
tig
158159
])
159160
(with pkgs.python3Packages; [
160161
gitpython

‎lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.dtg.toml‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,8 @@ includes = [
2424
]
2525

2626
src_includes = [
27-
"utils/hash/unordered_map.h",
28-
"utils/fmt/unordered_map.h",
29-
"utils/ord/unordered_map.h",
27+
"utils/hash/map.h",
28+
"utils/fmt/map.h",
3029
]
3130

3231
[[fields]]
@@ -35,15 +34,15 @@ type = "::FlexFlow::PCGOperatorAttrs"
3534

3635
[[fields]]
3736
name = "input_shapes"
38-
type = "std::unordered_map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
37+
type = "std::map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
3938

4039
[[fields]]
4140
name = "weight_shapes"
42-
type = "std::unordered_map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
41+
type = "std::map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
4342

4443
[[fields]]
4544
name = "output_shapes"
46-
type = "std::unordered_map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
45+
type = "std::map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
4746

4847
[[fields]]
4948
name = "optimizer_attrs"

‎lib/compiler/include/compiler/cost_estimator/op_cost_metrics.dtg.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ name = "OpCostMetrics"
33
type = "struct"
44
features = [
55
"eq",
6+
"ord",
67
"fmt",
78
"hash",
89
]

‎lib/compiler/include/compiler/cost_estimator/op_cost_metrics.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
namespace FlexFlow {
88

99
bool is_pareto_optimal_in(OpCostMetrics const &,
10-
std::unordered_set<OpCostMetrics> const &);
10+
std::set<OpCostMetrics> const &);
1111

1212
OpCostMetrics make_op_cost_metrics_from_runtime_only(
1313
RuntimeOnlyOpCostMetrics const &runtime_only,

‎lib/compiler/include/compiler/cost_estimator/runtime_only_op_cost_estimate_key.dtg.toml‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,8 @@ includes = [
1717
]
1818

1919
src_includes = [
20-
"utils/hash/unordered_map.h",
21-
"utils/ord/unordered_map.h",
22-
"utils/fmt/unordered_map.h",
20+
"utils/hash/map.h",
21+
"utils/fmt/map.h",
2322
]
2423

2524
[[fields]]
@@ -28,15 +27,15 @@ type = "::FlexFlow::PCGOperatorAttrs"
2827

2928
[[fields]]
3029
name = "input_shapes"
31-
type = "std::unordered_map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
30+
type = "std::map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
3231

3332
[[fields]]
3433
name = "weight_shapes"
35-
type = "std::unordered_map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
34+
type = "std::map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
3635

3736
[[fields]]
3837
name = "output_shapes"
39-
type = "std::unordered_map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
38+
type = "std::map<::FlexFlow::TensorSlotName, ::FlexFlow::ParallelTensorShape>"
4039

4140
[[fields]]
4241
name = "machine_view"

0 commit comments

Comments
 (0)