support cpu offload for stage3 (#49196)

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage3.py

@@ -428,10 +428,11 @@ class GroupShardedStage3(nn.Layer):
                 place=core.CPUPlace(),
                 name="slice@" + param.name,
             )
-            with device_guard():
-                param.master_weight = paddle.cast(
-                    param.fw_storage, Type.fp32.value
-                )
+            if param.trainable:
+                with device_guard():
+                    param.master_weight = paddle.cast(
+                        param.fw_storage, Type.fp32.value
+                    )
         else:
             param.fw_storage = core.eager.Tensor(
                 value=buffer._slice(start, end), name="slice@" + param.name

python/paddle/fluid/tests/unittests/collective/fleet/dygraph_group_sharded_stage3_offload.py

@@ -40,6 +40,9 @@ class MLP(fluid.Layer):
 
         self._linear1 = Linear(linear_size, linear_size)
         self._linear2 = Linear(linear_size, linear_size)
+        # test for trainable & untrainable offload
+        self._linear2.weight.stop_gradient = False
+        self._linear2.bias.stop_gradient = False
         self._linear3 = Linear(linear_size, 10)
 
     def forward(self, inputs):
@@ -119,7 +122,7 @@ def train_mlp(
             img, label = data
             label.stop_gradient = True
             img.stop_gradient = True
-            with paddle.amp.auto_cast(True, level='O2'):
+            with paddle.amp.auto_cast(use_pure_fp16, level='O2'):
                 out = model(img)
                 loss = paddle.nn.functional.cross_entropy(
                     input=out, label=label