| 477 | |
| 478 | @staticmethod |
| 479 | def try_load_train_state(resume_from, optimizer, scaler, scheduler, fsdp_config): |
| 480 | if resume_from is not None and os.path.exists(resume_from): |
| 481 | if fsdp_config.sharding_strategy == "FULL_SHARD": |
| 482 | shard_index = dist.get_rank() |
| 483 | total_shards = dist.get_world_size() |
| 484 | elif fsdp_config.sharding_strategy == "HYBRID_SHARD": |
| 485 | shard_index = dist.get_rank() % fsdp_config.num_shard |
| 486 | total_shards = fsdp_config.num_shard |
| 487 | else: |
| 488 | raise NotImplementedError |
| 489 | |
| 490 | optimizer_state_dict_path = os.path.join( |
| 491 | resume_from, f"optimizer.{shard_index:05d}-of-{total_shards:05d}.pt" |
| 492 | ) |
| 493 | optimizer_state_dict = torch.load(optimizer_state_dict_path, map_location="cpu", weights_only=True) |
| 494 | optimizer.load_state_dict(optimizer_state_dict) |
| 495 | del optimizer_state_dict |
| 496 | |
| 497 | scaler_state_dict_path = os.path.join(resume_from, "scaler.pt") |
| 498 | scaler_state_dict = torch.load(scaler_state_dict_path, weights_only=True, map_location="cpu") |
| 499 | scaler.load_state_dict(scaler_state_dict) |
| 500 | del scaler_state_dict |
| 501 | |
| 502 | scheduler_state_dict_path = os.path.join(resume_from, "scheduler.pt") |
| 503 | scheduler_state_dict = torch.load(scheduler_state_dict_path, weights_only=True, map_location="cpu") |
| 504 | scheduler.load_state_dict(scheduler_state_dict) |
| 505 | del scheduler_state_dict |
| 506 | |
| 507 | train_steps = int(os.path.basename(os.path.normpath(resume_from))) + 1 |
| 508 | """ |
| 509 | data_status = [ |
| 510 | { |
| 511 | dataset_name: { |
| 512 | worker_id: [parquet_idx, row_group_id, row_idx], |
| 513 | }, |
| 514 | }, |
| 515 | ] |
| 516 | """ |
| 517 | data_status_path = os.path.join(resume_from, "data_status.pt") |
| 518 | if os.path.exists(data_status_path): |
| 519 | data_status = torch.load(data_status_path, weights_only=True, map_location="cpu") |
| 520 | local_rank = dist.get_rank() |
| 521 | if local_rank < len(data_status): |
| 522 | data_status = data_status[local_rank] |
| 523 | else: |
| 524 | data_status = None |
| 525 | else: |
| 526 | data_status = None |
| 527 | else: |
| 528 | train_steps = 0 |
| 529 | data_status = None |
| 530 | return optimizer, scaler, scheduler, train_steps, data_status |
| 531 | |
| 532 | |
| 533 | def grad_checkpoint_check_fn(module): |