actually enable the test

Author: ngc92

.github/workflows/wheel.yml

@@ -127,6 +127,8 @@ jobs:
             args: "--recompute-swiglu --recompute-norm"
           - name: "Offload Opt"
             args: "--offload-opt-m --offload-opt-v --offload-master"
+          - name: "Offload Gradient"
+            args: "--shard-gradients --offload-grads"
           # While not strictly a recomputation, chunked attention should be bitwise identical, too
           - name: "Chunked attention"
             args: "--recompute-att --attn-bwd-chunks=2"

src/binding/python/tests/run.py

@@ -75,6 +75,9 @@ def _create_options(config: TrainingConfig) -> pyllmq.LLamaOptions:
     options.offload_grads = config.offload_grads
     options.persistent_quants = config.persistent_quants
 
+    options.shard_gradients = config.shard_gradients
+    options.shard_weights = config.shard_weights
+
     if config.matmul_dtype:
         options.matmul_type = config.matmul_dtype
     if config.gradient_dtype: