Loader ckp fixes (#119)

Updates `get_latest` and `get_oldest` to use the same sorting function, and allows the dataloader ckp handler to pass in its custom sort manually. Removes the bug where excessive path joins lead to repeated path prefixes in dataloader ckp loading. Fixes GPTBigCode signatures used for speculator training, to match superclass signatures (currently preventing other PRs from landing).

Includes and subsumes #110 and #96. Full credit to @weiji14 and @Akash-Nayak respectively

Author: daviswer

fms_fsdp/utils/checkpointing_utils.py

@@ -20,35 +20,45 @@
 from torch.distributed.fsdp import StateDictType
 
 
-def get_latest(targdir, qualifier=lambda x: True):
-    """Fetch the latest file or folder written to target directory, subject to name passing the qualifier fn.
-    If directory is empty or nonexistent or no items qualify, return None."""
+def get_latest(targdir, qualifier=lambda x: True, key=os.path.getctime):
+    """
+    Fetch the full path of the latest file or folder written to target directory,
+    subject to name passing the qualifier fn.
+    Optional key fn can be used for custom sorting.
+    Both functions take full path arguments.
+    If directory is empty or nonexistent or no items qualify, return None.
+    """
     if os.path.exists(targdir) and len(os.listdir(targdir)) > 0:
         latest = max(
             [
                 os.path.join(targdir, x)
                 for x in os.listdir(targdir)
                 if qualifier(os.path.join(targdir, x))
             ],
-            key=lambda path: int(path.split("/")[-1].split("_")[1]),
+            key=key,
         )
-        return os.path.join(targdir, latest)
+        return latest
     return None
 
 
-def get_oldest(targdir, qualifier=lambda x: True):
-    """Fetch the oldest file or folder written to target directory, subject to name passing the qualifier fn.
-    If directory is empty or nonexistent or no items qualify, return None."""
+def get_oldest(targdir, qualifier=lambda x: True, key=os.path.getctime):
+    """
+    Fetch the full path of the oldest file or folder written to target directory,
+    subject to name passing the qualifier fn.
+    Optional key fn can be used for custom sorting.
+    Both functions take full path arguments.
+    If directory is empty or nonexistent or no items qualify, return None.
+    """
     if os.path.exists(targdir) and len(os.listdir(targdir)) > 0:
         oldest = min(
             [
                 os.path.join(targdir, x)
                 for x in os.listdir(targdir)
                 if qualifier(os.path.join(targdir, x))
             ],
-            key=os.path.getctime,
+            key=key,
         )
-        return os.path.join(targdir, oldest)
+        return oldest
     return None
 
 

fms_fsdp/utils/dataset_utils.py

@@ -32,7 +32,7 @@
     rescaling (i.e. counters, RNG states), and `reshard_params`, which are lists that can be 
     re-distributed over workers (i.e. buffers).
 
-Our loaders obey the following type heirarchy: 
+Our loaders obey the following type hierarchy: 
 torch.data.IterableDataset -> _StatefulDataset -> _WrapperDataset. 
 `_StatefulDataset` implements state and checkpointing logic. A `_WrapperDataset` holds a 
 single `_StatefulDataset` and iterates via calling its wrapped dataset any number of times, 
@@ -510,8 +510,8 @@ def _validate_ckp_path(self, path: str, verbose: bool = False):
                     f"  Dataset: No valid checkpoint detected at {path}, dataset starting from scratch."
                 )
             return ""
-        # Check latest path
-        latest = os.path.join(path, get_latest(path))
+        # Check latest path, using ckp naming syntax
+        latest = get_latest(path, key=lambda path: int(path.split("_")[-2]))
         if verbose:
             self.report(f"Checkpoint detected at {latest}")
         # If item is not a folder, exit early

speculator/train_speculator_utils.py

@@ -1,7 +1,7 @@
 import os
 import re
 import time
-from typing import Any, Callable, Mapping, MutableMapping, Optional, Tuple, Union
+from typing import Any, Callable, List, MutableMapping, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
@@ -437,11 +437,12 @@ class EmbedGPTBigCode(GPTBigCode):
     # Overrides the forward function of GPTBigCode to allow returning embedding vectors
     def forward(
         self,
-        x: torch.LongTensor,
+        x: torch.Tensor,
         mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value_states: Optional[Tuple[torch.FloatTensor,]] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value_states: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
         use_cache: bool = False,
+        only_last_token: bool = False,
         attn_algorithm: Optional[str] = None,
         include_embeds: bool = False,
     ):