wehub-resource-sync
96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
# Copyright (c) 2023 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.
|
|
|
|
from argparse import ArgumentParser
|
|
|
|
|
|
def parse_arguments():
|
|
parser = ArgumentParser()
|
|
|
|
# for distributed strategy
|
|
parser.add_argument(
|
|
"--dp_degree", type=int, required=True, help="dp degree"
|
|
)
|
|
parser.add_argument(
|
|
"--mp_degree", type=int, required=True, help="mp degree"
|
|
)
|
|
parser.add_argument(
|
|
"--pp_degree", type=int, required=True, help="pp degree"
|
|
)
|
|
parser.add_argument(
|
|
"--vpp_degree", type=int, required=True, help="vpp degree"
|
|
)
|
|
parser.add_argument(
|
|
"--sharding_degree", type=int, required=True, help="sharding degree"
|
|
)
|
|
parser.add_argument(
|
|
"--sharding_stage", type=int, required=True, help="sharding stage"
|
|
)
|
|
parser.add_argument(
|
|
"--micro_batch_size", type=int, required=True, help="micro batch size"
|
|
)
|
|
parser.add_argument(
|
|
"--use_recompute", type=bool, required=True, help="use recompute"
|
|
)
|
|
parser.add_argument(
|
|
"--recompute_granularity",
|
|
type=str,
|
|
required=True,
|
|
choices=["None", "core_attn", "full_attn", "full"],
|
|
help="recompute granularity",
|
|
)
|
|
|
|
# for model config
|
|
parser.add_argument(
|
|
"--hidden_size", type=int, required=False, help="hidden size"
|
|
)
|
|
parser.add_argument(
|
|
"--num_attention_heads",
|
|
type=int,
|
|
required=False,
|
|
help="number of attention heads",
|
|
)
|
|
parser.add_argument(
|
|
"--num_layers", type=int, required=False, help="number of hidden layers"
|
|
)
|
|
parser.add_argument(
|
|
"--max_sequence_length",
|
|
type=int,
|
|
required=False,
|
|
help="maximum sequence length",
|
|
)
|
|
parser.add_argument(
|
|
"--vocab_size", type=int, required=False, help="vocabulary size"
|
|
)
|
|
parser.add_argument(
|
|
"--intermediate_size",
|
|
type=int,
|
|
required=False,
|
|
help="intermediate size",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def get_model_memory_usage(args):
|
|
# evaluate model memory usage based on distributed strategy and model setting
|
|
raise NotImplementedError(
|
|
"Please implement this function for memory usage estimation based on distributed strategy and model setting."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_arguments()
|
|
print(get_model_memory_usage(args))
|