[doc] update nvme offload doc

docs/requirements-doc-test.txt

@@ -1,2 +1,6 @@
 colossalai
 torch
+packaging
+tensornvme
+psutil
+transformers

docs/source/en/features/nvme_offload.md

@@ -1,3 +1,4 @@
+<!-- doc-test-command: torchrun --standalone --nproc_per_node=1 nvme_offload.py  -->
 # NVMe offload
 
 Author: Hongxin Liu
@@ -36,12 +37,225 @@ pip install tensornvme
 
 We implement NVMe offload of optimizer states for Adam ([CPUAdam](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.nn.optimizer.cpu_adam.html) and [HybridAdam](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.nn.optimizer.hybrid_adam.html)).
 
+
+<!--- doc-test-ignore-start -->
+
 ```python
 from colossalai.nn.optimizer import CPUAdam, HybridAdam
 
 optimizer = HybridAdam(model.parameters(), lr=1e-3, nvme_offload_fraction=1.0, nvme_offload_dir='./')
 ```
 
+<!--- doc-test-ignore-end -->
+
 `nvme_offload_fraction` is the fraction of optimizer states to be offloaded to NVMe. `nvme_offload_dir` is the directory to save NVMe offload files. If `nvme_offload_dir` is `None`, a random temporary directory will be used.
 
 It's compatible with all parallel methods in ColossalAI.
+
+> ⚠ It only offloads optimizer states on CPU. This means it only affects CPU training or Zero/Gemini with offloading.
+
+## Exampls
+
+Let's start from two simple examples -- training GPT with different methods. These examples relies on `transformers`.
+
+We should install denpendencies first:
+
+```shell
+pip install psutil transformers
+```
+
+First, we import essential packages and modules:
+
+```python
+import os
+import time
+from typing import Dict, Optional
+
+import psutil
+import torch
+import torch.nn as nn
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
+
+import colossalai
+from colossalai.nn.optimizer import HybridAdam
+from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
+from colossalai.utils.model.colo_init_context import ColoInitContext
+```
+
+Then we define a loss function:
+
+```python
+class GPTLMLoss(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def forward(self, logits, labels):
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        # Flatten the tokens
+        return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)),
+                            shift_labels.view(-1))
+```
+
+And we define some utility functions, which generates random data, computes the number of paramters of a model and get memory usage of current process:
+
+```python
+def get_data(batch_size: int, seq_len: int,
+             vocab_size: int, device: Optional[str] = None) -> Dict[str, torch.Tensor]:
+    device = torch.cuda.current_device() if device is None else device
+    input_ids = torch.randint(vocab_size, (batch_size, seq_len),
+                              device=device)
+    attn_mask = torch.ones_like(input_ids)
+    return dict(input_ids=input_ids, attention_mask=attn_mask)
+
+
+def get_model_numel(model: nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters())
+
+
+def get_mem_usage() -> int:
+    proc = psutil.Process(os.getpid())
+    return proc.memory_info().rss
+```
+
+We first try to train GPT model on CPU:
+
+```python
+def train_cpu(nvme_offload_fraction: float = 0.0):
+    config = GPT2Config()
+    model = GPT2LMHeadModel(config)
+    criterion = GPTLMLoss()
+    optimizer = HybridAdam(model.parameters(), nvme_offload_fraction=nvme_offload_fraction)
+    print(f'Model numel: {get_model_numel(model) / 1024**3:.3f} B')
+
+    start = time.time()
+    for step in range(3):
+        data = get_data(4, 128, config.vocab_size, device='cpu')
+        outputs = model(**data)
+        loss = criterion(outputs.logits, data['input_ids'])
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        print(f'[{step}] loss: {loss.item():.3f}')
+
+    print(f'Time: {time.time() - start:.3f} s')
+    print(f'Mem usage: {get_mem_usage() / 1024**2:.3f} MB')
+```
+
+Run without NVME offload:
+
+```python
+train_cpu(0.0)
+```
+
+We may get below output:
+
+```
+Model numel: 0.116 B
+[0] loss: 10.953
+[1] loss: 10.974
+[2] loss: 10.965
+Time: 7.739 s
+Mem usage: 5966.445 MB
+```
+
+And then run with (full) NVME offload:
+
+```python
+train_cpu(1.0)
+```
+
+We may get:
+
+```
+Model numel: 0.116 B
+[0] loss: 10.951
+[1] loss: 10.994
+[2] loss: 10.984
+Time: 8.527 s
+Mem usage: 4968.016 MB
+```
+
+For GPT2-S, which has 0.116 billion parameters, its optimizer states take about 0.928 GB memory. And NVME offload saves about 998 MB memory, which meets our expectations.
+
+Then we can train GPT model with Gemini. The placement policy of Gemini should be `"auto"`, `"cpu"` or `"const"`.
+
+```python
+def train_gemini_cpu(nvme_offload_fraction: float = 0.0):
+    colossalai.launch_from_torch({})
+    config = GPT2Config()
+    with ColoInitContext(device=torch.cuda.current_device()):
+        model = GPT2LMHeadModel(config)
+    criterion = GPTLMLoss()
+    optimizer = HybridAdam(model.parameters(), nvme_offload_fraction=nvme_offload_fraction)
+    print(f'Model numel: {get_model_numel(model) / 1024**3:.3f} B')
+
+    gemini_config = dict(strict_ddp_mode=True, device=torch.cuda.current_device(),
+                         placement_policy='cpu', pin_memory=True, hidden_dim=config.n_embd)
+    model = zero_model_wrapper(model, zero_stage=3, gemini_config=gemini_config)
+    optimizer = zero_optim_wrapper(model, optimizer, initial_scale=2**5)
+
+    start = time.time()
+    for step in range(3):
+        data = get_data(4, 128, config.vocab_size)
+        outputs = model(**data)
+        loss = criterion(outputs.logits, data['input_ids'])
+        optimizer.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+        print(f'[{step}] loss: {loss.item():.3f}')
+
+    print(f'Time: {time.time() - start:.3f} s')
+    print(f'Mem usage: {get_mem_usage() / 1024**2:.3f} MB')
+```
+
+Run without NVME offload:
+
+```python
+train_gemini_cpu(0.0)
+```
+
+We may get:
+
+```
+Model numel: 0.116 B
+searching chunk configuration is completed in 0.27 s.
+used number: 118.68 MB, wasted number: 0.75 MB
+total wasted percentage is 0.63%
+[0] loss: 10.953
+[1] loss: 10.938
+[2] loss: 10.969
+Time: 2.997 s
+Mem usage: 5592.227 MB
+```
+
+And run with (full) NVME offload:
+
+```python
+train_gemini_cpu(1.0)
+```
+
+We may get:
+
+```
+Model numel: 0.116 B
+searching chunk configuration is completed in 0.27 s.
+used number: 118.68 MB, wasted number: 0.75 MB
+total wasted percentage is 0.63%
+[0] loss: 10.953
+[1] loss: 10.938
+[2] loss: 10.969
+Time: 3.691 s
+Mem usage: 5298.344 MB
+```
+
+NVME offload saves about 294 MB memory. Note that enabling `pin_memory` of Gemini can accelerate training but increase memory usage. So this result also meets our expectation. If we disable `pin_memory`, we can aslo observe a memory usage drop about 900 MB.
+
+## API Reference
+
+{{ autodoc:colossalai.nn.optimizer.HybridAdam }}
+
+{{ autodoc:colossalai.nn.optimizer.CPUAdam }}