Add dataloaders for all training phases

src/utils/dataloader.py

@@ -0,0 +1,53 @@
+import os
+import numpy as np
+import torch
+
+def load_tokens(filename):
+    # simple numpy array extraction into torch tensor matrix memory
+    npt = np.load(filename)
+    npt = npt.astype(np.int32)
+    ptt = torch.tensor(npt, dtype=torch.long)
+    return ptt
+
+class DataLoaderLite:
+    def __init__(self, B, T, process_rank, num_processes, split, master_process):
+        self.B = B
+        self.T = T
+        self.process_rank = process_rank
+        self.num_processes = num_processes
+        assert split in {'train', 'val'}
+        
+        # swap directory to parse the dedicated edu fineweb repository location
+        data_root = "data/edu_fineweb10B"
+        shards = os.listdir(data_root)
+        shards = [s for s in shards if split in s]
+        shards = sorted(shards)
+        shards = [os.path.join(data_root, s) for s in shards]
+        self.shards = shards
+        assert len(shards) > 0, f"no shards found for split {split}"
+        if master_process:
+            print(f"found {len(shards)} shards for split {split}")
+        self.reset()
+
+    def reset(self):
+        # restart cursor coordinates back to default baseline states
+        self.current_shard = 0
+        self.tokens = load_tokens(self.shards[self.current_shard])
+        self.current_position = self.B * self.T * self.process_rank
+    
+    def next_batch(self):
+        B, T = self.B, self.T
+        # slice out an extra token coordinate to establish lookahead targets
+        buf = self.tokens[self.current_position : self.current_position + B * T + 1]
+        x = (buf[:-1]).view(B, T) 
+        y = (buf[1:]).view(B, T) 
+        
+        # advance sequence coordinates forward across total active processes
+        self.current_position += B * T * self.num_processes
+        
+        # boundary loop detection to wrap safely over to the next shard asset
+        if self.current_position + B * T * self.num_processes + 1 > len(self.tokens):
+            self.current_shard = (self.current_shard + 1) % len(self.shards)
+            self.tokens = load_tokens(self.shards[self.current_shard])
+            self.current_position = B * T * self.process_rank
+        return x, y

src/utils/dataloader_masked.py

@@ -0,0 +1,54 @@
+import os
+import numpy as np
+import torch
+
+def load_data(filename):
+    # extract two column array containing tracking tokens and target masks
+    npt = np.load(filename)
+    npt = npt.astype(np.int32)
+    ptt = torch.tensor(npt, dtype=torch.long)
+    return ptt
+
+class DataLoaderMasked:
+    def __init__(self, B, T, process_rank, num_processes, split, master_process):
+        self.B = B
+        self.T = T
+        self.process_rank = process_rank
+        self.num_processes = num_processes
+        assert split in {'train', 'val'}
+        
+        data_root = "data/sft_dataset"
+        shards = os.listdir(data_root)
+        shards = [s for s in shards if split in s]
+        shards = sorted(shards)
+        shards = [os.path.join(data_root, s) for s in shards]
+        self.shards = shards
+        assert len(shards) > 0, f"no shards found for split {split}"
+        if master_process:
+            print(f"found {len(shards)} shards for split {split}")
+        self.reset()
+
+    def reset(self):
+        self.current_shard = 0
+        self.data = load_data(self.shards[self.current_shard])
+        self.current_position = self.B * self.T * self.process_rank
+    
+    def next_batch(self):
+        B, T = self.B, self.T
+        buf = self.data[self.current_position : self.current_position + B * T + 1]
+        
+        # partition raw text streams away from structural response mask columns
+        x = buf[:-1, 0].view(B, T) 
+        y = buf[1:, 0].view(B, T)  
+        m = buf[1:, 1].view(B, T)  
+        
+        # replace unmasked prompt context fields with pytorch ignore token index
+        y = y.clone()
+        y[m == 0] = -100
+        
+        self.current_position += B * T * self.num_processes
+        if self.current_position + B * T * self.num_processes + 1 > len(self.data):
+            self.current_shard = (self.current_shard + 1) % len(self.shards)
+            self.data = load_data(self.shards[self.current_shard])
+            self.current_position = B * T * self.process_rank
+        return x, y

src/utils/dataloader_ppo.py

@@ -0,0 +1,40 @@
+import os
+import numpy as np
+import torch
+
+def load_data(filename):
+    npt = np.load(filename)
+    npt = npt.astype(np.int32)
+    return torch.tensor(npt, dtype=torch.long)
+
+class DataLoaderPPO:
+    def __init__(self, B, process_rank, num_processes, split, master_process):
+        self.B = B
+        self.process_rank = process_rank
+        self.num_processes = num_processes
+        
+        data_root = "data/ppo_dataset"
+        shards = [os.path.join(data_root, s) for s in os.listdir(data_root) if split in s]
+        self.shards = sorted(shards)
+        assert len(self.shards) > 0, f"no shards found for split {split}"
+        
+        if master_process:
+            print(f"found {len(self.shards)} ppo shards for split {split}")
+        self.reset()
+
+    def reset(self):
+        self.current_shard = 0
+        self.data = load_data(self.shards[self.current_shard])
+        self.current_position = self.B * self.process_rank
+    
+    def next_batch(self):
+        B = self.B
+        # isolate the exact prompt token batch chunks out of static arrays
+        prompts = self.data[self.current_position : self.current_position + B]
+        
+        self.current_position += B * self.num_processes
+        if self.current_position + B * self.num_processes > len(self.data):
+            self.current_shard = (self.current_shard + 1) % len(self.shards)
+            self.data = load_data(self.shards[self.current_shard])
+            self.current_position = B * self.process_rank
+        return prompts

src/utils/dataloader_reward.py

@@ -0,0 +1,47 @@
+import os
+import numpy as np
+import torch
+
+def load_data(filename):
+    npt = np.load(filename)
+    npt = npt.astype(np.int32)
+    return torch.tensor(npt, dtype=torch.long)
+
+class DataLoaderReward:
+    def __init__(self, B, T, process_rank, num_processes, split, master_process):
+        self.B = B
+        self.T = T
+        self.process_rank = process_rank
+        self.num_processes = num_processes
+        assert split in {'train', 'val'}
+        
+        data_root = "data/rm_dataset"
+        shards = os.listdir(data_root)
+        shards = [s for s in shards if split in s]
+        shards = sorted(shards)
+        shards = [os.path.join(data_root, s) for s in shards]
+        self.shards = shards
+        assert len(shards) > 0, f"no shards found for split {split}"
+        if master_process:
+            print(f"found {len(shards)} reward shards for split {split}")
+        self.reset()
+
+    def reset(self):
+        self.current_shard = 0
+        self.data = load_data(self.shards[self.current_shard])
+        self.current_position = self.B * self.process_rank
+    
+    def next_batch(self):
+        B = self.B
+        buf = self.data[self.current_position : self.current_position + B]
+        
+        # split custom preference sequence records into winning and losing tensors
+        x_chosen = buf[:, 0, :]   
+        x_rejected = buf[:, 1, :] 
+        
+        self.current_position += B * self.num_processes
+        if self.current_position + B * self.num_processes > len(self.data):
+            self.current_shard = (self.current_shard + 1) % len(self.shards)
+            self.data = load_data(self.shards[self.current_shard])
+            self.current_position = B * self.process_rank
+        return x_chosen, x_rejected