| #!/Users/james/src/MMaDA/venv/bin/python | |
| from benchmarks.communication.run_all import main | |
| from benchmarks.communication.constants import * | |
| from benchmarks.communication.utils import * | |
| import os | |
| import sys | |
| # Run the same file with deepspeed launcher. This is required since setuptools will auto-detect python files and insert a python shebang for both 'scripts' and 'entry_points', and this benchmarks require the DS launcher | |
| required_env = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"] | |
| if not all(map(lambda v: v in os.environ, required_env)): | |
| import subprocess | |
| r = subprocess.check_output(["which", "ds_bench"]) | |
| ds_bench_bin = r.decode('utf-8').strip() | |
| safe_cmd = ["deepspeed", ds_bench_bin] + sys.argv[1:] | |
| subprocess.run(safe_cmd) | |
| else: | |
| args = benchmark_parser().parse_args() | |
| rank = args.local_rank | |
| main(args, rank) | |