Update deepspeed model-parallel docs (#18091)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
b8d4a70db7
commit
d79eaae334
|
@ -187,12 +187,10 @@ Here's an example using that uses ``wrap`` to create your model:
|
|||
|
||||
|
||||
class MyModel(pl.LightningModule):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
def configure_model(self):
|
||||
self.linear_layer = nn.Linear(32, 32)
|
||||
self.block = nn.Sequential(nn.Linear(32, 32), nn.Linear(32, 32))
|
||||
|
||||
def configure_model(self):
|
||||
# modules are sharded across processes
|
||||
# as soon as they are wrapped with `wrap`.
|
||||
# During the forward/backward passes, weights get synced across processes
|
||||
|
|
Loading…
Reference in New Issue