| 328 | self.arrow_writer = pa.ipc.new_stream(self.arrow_file, schema) |
| 329 | |
| 330 | def _save_step_data(self, keys, vals): |
| 331 | # Detach from computation graph and ensure contiguous memory |
| 332 | keys_tensor = keys.detach().contiguous() |
| 333 | vals_tensor = vals.detach().contiguous() |
| 334 | |
| 335 | # Convert to desired types |
| 336 | keys_tensor = keys_tensor.to(dtype=torch.float16) |
| 337 | vals_tensor = vals_tensor.to(dtype=torch.int32) |
| 338 | |
| 339 | # Gather from all processes |
| 340 | all_keys = self.accelerator.gather_for_metrics(keys_tensor) |
| 341 | all_vals = self.accelerator.gather_for_metrics(vals_tensor) |
| 342 | |
| 343 | shift = 0 if self.is_encoder_decoder else 1 |
| 344 | if shift == 1: |
| 345 | all_keys = all_keys[:, :-shift] |
| 346 | all_keys = all_keys.flatten(0, 1) # (batch * time, dim) |
| 347 | all_vals = all_vals[:, shift:].flatten(0, 1) # (batch * time) |
| 348 | |
| 349 | nonpad_mask = all_vals != -100 |
| 350 | all_keys = all_keys[nonpad_mask] |
| 351 | all_vals = all_vals[nonpad_mask] |
| 352 | |
| 353 | # Only main process writes to file |
| 354 | if self.process_index == 0: |
| 355 | # Convert tensors back to numpy arrays |
| 356 | all_keys_np = all_keys.cpu().numpy().astype(np.float16) |
| 357 | all_vals_np = all_vals.cpu().numpy().astype(np.int32) |
| 358 | |
| 359 | # Create Arrow arrays |
| 360 | keys_list = [all_keys_np[i] for i in range(all_keys_np.shape[0])] |
| 361 | keys_array = pa.array(keys_list, type=pa.list_(pa.float16(), self.dimension)) |
| 362 | vals_array = pa.array(all_vals_np, type=pa.int32()) |
| 363 | |
| 364 | # Create batch and write |
| 365 | batch = pa.RecordBatch.from_arrays([keys_array, vals_array], ['keys', 'vals']) |
| 366 | self.arrow_writer.write_batch(batch) |
| 367 | |
| 368 | logger.info(f"Main process: Flushed buffer, total saved: {self.dstore_idx + all_keys_np.shape[0]}") |
| 369 | self.dstore_idx += all_keys_np.shape[0] |
| 370 | |
| 371 | self.accelerator.wait_for_everyone() |
| 372 | |
| 373 | def break_into(self, model): |
| 374 | self.model = model |