Update tensor_server.py

tensor_server.py

@@ -151,6 +151,8 @@ def load_chunk(chunk: ModelChunk) -> torch.nn.Module:
                 self.config = config
                 self.start_offset = config.get('start_offset', 0)
                 self.size = config.get('size_bytes', 0)
+                # expose vocab_offset on the module for aggregator use
+                self.vocab_offset = int(config.get('vocab_offset', 0))
                 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
                 # In a real implementation, this would process the input
@@ -160,8 +162,10 @@ def load_chunk(chunk: ModelChunk) -> torch.nn.Module:
                 
         # Create and return the chunk buffer
         chunk_model = ChunkBuffer(chunk_file, chunk_config)
+        # Ensure the chunk_model.config is the up-to-date config (including any assigned offsets)
+        chunk_model.config = chunk_config
         print(f"[INFO] Loaded chunk {chunk.chunk_id} ({chunk_config.get('size_bytes', 0)} bytes) from {chunk.files[0]}")
-        
+
         return chunk_model
         
     except Exception as e:
@@ -206,6 +210,28 @@ async def load_model_chunk(chunk: ModelChunk):
         # Store the chunk metadata
         chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])
         state.chunk_configs = getattr(state, 'chunk_configs', {})
+
+        # Ensure a vocab_offset is present; if not, assign a non-overlapping offset
+        cfg = chunk.config or {}
+        if 'vocab_offset' not in cfg:
+            # Compute next available offset from existing registered chunks
+            max_end = 0
+            for existing in state.chunk_configs.values():
+                try:
+                    e_cfg = existing.config if hasattr(existing, 'config') else existing
+                    e_offset = int(e_cfg.get('vocab_offset', 0))
+                    e_shard = int(e_cfg.get('shard_dim', e_cfg.get('size', 1) or 1))
+                    max_end = max(max_end, e_offset + e_shard)
+                except Exception:
+                    continue
+
+            # If this chunk declares a shard_dim, use it; otherwise default to 1
+            shard_dim = int(cfg.get('shard_dim', cfg.get('size', 1) or 1))
+            cfg['vocab_offset'] = max_end
+            cfg['shard_dim'] = cfg.get('shard_dim', shard_dim)
+
+        # Store back the possibly-updated config
+        chunk.config = cfg
         state.chunk_configs[chunk.chunk_id] = chunk
         
         print(f"[INFO] Registered chunk {chunk.chunk_id} configuration")
@@ -239,6 +265,20 @@ async def upload_chunk_data(chunk_id: int, file: UploadFile = File(...)):
             
         # Now load the chunk
         chunk_model = load_chunk(chunk)
+        # Ensure the loaded module has the registered config (including vocab_offset)
+        try:
+            registered = getattr(state, 'chunk_configs', {}).get(chunk_id)
+            if registered is not None:
+                # registered is a ModelChunk; merge config into module
+                reg_cfg = registered.config or {}
+                if hasattr(chunk_model, 'config'):
+                    chunk_model.config.update(reg_cfg)
+                else:
+                    chunk_model.config = reg_cfg
+                # expose vocab_offset on module
+                chunk_model.vocab_offset = int(reg_cfg.get('vocab_offset', 0))
+        except Exception:
+            pass
         state.loaded_chunks[chunk_id] = chunk_model
         
         file_size = os.path.getsize(chunk_file)