wehub-resource-sync
73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import List, Tuple
|
|
|
|
import torch
|
|
from torch.fx import GraphModule
|
|
|
|
from ..util import get_deepcompile_handle
|
|
from ..fx import (add_postprocess, move_primals_to_head, _make_node_meta, add_end_backward,
|
|
replace_reduce_outputs_with_none, should_release_reduce_buckets)
|
|
|
|
NAME = "zero1_compile"
|
|
|
|
|
|
def add_z1_reduce_fw(gm: GraphModule, graph_id: int, profiling_results, param_manager, use_z2=False) -> GraphModule:
|
|
|
|
dc = get_deepcompile_handle()
|
|
param_indices = profiling_results[graph_id].param_indices
|
|
# Need this before profiling
|
|
if use_z2:
|
|
dc.register_graph_z2(graph_id, [v[1] for v in param_indices])
|
|
else:
|
|
dc.register_graph_z1(graph_id, [v[1] for v in param_indices])
|
|
|
|
return gm
|
|
|
|
|
|
def add_z1_reduce_bw(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int, bool]],
|
|
param_manager) -> GraphModule:
|
|
|
|
graph = gm.graph
|
|
pm = param_manager[graph_id]
|
|
_, param_name_to_grad = pm.get_bwd_mapping(graph)
|
|
|
|
for param_name in pm.param_names:
|
|
|
|
grad_node = param_name_to_grad[param_name]
|
|
|
|
assert param_name in pm.ds_ids, f"param_name={param_name} not in ds_ids"
|
|
ds_id = pm.ds_ids[param_name]
|
|
|
|
new_node = add_postprocess(graph,
|
|
grad_node,
|
|
torch.ops.dc.reduce_grad.default,
|
|
extra_args=[graph_id, ds_id],
|
|
name=f"reduce_param_{param_name}",
|
|
meta=_make_node_meta(grad_node, param_name, True))
|
|
new_node.meta["val"] = None
|
|
|
|
gm.graph = move_primals_to_head(graph)
|
|
|
|
add_end_backward(gm.graph, graph_id, should_release_reduce_buckets(graph_order, graph_id))
|
|
replace_reduce_outputs_with_none(gm.graph)
|
|
|
|
return gm
|
|
|
|
|
|
def add_z1_reduce(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int, bool]], profiling_results,
|
|
create_inputs_fn, mem_budget: float, param_manager, bwd: bool) -> GraphModule:
|
|
if bwd:
|
|
return add_z1_reduce_bw(gm, graph_id, graph_order, param_manager)
|
|
return add_z1_reduce_fw(gm, graph_id, profiling_results, param_manager, use_z2=False)
|
|
|
|
|
|
def add_z2_reduce(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int, bool]], profiling_results,
|
|
create_inputs_fn, mem_budget: float, param_manager, bwd: bool) -> GraphModule:
|
|
if bwd:
|
|
return add_z1_reduce_bw(gm, graph_id, graph_order, param_manager)
|
|
return add_z1_reduce_fw(gm, graph_id, profiling_results, param_manager, use_z2=True)
|