feat: FSDP training failures now fallback to next strategy
This commit is contained in:
16
train.py
16
train.py
@@ -176,6 +176,7 @@ def train(config_path):
|
||||
)
|
||||
|
||||
# Manually wrap model with FSDP before passing to trainer
|
||||
try:
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
||||
from functools import partial
|
||||
@@ -191,7 +192,6 @@ def train(config_path):
|
||||
return partial(
|
||||
transformer_auto_wrap_policy,
|
||||
transformer_layer_cls={Qwen3_5MoeDecoderLayer},
|
||||
min_num_params=1,
|
||||
)
|
||||
|
||||
# Wrap model with FSDP on CPU first
|
||||
@@ -247,6 +247,20 @@ def train(config_path):
|
||||
|
||||
print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
|
||||
|
||||
# Cleanup distributed process group
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
# Return early - success
|
||||
return
|
||||
except Exception as e:
|
||||
errors.append(("QLoRA 4-bit FSDP training", e))
|
||||
print(f"✗ Failed during FSDP training: {e}")
|
||||
# Cleanup distributed process group if initialized
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
# Fall through to next strategy
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Train LoRA adapter")
|
||||
|
||||
Reference in New Issue
Block a user