Building block for specifying pipeline-parallel modules. LayerSpec stores the type information and parameters for each stage in a PipelineModule. For example: .. code-block:: python nn.Sequence( torch.nn.Linear(self.in_dim, self.hidden_dim, bias=False),
| 28 | |
| 29 | |
| 30 | class LayerSpec: |
| 31 | """Building block for specifying pipeline-parallel modules. |
| 32 | |
| 33 | LayerSpec stores the type information and parameters for each stage in a |
| 34 | PipelineModule. For example: |
| 35 | |
| 36 | .. code-block:: python |
| 37 | |
| 38 | nn.Sequence( |
| 39 | torch.nn.Linear(self.in_dim, self.hidden_dim, bias=False), |
| 40 | torch.nn.Linear(self.hidden_hidden, self.out_dim) |
| 41 | ) |
| 42 | |
| 43 | becomes |
| 44 | |
| 45 | .. code-block:: python |
| 46 | |
| 47 | layer_specs = [ |
| 48 | LayerSpec(torch.nn.Linear, self.in_dim, self.hidden_dim, bias=False), |
| 49 | LayerSpec(torch.nn.Linear, self.hidden_hidden, self.out_dim)] |
| 50 | ] |
| 51 | """ |
| 52 | |
| 53 | def __init__(self, typename, *module_args, **module_kwargs): |
| 54 | self.typename = typename |
| 55 | self.module_args = module_args |
| 56 | self.module_kwargs = module_kwargs |
| 57 | |
| 58 | if not issubclass(typename, nn.Module): |
| 59 | raise RuntimeError('LayerSpec only supports torch.nn.Module types.') |
| 60 | |
| 61 | if dist.is_initialized(): |
| 62 | self.global_rank = dist.get_rank() |
| 63 | else: |
| 64 | self.global_rank = -1 |
| 65 | |
| 66 | def __repr__(self): |
| 67 | return ds_utils.call_to_str(self.typename.__name__, self.module_args, self.module_kwargs) |
| 68 | |
| 69 | def build(self, log=False): |
| 70 | """Build the stored specification.""" |
| 71 | if log: |
| 72 | logger.info(f'RANK={self.global_rank} building {repr(self)}') |
| 73 | |
| 74 | return self.typename(*self.module_args, **self.module_kwargs) |
| 75 | |
| 76 | |
| 77 | class TiedLayerSpec(LayerSpec): |