ශුන්ය-ඩීපීමතක ප්රශස්තිකරණය

32import functools
33from typing import List, Optional, Tuple
34
35import torch
36import torch.distributed as dist
37from torch import nn

40class Zero3Layer(nn.Module):

49    chunk: List[nn.Parameter]

51    chunk_size: List[int]

53    TRAINING_PARAMS_IDX = 0

56    param_refs: List[List[nn.Parameter]]

59    fetch_stream: Optional[torch.cuda.Stream]

61    backup_stream: Optional[torch.cuda.Stream]

63    prev_layer: List['Zero3Layer']

65    next_layer: List['Zero3Layer']

67    layer_idx: int

70    is_fetched: bool

73    device: torch.device

75    dtype: torch.dtype

77    module: nn.Module

79    world_size: int

81    def __init__(self, module: nn.Module, rank: int, world_size: int, device: torch.device, dtype: torch.dtype):

89        super().__init__()

92        self.device = device
93        self.dtype = dtype
94        self.module = module
95        self.prev_layer = []
96        self.next_layer = []
97        self.is_fetched = False
98        self.world_size = world_size
99        self.layer_idx = -1
100        self.fetch_stream = None
101        self.backup_stream = None
102
103        with torch.no_grad():

105            all_param_refs = [p for p in self.parameters()]

108            for p in all_param_refs:
109                p._orig_shape = p.shape

112            for p in all_param_refs:
113                assert p.dtype == dtype, "All parameters should have same dtype"

116            self.param_refs = [[p for p in all_param_refs if p.requires_grad],
117                               [p for p in all_param_refs if not p.requires_grad]]
118            del all_param_refs

122            if rank == 0:

124                merged_params = [self._merge_and_pad_params(ps) for ps in self.param_refs]

