refactor(args): validate the page-cache cap inside the Megatron stream branch

The cap's checks ran before the block that still allows torchtitan streaming.
The non-negative check now sits in the Adam branch, and one assert after the
Megatron branch rejects the cap everywhere else.

Author: guapisolo

miles/utils/arguments.py

@@ -1375,17 +1375,6 @@ def miles_validate_args(args):
             f"chunk={args.offload_train_disk_chunk_mb}MB"
         )
 
-    if args.stream_optimizer_state_page_cache_mb is not None:
-        assert args.stream_optimizer_state_to_disk, (
-            "--stream-optimizer-state-page-cache-mb caps the page cache of the streamed optimizer state; "
-            "it needs --stream-optimizer-state-to-disk"
-        )
-        assert args.stream_optimizer_state_page_cache_mb >= 0, "--stream-optimizer-state-page-cache-mb must be >= 0"
-        assert args.train_backend == "megatron", (
-            "--stream-optimizer-state-page-cache-mb caps the Megatron Adam bucket store; "
-            "it is not supported on other backends"
-        )
-
     if args.stream_optimizer_state_to_disk:
         assert args.train_backend in (
             "megatron",
@@ -1421,6 +1410,9 @@ def miles_validate_args(args):
                 "file-backed mappings the chunked offloader owns"
             )
         else:
+            assert (
+                args.stream_optimizer_state_page_cache_mb is None or args.stream_optimizer_state_page_cache_mb >= 0
+            ), "--stream-optimizer-state-page-cache-mb must be >= 0"
             assert (
                 args.use_distributed_optimizer
             ), "--stream-optimizer-state-to-disk requires the distributed optimizer"
@@ -1469,6 +1461,11 @@ def miles_validate_args(args):
             f"chunk={args.offload_train_disk_chunk_mb}MB, moments={args.stream_optimizer_state_moment_dtype}, "
             f"page cache {'left to the kernel' if page_cache_mb is None else f'capped at {page_cache_mb}MB'}"
         )
+    else:
+        assert args.stream_optimizer_state_page_cache_mb is None, (
+            "--stream-optimizer-state-page-cache-mb caps the Megatron Adam bucket store; it needs "
+            "--stream-optimizer-state-to-disk and is not supported on other backends"
+        )
 
     if args.async_max_concurrent_samples is not None:
         assert args.async_max_concurrent_samples >= args.n_samples_per_prompt, (