<commit_before>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp distributed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai2.basics import *\n",
    "from fastai2.callback.progress import ProgressCallback\n",
    "from torch.nn.parallel import DistributedDataParallel, DataParallel\n",
    "from torch.utils.data.distributed import DistributedSampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed and parallel training\n",
    "\n",
    "> Callbacks and helper functions to train in parallel or use distributed training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Patch the parallel models so they work with RNNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def reset(self: DataParallel):\n",
    "    if hasattr(self.module, 'reset'): self.module.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ParallelTrainer(Callback):\n",
    "    run_after,run_before = TrainEvalCallback,Recorder\n",
    "    def __init__(self, device_ids): self.device_ids = device_ids\n",
    "    def begin_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,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def to_parallel(self: Learner, device_ids=None):\n",
    "    self.add_cb(ParallelTrainer(device_ids))\n",
    "    return self"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Patch the parallel models so they work with RNNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def reset(self: DistributedDataParallel):\n",
    "    if hasattr(self.module, 'reset'): self.module.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def setup_distrib(gpu=None):\n",
    "    if gpu is None: return gpu\n",
    "    gpu = int(gpu)\n",
    "    torch.cuda.set_device(int(gpu))\n",
    "    if num_distrib() > 1:\n",
    "        torch.distributed.init_process_group(backend='nccl', init_method='env://')\n",
    "    return gpu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to change the dataloaders so that they only get one part of the batch each (otherwise tehre is not point in using distributed training)."
   ]
  },<commit_msg>Fix typo (#180)<commit_after>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp distributed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai2.basics import *\n",
    "from fastai2.callback.progress import ProgressCallback\n",
    "from torch.nn.parallel import DistributedDataParallel, DataParallel\n",
    "from torch.utils.data.distributed import DistributedSampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed and parallel training\n",
    "\n",
    "> Callbacks and helper functions to train in parallel or use distributed training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Patch the parallel models so they work with RNNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def reset(self: DataParallel):\n",
    "    if hasattr(self.module, 'reset'): self.module.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ParallelTrainer(Callback):\n",
    "    run_after,run_before = TrainEvalCallback,Recorder\n",
    "    def __init__(self, device_ids): self.device_ids = device_ids\n",
    "    def begin_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,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def to_parallel(self: Learner, device_ids=None):\n",
    "    self.add_cb(ParallelTrainer(device_ids))\n",
    "    return self"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Patch the parallel models so they work with RNNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def reset(self: DistributedDataParallel):\n",
    "    if hasattr(self.module, 'reset'): self.module.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def setup_distrib(gpu=None):\n",
    "    if gpu is None: return gpu\n",
    "    gpu = int(gpu)\n",
    "    torch.cuda.set_device(int(gpu))\n",
    "    if num_distrib() > 1:\n",
    "        torch.distributed.init_process_group(backend='nccl', init_method='env://')\n",
    "    return gpu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to change the dataloaders so that they only get one part of the batch each (otherwise there is no point in using distributed training)."
   ]
  },
{'fraction_non_alphanumeric': 0.193104491153788, 'fraction_numerical': 0.0013609556933313171, 'mean_word_length': 2.781589479702687, 'pattern_counts': {'":': 114, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4195193', 'n_tokens_mistral': 2402, 'n_tokens_neox': 2372, 'n_words': 458}