126                self.chunk_size = [(len(p) // world_size if p is not None else 0) for p in merged_params]

128                dist.broadcast(torch.tensor(self.chunk_size, device=device), src=0)
129            else:

131                chunk_size = torch.tensor([0, 0], device=device)

133                dist.broadcast(chunk_size, src=0)
134                self.chunk_size = chunk_size.tolist()

138            self.chunk = [nn.Parameter(self._empty((s,)), requires_grad=i == self.TRAINING_PARAMS_IDX)
139                          for i, s in enumerate(self.chunk_size)]

142            chunk = self._empty((sum(self.chunk_size),))
143
144            if rank == 0:

146                all_params = torch.cat([p.view(world_size, -1) for p in merged_params], dim=-1).view(-1)
147                del merged_params

150                dist.scatter(chunk, list(all_params.split(sum(self.chunk_size))))
151                del all_params
152            else:

154                dist.scatter(chunk)

157            chunk = chunk.split(self.chunk_size)
158            for i, c in enumerate(chunk):
159                self.chunk[i].data[:] = c
160            del chunk

163            self._cleanup_params()

166            self._backward_hook_ref = self.register_full_backward_hook(self._backward_hook)  # type: ignore

168    def _merge_and_pad_params(self, params: List[nn.Parameter]) -> torch.Tensor:

173        size = sum(p.shape.numel() for p in params)

176        if size % self.world_size != 0:
177            padding_fixed = self.world_size - (size % self.world_size)

179        else:
180            padding_fixed = 0

182        padding = self._empty((padding_fixed,))

184        return torch.cat([p.view(-1) for p in params] + [padding], dim=0)

186    def get_trainable_chunk(self) -> List[nn.Parameter]:

193        if len(self.chunk[self.TRAINING_PARAMS_IDX]) == 0:
194            return []

197        return [self.chunk[self.TRAINING_PARAMS_IDX]]

199    def _empty(self, shape: Tuple[int, ...]) -> torch.Tensor:

203        return torch.empty(shape, device=self.device, dtype=self.dtype)

205    @torch.no_grad()
206    def _cleanup_params(self):

214        self.is_fetched = False

217        for ps in self.param_refs:
218            for p in ps:

220                p.data.record_stream(torch.cuda.current_stream())

222                assert p.data.storage_offset() == 0, "The tensor is not the sole occupant of the storage."

226                p.data.storage().resize_(0)  # This is what actually clears the memory

228                assert p.grad is None, 'Gradients should be None'

230    @torch.no_grad()
231    def fetch_params(self):

239        if self.is_fetched:
240            return

243        self.is_fetched = True

246        if sum(self.chunk_size) == 0:
247            return

250        with torch.cuda.stream(self.fetch_stream):

252            buffer = self._empty((self.world_size * sum(self.chunk_size),))

254            buffers = list(buffer.split(sum(self.chunk_size)))

257            chunk = torch.cat(self.chunk, dim=0)

260            dist.all_gather(buffers, chunk)

263            params = buffer.view(-1, sum(self.chunk_size)).split(self.chunk_size, dim=1)

265            buffer.record_stream(self.fetch_stream)
266            for b in buffers:
267                b.record_stream(self.fetch_stream)
268            buffer.record_stream(self.fetch_stream)
269            del buffer
270            del buffers

273            params = [p.reshape(-1) for p in params]

276            for cont, ps in zip(params, self.param_refs):

278                if not ps:
279                    continue

282                offset = 0

284                for p in ps:

286                    shape = p._orig_shape  # type: ignore[attr-defined]

288                    p.data.storage().resize_(shape.numel())

290                    p.data[:] = cont[offset: offset + shape.numel()].reshape(shape)

292                    p.data.record_stream(self.fetch_stream)

294                    offset += shape.numel()

297                cont.record_stream(self.fetch_stream)

300            del params

302    def forward(self, *args, **kwargs):

309        self.fetch_params()

312        torch.cuda.current_stream().wait_stream(self.fetch_stream)

316        for layer in self.next_layer:
317            layer.fetch_params()

320        if torch.is_grad_enabled():
321            self._add_backward_hooks()

324        res = self.module(*args, **kwargs)

330        if not torch.is_grad_enabled() or self.next_layer:
331            self._cleanup_params()
332
333        return res

335    def _add_backward_hooks(self):

341        self._backward_hook_handles = 0

344        for p in self.param_refs[self.TRAINING_PARAMS_IDX]:

346            assert not hasattr(p, "_hook_handle"), 'Parameter has already been hooked'

348            p_tmp = p.expand_as(p)

351            grad_acc = p_tmp.grad_fn.next_functions[0][0]

353            handle = grad_acc.register_hook(
354                functools.partial(self._post_backward_hook, p))

356            p._hook_handle = handle

358            self._backward_hook_handles += 1

360    def _backward_event(self):

368        self._backward_hook_handles -= 1

372        if self._backward_hook_handles == -1:
373            self._backup_grads()
374            self._cleanup_params()

377        for layer in self.prev_layer:
378            layer.fetch_params()

380    def _post_backward_hook(self, p: nn.Parameter, *args):

385        p._hook_handle.remove()  # type: ignore[attr-defined]
386        delattr(p, "_hook_handle")

389        self._backward_event()

391    def _backward_hook(self, *args, **kwargs):

396        self._backward_event()

399        torch.cuda.current_stream().wait_stream(self.fetch_stream)

402        return None

404    @torch.no_grad()
405    def _backup_grads(self):

410        if self.chunk_size[self.TRAINING_PARAMS_IDX] == 0:
411            return

414        with torch.cuda.stream(self.backup_stream):

416            buffer = self._empty((self.world_size * self.chunk_size[self.TRAINING_PARAMS_IDX],))

418            buffers = list(buffer.split(self.chunk_size[self.TRAINING_PARAMS_IDX]))

421            offset = 0

423            for p in self.param_refs[self.TRAINING_PARAMS_IDX]:

425                shape = p._orig_shape  # type: ignore[attr-defined]
426                buffer[offset: offset + shape.numel()] = p.grad.view(-1)

428                offset += shape.numel()

430                p.grad = None

433            grad = self._empty((self.chunk_size[self.TRAINING_PARAMS_IDX],))

436            dist.reduce_scatter(grad, buffers)

439            for b in buffers:
440                b.record_stream(self.fetch_stream)
441            buffer.record_stream(self.fetch_stream)
442            del buffer
443            del buffers

446            self.chunk[self.TRAINING_PARAMS_IDX].grad = grad
447            del grad

450class Zero3Sequential(nn.Module):

454    def __init__(self, modules: List[Zero3Layer]):

458        super().__init__()

461        self.fetch_stream = torch.cuda.Stream()

463        self.backup_stream = torch.cuda.Stream()

466        for i in range(len(modules)):

468            modules[i].layer_idx = i

470            modules[i].fetch_stream = self.fetch_stream
471            modules[i].backup_stream = self.backup_stream

473            if i + 1 < len(modules):
474                modules[i].next_layer.append(modules[i + 1])

476            if i - 1 >= 0:
477                modules[i].prev_layer.append(modules[i - 1])

480        self.module_list = nn.ModuleList(modules)

482    def get_trainable_chunk(self):

484        return sum([m.get_trainable_chunk() for m in self.module_list], [])

486    def forward(self, x: torch.Tensor):

488        torch.cuda.current_stream().wait_stream(self.backup_stream)

491        for m in self.module_list:
492            x = m(x)

495        return x