[brief] Fixes the distributed flag when using torchrun.

[detailed]
- Due to the order of operations, we cannot rely on the distributed info
  from the core module to determine the world size when torchrun is
  used. Since the distributed package hasn't been initialized yet, we
  have to pull the world size from the environment variable directly.

src/helios/trainer.py

@@ -3,6 +3,7 @@
 import dataclasses
 import enum
 import itertools
+import os
 import pathlib
 import re
 import time
@@ -686,7 +687,7 @@ def _setup_device_flags(self, use_cpu: bool | None):
         self._is_distributed = (
             len(self._gpu_ids) > 1
             if not self._is_torchrun
-            else dist.get_dist_info().world_size > 1
+            else int(os.environ["WORLD_SIZE"]) > 1
         )
 
     def _save_checkpoint(self, state: TrainingState) -> None: