Optimal usage example

  1# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
  2# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  3# SPDX-License-Identifier: Apache-2.0
  4#
  5# Licensed under the Apache License, Version 2.0 (the "License");
  6# you may not use this file except in compliance with the License.
  7# You may obtain a copy of the License at
  8#
  9# http://www.apache.org/licenses/LICENSE-2.0
 10#
 11# Unless required by applicable law or agreed to in writing, software
 12# distributed under the License is distributed on an "AS IS" BASIS,
 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 14# See the License for the specific language governing permissions and
 15# limitations under the License.
 16
 17
 18# This example demonstrates how to integrate ``inprocess.Wrapper()`` into an
 19# existing PyTorch training codebase.
 20#
 21# This example show the optimal usage:
 22# - only the training loop and objects depending on a torch distributed process
 23# group are being restarted upon a failure
 24# - process-group-independent objects (e.g. TCPStore, Model, Optimizer) are
 25# created once, and reused between all restart iterations to minimize restart
 26# latency
 27#
 28# NOTE: inprocess.Wrapper is not fully compatible with modern
 29# ``torch.distributed.run``, because it automatically terminates all local
 30# workers upon any local worker process failure; in this case inprocess.Wrapper
 31# can only recover from transient faults that don't terminate any of the
 32# training processes
 33
 34import argparse
 35import datetime
 36import logging
 37import os
 38import pathlib
 39import random
 40import time
 41from typing import Optional
 42
 43os.environ['TORCH_CPP_LOG_LEVEL'] = 'error'
 44import torch
 45
 46import nvidia_resiliency_ext.inprocess as inprocess
 47
 48raise_timestamp = None
 49
 50
 51def parse_args():
 52    parser = argparse.ArgumentParser(
 53        description='Inprocess Restart Optimal Example',
 54        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
 55    )
 56
 57    parser.add_argument(
 58        '--size',
 59        default=64,
 60        type=int,
 61        help='model hidden size',
 62    )
 63    parser.add_argument(
 64        '--layers',
 65        default=4,
 66        type=int,
 67        help='number of layers',
 68    )
 69    parser.add_argument(
 70        '--log-interval',
 71        default=100,
 72        type=int,
 73        help='logging interval',
 74    )
 75    parser.add_argument(
 76        '--chkpt-interval',
 77        default=100,
 78        type=int,
 79        help='checkpointing interval',
 80    )
 81    parser.add_argument(
 82        '--total-iterations',
 83        default=1000000,
 84        type=int,
 85        help='total training iterations',
 86    )
 87    parser.add_argument(
 88        '--seed',
 89        default=None,
 90        type=int,
 91        help='random seed, time-based if None',
 92    )
 93    parser.add_argument(
 94        '--path',
 95        default='/tmp/',
 96        type=str,
 97        help='directory for the checkpoint file',
 98    )
 99    parser.add_argument(
100        '--fault-prob',
101        default=0.001,
102        type=float,
103        help='fault injection probability',
104    )
105    parser.add_argument(
106        '--device',
107        default='cpu',
108        choices=['cpu', 'cuda'],
109        help='device',
110    )
111    parser.add_argument(
112        '--log-level',
113        type=lambda s: logging._nameToLevel[s.upper()],
114        default=logging.INFO,
115        help='logging level',
116    )
117
118    return parser.parse_args()
119
120
121# TCPStore created by the Wrapper uses ``(MASTER_PORT + 1)`` port for the
122# internal Wrapper's TCPStore to avoid conflicts with application's TCPStore
123# listening on ``(MASTER_PORT + 2)``, and with a TCPStore created by
124# ``torch.distributed.run`` listening on ``MASTER_PORT``.
125#
126# An instance of ``inprocess.CallWrapper` is automatically injected into
127# wrapped function arguments when Wrapper is invoked.
128@inprocess.Wrapper(
129    store_kwargs={'port': int(os.getenv('MASTER_PORT', 29500)) + 1},
130    health_check=inprocess.health_check.CudaHealthCheck(),
131)
132def train(
133    base_store,
134    model,
135    opt,
136    backend,
137    device,
138    timeout,
139    args,
140    call_wrapper: Optional[inprocess.CallWrapper] = None,
141):
142    global raise_timestamp
143    if raise_timestamp is not None:
144        restart_latency = time.perf_counter() - raise_timestamp
145        logging.info(f'restart latency: {restart_latency:.3f}s')
146    raise_timestamp = None
147
148    log_interval = args.log_interval
149    chkpt_interval = args.chkpt_interval
150
151    rank = int(os.environ['RANK'])
152    world_size = int(os.environ['WORLD_SIZE'])
153
154    # Create a new Store by adding a prefix based on the current inprocess
155    # restart iteration. PrefixStore wraps the baseline TCPStore which is
156    # reused for all restart iterations
157    store = torch.distributed.PrefixStore(str(call_wrapper.iteration), base_store)
158
159    torch.distributed.init_process_group(
160        backend,
161        store=store,
162        rank=rank,
163        world_size=world_size,
164        timeout=timeout,
165    )
166    model_ddp = torch.nn.parallel.DistributedDataParallel(model)
167
168    iteration = 0
169    loss = torch.tensor(float('nan'))
170    checkpoint_path = pathlib.Path(args.path) / 'checkpoint.pt'
171
172    # Application loads state from the latest checkpoint on every restart
173    # iteration of the wrapped function.
174    if checkpoint_path.exists():
175        checkpoint = torch.load(checkpoint_path)
176        model.load_state_dict(checkpoint['model'])
177        opt.load_state_dict(checkpoint['opt'])
178        torch.set_rng_state(checkpoint['rng'])
179        iteration = checkpoint['iteration']
180
181    if args.seed is not None:
182        random.seed(args.seed + iteration * world_size + rank)
183    else:
184        random.seed(time.perf_counter_ns())
185
186    for iteration in range(iteration, args.total_iterations):
187
188        # Application periodically saves a checkpoint. The checkpoint allows
189        # the application to continue from previous state after a restart.
190        if iteration % chkpt_interval == chkpt_interval - 1:
191            torch.distributed.barrier()
192            if rank == 0:
193                checkpoint = {
194                    'model': model.state_dict(),
195                    'opt': opt.state_dict(),
196                    'rng': torch.get_rng_state(),
197                    'iteration': iteration,
198                }
199                # Saving the checkpoint is performed within atomic() context
200                # manager to ensure that the main thread won't execute
201                # torch.save while a restart procedure is in progress.
202                with call_wrapper.atomic():
203                    torch.save(checkpoint, checkpoint_path)
204
205        # Randomly trigger an example fault
206        if random.random() < args.fault_prob:
207            raise_timestamp = time.perf_counter()
208            raise RuntimeError(f'example fault at {iteration=} from {rank=}')
209
210        inp = torch.rand(args.size, args.size).to(device)
211        model.zero_grad()
212        out = model_ddp(inp)
213        loss = out.square().mean()
214        loss.backward()
215        opt.step()
216        loss.item()
217
218        if rank == 0 and iteration % log_interval == log_interval - 1:
219            logging.info(f'{rank=} {iteration=} {loss.item()=}')
220
221
222def main():
223    args = parse_args()
224    logging.basicConfig(
225        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
226        level=args.log_level,
227    )
228    logging.info(f'{args}')
229
230    local_rank = int(os.environ['LOCAL_RANK'])
231
232    if args.device == 'cuda':
233        torch.cuda.set_device(local_rank)
234        device = torch.device('cuda')
235        backend = 'nccl'
236        timeout = datetime.timedelta(seconds=150)
237    elif args.device == 'cpu':
238        device = torch.device('cpu')
239        backend = 'gloo'
240        timeout = datetime.timedelta(seconds=10)
241    else:
242        raise RuntimeError
243
244    # All objects created in ``main()`` are constructed only once, and reused
245    # for all restart iterations.
246    if args.seed is not None:
247        torch.manual_seed(args.seed)
248    model = torch.nn.Sequential(
249        *[torch.nn.Linear(args.size, args.size) for _ in range(args.layers)]
250    ).to(device)
251    opt = torch.optim.Adam(model.parameters(), lr=1e-5)
252
253    # TCPStore uses ``(MASTER_PORT + 2)`` to avoid conflicts with TCPStore
254    # created by ``torch.distributed.run`` and listening on ``MASTER_PORT``,
255    # and Wrapper's TCPStore listening on ``(MASTER_PORT + 1)``
256    store = torch.distributed.TCPStore(
257        host_name=os.environ['MASTER_ADDR'],
258        port=int(os.environ['MASTER_PORT']) + 2,
259        world_size=int(os.environ['WORLD_SIZE']),
260        is_master=(int(os.environ['RANK']) == 0),
261        multi_tenant=True,
262        wait_for_workers=True,
263        use_libuv=True,
264    )
265
266    # Call the wrapped function.
267    # ``train()`` is automatically restarted to recover from faults.
268    train(store, model, opt, backend, device, timeout, args)
269
270
271if __name__ == '__main__':
272    main()