Skip to content

initialize.rst

Latest commit

 

History

History
38 lines (26 loc) · 1.2 KB

initialize.rst

File metadata and controls

38 lines (26 loc) · 1.2 KB

Training Setup

Argument Parsing

DeepSpeed uses the argparse library to supply commandline configuration to the DeepSpeed runtime. Use deepspeed.add_config_arguments() to add DeepSpeed's builtin arguments to your application's parser.

parser = argparse.ArgumentParser(description='My training script.')
parser.add_argument('--local_rank', type=int, default=-1,
                    help='local rank passed from distributed launcher')
# Include DeepSpeed configuration arguments
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()
.. autofunction:: deepspeed.add_config_arguments

Training Initialization

The entrypoint for all training with DeepSpeed is deepspeed.initialize().

Example usage:

model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
                                                     model=net,
                                                     model_parameters=net.parameters())
.. autofunction:: deepspeed.initialize