()
| 54 | printonce(*args, **kwargs) |
| 55 | |
| 56 | def init_dist(): |
| 57 | if T.distributed.is_initialized(): |
| 58 | print0('Distributed already initialized') |
| 59 | rank = T.distributed.get_rank() |
| 60 | local_rank = int(os.environ.get('LOCAL_RANK', 0)) |
| 61 | world_size = T.distributed.get_world_size() |
| 62 | else: |
| 63 | try: |
| 64 | rank = int(os.environ['RANK']) |
| 65 | local_rank = int(os.environ['LOCAL_RANK']) |
| 66 | world_size = int(os.environ['WORLD_SIZE']) |
| 67 | device = f'cuda:{local_rank}' |
| 68 | T.cuda.set_device(device) |
| 69 | T.distributed.init_process_group(backend='nccl', timeout=timedelta(minutes=30), rank=rank, world_size=world_size, device_id=T.device(device)) |
| 70 | print(f'Rank {rank} of {world_size}.') |
| 71 | except Exception as e: |
| 72 | print0once(f'Not initializing distributed env: {e}') |
| 73 | rank = 0 |
| 74 | local_rank = 0 |
| 75 | world_size = 1 |
| 76 | return rank, local_rank, world_size |
| 77 | |
| 78 | def load_ckpt(load_from_location, expected_hash=None): |
| 79 | os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1' #Disable this to speed up debugging errors with downloading from the hub |
nothing calls this directly
no test coverage detected