fix: clarify strategy naming + init distributed process group for FSDP
This commit is contained in:
14
train.py
14
train.py
@@ -38,9 +38,9 @@ def train(config_path):
|
|||||||
errors = []
|
errors = []
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
# Strategy 1: QLoRA with FSDP (preferred)
|
# Strategy 1: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs)
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
print("\n[1/4] Trying: 4-bit QLoRA (FSDP)...")
|
print("\n[1/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...")
|
||||||
try:
|
try:
|
||||||
bnb_config = BitsAndBytesConfig(
|
bnb_config = BitsAndBytesConfig(
|
||||||
load_in_4bit=True,
|
load_in_4bit=True,
|
||||||
@@ -56,7 +56,7 @@ def train(config_path):
|
|||||||
trust_remote_code=True,
|
trust_remote_code=True,
|
||||||
low_cpu_mem_usage=True,
|
low_cpu_mem_usage=True,
|
||||||
)
|
)
|
||||||
print("✓ Success: QLoRA 4-bit loaded to CPU")
|
print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
errors.append(("QLoRA 4-bit", e))
|
errors.append(("QLoRA 4-bit", e))
|
||||||
print(f"✗ Failed: {e}")
|
print(f"✗ Failed: {e}")
|
||||||
@@ -156,6 +156,12 @@ def train(config_path):
|
|||||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||||
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
||||||
from functools import partial
|
from functools import partial
|
||||||
|
import torch.distributed as dist
|
||||||
|
|
||||||
|
# Initialize distributed process group (required for FSDP)
|
||||||
|
if not dist.is_initialized():
|
||||||
|
dist.init_process_group(backend="nccl")
|
||||||
|
print("✓ Distributed process group initialized")
|
||||||
|
|
||||||
def get_auto_wrap_policy(model):
|
def get_auto_wrap_policy(model):
|
||||||
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
|
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
|
||||||
@@ -175,7 +181,7 @@ def train(config_path):
|
|||||||
sync_module_states=True,
|
sync_module_states=True,
|
||||||
use_orig_params=True,
|
use_orig_params=True,
|
||||||
)
|
)
|
||||||
print("✓ Model wrapped with FSDP")
|
print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")
|
||||||
|
|
||||||
# Training arguments (no FSDP config - we handle it manually)
|
# Training arguments (no FSDP config - we handle it manually)
|
||||||
training_args = TrainingArguments(
|
training_args = TrainingArguments(
|
||||||
|
|||||||
Reference in New Issue
Block a user