deepspeed

1
#!/usr/bin/env python3
2

3
import argparse
4
import json
5

6
import deepspeed
7
from deepspeed.elasticity import compute_elastic_config
8

9
if __name__ == '__main__':
10
    parser = argparse.ArgumentParser()
11
    parser.add_argument('-c', '--config', type=str, help="DeepSpeed config json")
12
    parser.add_argument('-w', '--world-size', type=int, default=0, help="Intended/current world size")
13
    args = parser.parse_args()
14
    ds_config = json.load(open(args.config, 'r'))
15

16
    ds_version = deepspeed.__version__
17

18
    elastic_config = ds_config['elasticity']
19
    print('------------------------------------------')
20
    print("Elasticity config:")
21
    print('------------------------------------------')
22
    print(json.dumps(elastic_config, indent=4, sort_keys=True))
23

24
    if args.world_size > 0:
25
        final_batch_size, valid_gpus, micro_batch_size = compute_elastic_config(ds_config=ds_config,
26
                                                                                target_deepspeed_version=ds_version,
27
                                                                                world_size=args.world_size)
28
        print('------------------------------------------')
29
        print(f"Calculated results for world size {args.world_size}:")
30
        print('------------------------------------------')
31
        print(f'final_batch_size .... {final_batch_size}')
32
        print(f'valid_gpus .......... {valid_gpus}')
33
        print(f'micro_batch_size .... {micro_batch_size}')
34
    else:
35
        final_batch_size, valid_gpus = compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
36
        print('------------------------------------------')
37
        print("Calculated results:")
38
        print('------------------------------------------')
39
        print(f'final_batch_size .... {final_batch_size}')
40
        print(f'valid_gpus .......... {valid_gpus}')
41