wehub-resource-sync
745 lines
25 KiB
Plaintext
745 lines
25 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f14b034e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| eval: false\n",
|
|
"! [ -e /content ] && pip install -Uqq fastai # upgrade fastai on colab"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3809e462",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| default_exp distributed"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "bf47f64d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"from fastai.basics import *\n",
|
|
"from fastai.callback.progress import ProgressCallback\n",
|
|
"from torch.nn.parallel import DistributedDataParallel, DataParallel\n",
|
|
"from fastai.data.load import _FakeLoader,_loaders\n",
|
|
"from fastai.optimizer import OptimWrapper\n",
|
|
"try: from accelerate import Accelerator\n",
|
|
"except ModuleNotFoundError: pass"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "fb67a461",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Distributed training\n",
|
|
"\n",
|
|
"> Callbacks and helper functions to train in parallel or use distributed training"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "fa2bd944",
|
|
"metadata": {},
|
|
"source": [
|
|
"When using multiple GPUs, you will most probably want to fit using distributed training. \n",
|
|
"\n",
|
|
"Example use can be found:\n",
|
|
"\n",
|
|
"- In the form of a script with [examples/distrib.py](https://github.com/fastai/fastai/blob/master/nbs/examples/distrib.py)\n",
|
|
"- Across all the App Examples with the [Notebook Launcher](https://docs.fast.ai/distributed_app_examples.html)\n",
|
|
"- At the bottom of this notebook for more examples with `notebook_launcher`.\n",
|
|
"\n",
|
|
"To use distributed training, there are only three required steps:\n",
|
|
"\n",
|
|
"1. Add `with learn.distrib_ctx():` before your `learn.fit` call\n",
|
|
"2. Either config Accelerate yourself by running `accelerate config` from the command line, or run:\n",
|
|
"```python\n",
|
|
"from accelerate.utils import write_basic_config\n",
|
|
"write_basic_config()\n",
|
|
"```\n",
|
|
"3. Run your training script with `accelerate launch scriptname.py ...args...`\n",
|
|
"\n",
|
|
"\n",
|
|
"If you're using `untar_data`, or may be downloading or uncompressing data or models as part of your script, you should wrap that code with `rank0_first`, which forces that step to occur first just once on the master process, prior to the remaining processes running it in parallel. E.g. instead of:\n",
|
|
"\n",
|
|
"```python\n",
|
|
"path = untar_data(URLs.IMAGEWOOF_320)\n",
|
|
"```\n",
|
|
"\n",
|
|
"...you instead use:\n",
|
|
"\n",
|
|
"```python\n",
|
|
"path = rank0_first(untar_data, URLs.IMAGEWOOF_320)\n",
|
|
"```\n",
|
|
"\n",
|
|
"See below for details on the full API and underlying helper functions, if needed -- however, note that you will not need anything except the above unless you need to change how the distributed training is implemented."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "9ab29201",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Parallel"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "77dfee6b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def reset(self: DataParallel):\n",
|
|
" \"Patch required `reset` call into `DataParallel`\"\n",
|
|
" if hasattr(self.module, 'reset'): self.module.reset()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "da5948b2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"class ParallelTrainer(Callback):\n",
|
|
" \"Wrap a model `DataParallel` automatically\"\n",
|
|
" run_after,run_before = TrainEvalCallback,Recorder\n",
|
|
" def __init__(self, device_ids): self.device_ids = device_ids\n",
|
|
" def before_fit(self): self.learn.model = DataParallel(self.learn.model, device_ids=self.device_ids)\n",
|
|
" def after_fit(self): self.learn.model = self.learn.model.module"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a42c8240",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def to_parallel(self: Learner, device_ids=None):\n",
|
|
" \"Add `ParallelTrainer` callback to a `Learner`\"\n",
|
|
" self.add_cb(ParallelTrainer(device_ids))\n",
|
|
" return self"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1049d10f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def detach_parallel(self: Learner):\n",
|
|
" \"Remove `ParallelTrainer` callback from a Learner\"\n",
|
|
" self.remove_cb(ParallelTrainer)\n",
|
|
" return self"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "64a51629",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"@contextmanager\n",
|
|
"def parallel_ctx(self: Learner, device_ids=None):\n",
|
|
" \"A context manager to adapt a learner to train in data parallel mode.\"\n",
|
|
" try:\n",
|
|
" self.to_parallel(device_ids)\n",
|
|
" yield self\n",
|
|
" finally: self.detach_parallel()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "aac90696",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Distributed"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "da37422a",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Helper functions"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "84fd7468",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def reset(self: DistributedDataParallel):\n",
|
|
" \"Patch required `reset` call into `DistributedDataParallel`\"\n",
|
|
" if hasattr(self.module, 'reset'): self.module.reset()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3a5fabfb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def setup_distrib(gpu=None):\n",
|
|
" \"Setup this process to participate in distributed training\"\n",
|
|
" if gpu is None: return gpu\n",
|
|
" gpu = int(gpu)\n",
|
|
" torch.cuda.set_device(int(gpu))\n",
|
|
" if num_distrib() > 0: torch.distributed.init_process_group(backend='nccl', init_method='env://')\n",
|
|
" return gpu"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "aaf742dd",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def teardown_distrib():\n",
|
|
" \"Free distributed training resources\"\n",
|
|
" if torch.distributed.is_initialized(): torch.distributed.destroy_process_group()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f685c241",
|
|
"metadata": {},
|
|
"source": [
|
|
"### DataLoader"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5616de4e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def _round_to_multiple(number,multiple): return int(math.ceil(number/multiple)*multiple)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "95952eb9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"class DistributedDL(TfmdDL):\n",
|
|
" \"A `TfmdDL` which splits a batch into equal size pieces for each worker\"\n",
|
|
" def __init__(self,dl,rank=None,world_size=None,device=None):\n",
|
|
" if rank is None: rank=rank_distrib()\n",
|
|
" if world_size is None: world_size=num_distrib()\n",
|
|
" store_attr()\n",
|
|
" if type(dl) == torch.utils.data.DataLoader:\n",
|
|
" shuffle = True if eq(type(dl.sampler), torch.utils.data.RandomSampler) else False\n",
|
|
" self.dl = DataLoader(dataset=dl.dataset, bs=dl.batch_size, num_workers=dl.num_workers, \\\n",
|
|
" pin_memory=dl.pin_memory, timeout=dl.timeout, shuffle=shuffle, drop_last=dl.drop_last, persistent_workers=dl.persistent_workers)\n",
|
|
" self.bs,self.drop_last,self.dataset,fake,self.num_workers,self.offs,self.pin_memory = \\\n",
|
|
" attrgetter('bs','drop_last','dataset','fake_l','num_workers','offs','pin_memory')(self.dl)\n",
|
|
" if device is None: self.device = self.dl.device\n",
|
|
" self.fake_l = _FakeLoader(self, fake.pin_memory, fake.num_workers, fake.timeout, \n",
|
|
" persistent_workers=fake.persistent_workers, \n",
|
|
" pin_memory_device=fake.pin_memory_device)\n",
|
|
" \n",
|
|
" def _broadcast(self,t,rank):\n",
|
|
" \"Broadcasts t from rank `rank` to all other ranks. Returns t so t is same for all ranks after call.\"\n",
|
|
" t = LongTensor(t).cuda() # nccl only works with cuda tensors\n",
|
|
" torch.distributed.broadcast(t,rank)\n",
|
|
" return t.cpu().tolist()\n",
|
|
"\n",
|
|
" def _to_detach(self,b,cpu=True,gather=True): return to_detach(b,cpu,gather) # member func so we can override for test\n",
|
|
" def __len__(self): return _round_to_multiple(len(self.dl),self.world_size)//self.world_size\n",
|
|
" def get_idxs(self):\n",
|
|
" idxs = list(self.dl.get_idxs()) # compute get_idxs in all ranks (we'll only use rank 0 but size must be consistent)\n",
|
|
" idxs = self._broadcast(idxs,0) # broadcast and receive it from rank 0 to all\n",
|
|
" self.n = len(idxs) # we assumed n was dl.n but we really care about number of idxs\n",
|
|
" # add extra samples to make it evenly divisible\n",
|
|
" self.n_padded = _round_to_multiple(self.n,self.world_size)\n",
|
|
" idxs += (idxs * (self.n_padded//self.n))[:self.n_padded-self.n] # idx needs to be repeated when n_padded>>n\n",
|
|
" # slice padded idxs so that each rank gets self.n_padded//self.world_size tensors\n",
|
|
" return idxs[self.rank*self.n_padded//self.world_size:(self.rank+1)*self.n_padded//self.world_size]\n",
|
|
"\n",
|
|
" def before_iter(self):\n",
|
|
" self.i = 0\n",
|
|
" self.dl.before_iter()\n",
|
|
"\n",
|
|
" def randomize(self): self.dl.randomize()\n",
|
|
" def after_batch(self,b):\n",
|
|
" self.i += find_bs(b)\n",
|
|
" return self.dl.after_batch(b)\n",
|
|
"\n",
|
|
" def after_iter(self): self.dl.after_iter()\n",
|
|
" def create_batches(self,samps): return self.dl.create_batches(samps)\n",
|
|
" def to_detach(self,b, cpu=True, gather=True):\n",
|
|
" b = self._to_detach(b, cpu, gather)\n",
|
|
" def _inner(b):\n",
|
|
" if b.ndim>0:\n",
|
|
" # for each rank, compute overflow of read idxs vs self.n and accumulate them to unpad totals after gathering\n",
|
|
" n = sum([min(0,max(-len(b)//self.world_size,\n",
|
|
" self.n-(self.i+r*self.n_padded//self.world_size))) for r in range(self.world_size)])\n",
|
|
" b = b[:n or None]\n",
|
|
" return b\n",
|
|
" return apply(_inner,b) if gather and all(hasattr(self,o) for o in ('i','n','n_padded')) else b"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "59b75372",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"_tmp_file = tempfile.NamedTemporaryFile().name\n",
|
|
"# patch _broadcast with a mocked version so we can test DistributedDL w/o a proper DDP setup\n",
|
|
"@patch\n",
|
|
"def _broadcast(self:DistributedDL,t,rank):\n",
|
|
" t = LongTensor(t)\n",
|
|
" if rank == self.rank: torch.save(t,_tmp_file)\n",
|
|
" else: t.data = torch.load(_tmp_file)\n",
|
|
" return t.tolist()\n",
|
|
"# patch _to_detach with a mocked version that will return right gathered size but -100 for other rank tensors\n",
|
|
"@patch\n",
|
|
"def _to_detach(self:DistributedDL,b,cpu=True,gather=True):\n",
|
|
" b = to_detach(b,cpu,gather)\n",
|
|
" if not gather: return b\n",
|
|
" def _inner(b, cpu, gather):\n",
|
|
" if b.ndim == 0: b=b[None]\n",
|
|
" b = torch.cat([b if i==self.rank else torch.full_like(b,-100) for i in range(self.world_size)])\n",
|
|
" return b if b.ndim > 0 else b.mean()\n",
|
|
" return apply(_inner,b,cpu,gather)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "7df1387e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"dl = TfmdDL(list(range(50)), bs=12, num_workers=2)\n",
|
|
"for i in range(4):\n",
|
|
" dl1 = DistributedDL(dl, i, 4)\n",
|
|
" test_eq(list(dl1), (torch.arange(i*13, i*13+12)%50,torch.tensor([i*13+12])%50))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "78c58ae8",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"dl = torch.utils.data.DataLoader(list(range(50)), batch_size=12, num_workers=2)\n",
|
|
"for i in range(4):\n",
|
|
" dl1 = DistributedDL(dl, i, 4)\n",
|
|
" test_eq(list(dl1), (torch.arange(i*13, i*13+12)%50,torch.tensor([i*13+12])%50))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b536c7e4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"dl = TfmdDL(list(zip(range(50),range(100,150))), bs=12, num_workers=4)\n",
|
|
"for i in range(4):\n",
|
|
" dl1 = DistributedDL(dl, i, 4)\n",
|
|
" test_eq(list(dl1), [(torch.arange(i*13, i*13+12)%50,100+torch.arange(i*13, i*13+12)%50),\n",
|
|
" ((torch.tensor([i*13+12])%50),100+torch.tensor([i*13+12])%50)])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "33a06d92",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"dl = TfmdDL(list(range(50)), bs=12, num_workers=2,drop_last=True)\n",
|
|
"for i in range(4):\n",
|
|
" dl1 = DistributedDL(dl, i, 4)\n",
|
|
" test_eq(list(dl1), [torch.arange(i*13, i*13+12)%50])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3f941f29",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"dl = TfmdDL(list(zip(range(12),range(100,112))), bs=12, num_workers=4)\n",
|
|
"res,dls = [],[]\n",
|
|
"for i in range(5): dls.append(DistributedDL(dl, i, 5))\n",
|
|
"for b in zip(*dls):\n",
|
|
" for r in range(5):\n",
|
|
" d=L(dls[r].to_detach(b[r]))\n",
|
|
" test_eq(d.map(lambda x:(x!=-100).sum().item()),(3,3) if r!=4 else (0,0))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1fbff277",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"dl = TfmdDL(list(range(10)), bs=4, num_workers=2, shuffle=True)\n",
|
|
"res = []\n",
|
|
"for i in range(3):\n",
|
|
" dl1 = DistributedDL(dl, i, 3)\n",
|
|
" b = list(dl1)[0]\n",
|
|
" bd = dl1.to_detach(b)\n",
|
|
" test_eq(b[:None if i<2 else 2],bd[4*i:4*(i+1)])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3eaadcf6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"from fastai.callback.data import WeightedDL"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "90008ad4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"dl = WeightedDL(list(range(50)), bs=16, num_workers=2, shuffle=True,wgts=list(np.arange(50)>=25))\n",
|
|
"res = []\n",
|
|
"for i in range(4):\n",
|
|
" dl1 = DistributedDL(dl, i, 4)\n",
|
|
" res += list(dl1)[0].tolist()\n",
|
|
"test(res,[25]*len(res),operator.ge) # all res >=25\n",
|
|
"test(res,[25]*len(res),lambda a,b: not (a<b)) # all res NOT < 25"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "54a10a9a",
|
|
"metadata": {},
|
|
"source": [
|
|
"### DistributedTrainer -"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "bf701a2a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"_hidden_params = [\"mixed_precision\", \"fp16\", \"log_with\", \"logging_dir\", \"step_scheduler_with_optimizer\"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9a5bc62a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"class DistributedTrainer(Callback):\n",
|
|
" \"Wrap `model` in `DistributedDataParallel` and `dls` in `DistributedDL`\"\n",
|
|
" order = 11\n",
|
|
" @delegates(Accelerator, but=_hidden_params)\n",
|
|
" def __init__(self,\n",
|
|
" sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`\n",
|
|
" **kwargs\n",
|
|
" ):\n",
|
|
" store_attr()\n",
|
|
" self.accelerator = Accelerator(**kwargs)\n",
|
|
" def before_fit(self):\n",
|
|
" self.learn.model = self.accelerator.prepare(\n",
|
|
" nn.SyncBatchNorm.convert_sync_batchnorm(self.model) if self.sync_bn else self.model\n",
|
|
" )\n",
|
|
" self.old_dls = list(self.dls)\n",
|
|
" self.learn.dls.loaders = [self._wrap_dl(dl) for dl in self.dls]\n",
|
|
" if rank_distrib(): self.learn.logger=noop\n",
|
|
"\n",
|
|
" def _wrap_dl(self, dl): return dl if isinstance(dl,DistributedDL) else DistributedDL(dl, device=self.learn.model.device)\n",
|
|
" def _backward(self): self.accelerator.backward(self.learn.loss_grad)\n",
|
|
" \n",
|
|
" def before_train(self): self.learn.dl = self._wrap_dl(self.learn.dl)\n",
|
|
" def before_validate(self): self.learn.dl = self._wrap_dl(self.learn.dl)\n",
|
|
" def after_fit(self): self.learn.model,self.learn.dls.loaders = self.learn.model.module,self.old_dls"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "972de58c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"@delegates(Accelerator, but=_hidden_params)\n",
|
|
"def to_distributed(self: Learner,\n",
|
|
" sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`\n",
|
|
" **kwargs\n",
|
|
" ):\n",
|
|
" \"Add `AcceleratedTrainer` to a learner, and configures an Accelerator\"\n",
|
|
" self.add_cb(DistributedTrainer(sync_bn, **kwargs))\n",
|
|
" if rank_distrib(): self.remove_cb(ProgressCallback)\n",
|
|
" return self"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "be731f5d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def detach_distributed(self: Learner):\n",
|
|
" \"Remove `DistributedTrainer` from a learner\"\n",
|
|
" if num_distrib() <=1: return self\n",
|
|
" self.remove_cb(DistributedTrainer)\n",
|
|
" if rank_distrib() and not hasattr(self, 'progress'): self.add_cb(ProgressCallback())\n",
|
|
" return self"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d6301077",
|
|
"metadata": {},
|
|
"source": [
|
|
"### `distrib_ctx` context manager"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "993a237b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"@contextmanager\n",
|
|
"@delegates(Accelerator, but=_hidden_params)\n",
|
|
"def distrib_ctx(self: Learner,\n",
|
|
" sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`\n",
|
|
" in_notebook=False, # Whether we are launching from a notebook or not\n",
|
|
" **kwargs\n",
|
|
" ):\n",
|
|
" \"A context manager to adapt a learner to train in distributed data parallel mode.\"\n",
|
|
" try: import accelerate\n",
|
|
" except ImportError as e: \n",
|
|
" e.args = [\"Accelerate is required. Install with `pip install accelerate`\"]\n",
|
|
" raise\n",
|
|
" # Adapt self to DistributedDataParallel, yield, and cleanup afterwards.\n",
|
|
" cleanup_dpg = False\n",
|
|
" try:\n",
|
|
" if in_notebook:\n",
|
|
" cuda_id = rank_distrib()\n",
|
|
" if not torch.distributed.is_initialized():\n",
|
|
" setup_distrib(cuda_id)\n",
|
|
" cleanup_dpg = torch.distributed.is_initialized()\n",
|
|
" if not rank_distrib(): print(\"Training Learner...\")\n",
|
|
" if num_distrib(): self.to_distributed(sync_bn, **kwargs)\n",
|
|
" yield self\n",
|
|
" finally:\n",
|
|
" self.detach_distributed()\n",
|
|
" if cleanup_dpg: teardown_distrib()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ae3a345a",
|
|
"metadata": {},
|
|
"source": [
|
|
"`distrib_ctx` prepares a learner to train in distributed data parallel mode. It assumes the script/code will either be ran through the command line via `accelerate launch` or through the `notebook_launcher` function from Accelerate. It also assumes that `accelerate` has been configured through either running `write_basic_config()` or calling `accelerate config` through the CLI and answering the prompts.\n",
|
|
"\n",
|
|
"Typical usage:\n",
|
|
"\n",
|
|
"```\n",
|
|
"with learn.distrib_ctx(): learn.fit(.....)\n",
|
|
"```\n",
|
|
"\n",
|
|
"It attaches a `DistributedTrainer` callback and `DistributedDL` data loader to the learner, then executes `learn.fit(.....)`. Upon exiting the context, it removes the `DistributedTrainer` and `DistributedDL`, and destroys any locally created distributed process group. The process is still attached to the GPU though."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c0640003",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def rank0_first(func, *args, **kwargs):\n",
|
|
" \"Execute `func` in the Rank-0 process first, then in other ranks in parallel.\"\n",
|
|
" if args or kwargs: func = partial(func, *args, **kwargs)\n",
|
|
" dummy_l = Learner(DataLoaders(device='cpu'), nn.Linear(1,1), loss_func=lambda: 0)\n",
|
|
" with dummy_l.distrib_ctx():\n",
|
|
" if not rank_distrib(): res = func()\n",
|
|
" distrib_barrier()\n",
|
|
" if rank_distrib(): res = func()\n",
|
|
" return res"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "683e77a3",
|
|
"metadata": {},
|
|
"source": [
|
|
"`rank0_first` calls `f()` in rank-0 process first, then in parallel on the rest, in distributed training mode. In single process, non-distributed training mode, `f()` is called only once as expected.\n",
|
|
"\n",
|
|
"One application of `rank0_first()` is to make fresh downloads via `untar_data` safe in distributed training scripts launched by `python -m fastai.launch <script>`:\n",
|
|
"\n",
|
|
"<code>path = untar_data(URLs.IMDB)</code>\n",
|
|
"\n",
|
|
"becomes:\n",
|
|
"\n",
|
|
"<code>path = rank0_first(lambda: untar_data(URLs.IMDB))</code>\n",
|
|
"\n",
|
|
"Some learner factory methods may use `untar_data` to download pretrained models:\n",
|
|
"\n",
|
|
"<code>learn = text_classifier_learner(dls, AWD_LSTM, drop_mult=0.5, metrics=accuracy)</code>\n",
|
|
"\n",
|
|
"becomes:\n",
|
|
"\n",
|
|
"<code>learn = rank0_first(lambda: text_classifier_learner(dls, AWD_LSTM, drop_mult=0.5, metrics=accuracy))</code>\n",
|
|
"\n",
|
|
"Otherwise, multiple processes will download at the same time and corrupt the data."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "b9aa3a1a",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Notebook Launcher\n",
|
|
"\n",
|
|
"Accelerate provides a [notebook_launcher](https://huggingface.co/docs/accelerate/launcher) functionality to let you keep using your Jupyter Notebook as you would, but train in a distributed setup!\n",
|
|
"\n",
|
|
"First, make sure accelerate is properly configured. You can either run `accelerate config` from the command line, or have an autofilled configuration setup by running in the first cell of your notebook:\n",
|
|
"\n",
|
|
"```python\n",
|
|
"from accelerate.utils import write_basic_config\n",
|
|
"write_basic_config()\n",
|
|
"```\n",
|
|
"After Accelerate is configured, to utilize the `notebook_launcher` functionality migrate your training into a function, and pass this to `notebook_launcher`, such as:\n",
|
|
"\n",
|
|
"```python\n",
|
|
"---\n",
|
|
"from fastai.vision.all import *\n",
|
|
"from fastai.distributed import *\n",
|
|
"\n",
|
|
"def train():\n",
|
|
" set_seed(99, True)\n",
|
|
" path = untar_data(URLs.PETS)/'images'\n",
|
|
" dls = ImageDataLoaders.from_name_func(\n",
|
|
" path, get_image_files(path), valid_pct=0.2,\n",
|
|
" label_func=lambda x: x[0].isupper(), item_tfms=Resize(224))\n",
|
|
" \n",
|
|
" learn = vision_learner(dls, resnet34, metrics=error_rate).to_fp16()\n",
|
|
" with learn.distrib_ctx(in_notebook=True):\n",
|
|
" learn.fine_tune(1)\n",
|
|
"---\n",
|
|
"from accelerate import notebook_launcher\n",
|
|
"notebook_launcher(train, num_processes=2)\n",
|
|
"---\n",
|
|
"```"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7cc4c55c",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Export -"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ec8ed782",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"from nbdev import nbdev_export\n",
|
|
"nbdev_export()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e0076d87",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"jupytext": {
|
|
"split_at_heading": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "python3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|