github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 02:26:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[doc] polish shardformer doc (#4779)

Browse Source 
* fix example format in docstring

* polish shardformer doc
This commit is contained in:
Baizhou Zhang
2023-09-26 10:57:47 +08:00
committed by GitHub
parent 26cd6d850c
commit a2db75546d
8 changed files with 215 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
colossalai/booster/plugin/torch_fsdp_plugin.py
View File

@@ -143,16 +143,17 @@ class TorchFSDPPlugin(DPPluginBase):
"""
Plugin for PyTorch FSDP.
Example:
>>> from colossalai.booster import Booster
>>> from colossalai.booster.plugin import TorchFSDPPlugin
>>>
>>> model, train_dataset, optimizer, criterion = ...
>>> plugin = TorchFSDPPlugin()
```python
from colossalai.booster import Booster
from colossalai.booster.plugin import TorchFSDPPlugin
>>> train_dataloader = plugin.prepare_train_dataloader(train_dataset, batch_size=8)
>>> booster = Booster(plugin=plugin)
>>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)
model, train_dataset, optimizer, criterion = ...
plugin = TorchFSDPPlugin()
train_dataloader = plugin.prepare_train_dataloader(train_dataset, batch_size=8)
booster = Booster(plugin=plugin)
model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)
```
Args:
See https://pytorch.org/docs/stable/fsdp.html for details.