Refactor/dp only remove tp pp ep sp - #17
Merged
Merged
Conversation
Narrow the framework to data parallelism (DDP / FSDP / DeepSpeed ZeRO) and remove tensor/pipeline/expert/sequence parallelism, which were never real-GPU validated and carried maintenance cost without a clear user. Removed: - builtin_plugins/distributed/model_parallel/ (tp_auto, tp_aware, sp, ep) and pipeline/ (schedules); recipes 3d_parallel.yaml, tp_ddp.yaml - ParallelSection tp/pp/ep/sp + tensor_parallel/pipeline; TPConfig/PipelineConfig - ModelParallelStrategy/PipelineSchedule protocols + registry categories - _runtime _build_model_parallel_strategy/_build_pipeline_schedule/_preflight_parallel - ParallelContext reduced to data-parallel only (1-D dp DeviceMesh) Bundled fixes (intermingled in the same files): - fix(distributed): broadcast run_dir from rank 0 so all ranks share one dir; a per-rank datetime.now() timestamp split ranks across sibling dirs when a launch straddled a one-second boundary. Gate the lineage store to rank 0. - fix(cli): destroy_process_group() on train/resume exit (NCCL teardown warning) - fix(recipes): quote mixed_precision "no" in fsdp/zero2 (YAML bool coercion) - docs: rewrite operations/distributed(.zh-CN) + recipes tables to DP-only scope Validated on a real single-node 4xH20 box: DDP + FSDP run clean, loss drops, all ranks share one run_dir. Full suite 4465 passed; ruff + mypy clean. CanGeng
ZeRO had no tests (not even CPU) and the lighttrain<->DeepSpeed seam broke on real multi-GPU. Fixes to make ZeRO-2 actually run: - zero.py: add train_micro_batch_size_per_gpu / gradient_accumulation_steps so DeepSpeed derives train_batch_size from the live world_size (the hardcoded train_batch_size=1 was rejected when world_size > 1). - zero.py: clip_grad_norm returns 0.0 when get_global_grad_norm() is None (lighttrain clips before optimizer_step; DeepSpeed only has the norm after engine.step()). - route zero_grad through GradSyncStrategy: DeepSpeedEngine.zero_grad() lacks the set_to_none kwarg and step() already zeroed grads. Add zero_grad to the protocol + ddp/fsdp/noop (set_to_none=True) and zero (no-op); the update loop calls grad_sync.zero_grad(optimizer) instead of optimizer.zero_grad directly. - recipes/nano_zero2.yaml: complete tiny_lm ZeRO-2 smoke recipe. Validated on a real single-node 4xH20 box: ZeRO-2 trains, loss drops, ranks exit clean. Full suite 4466 passed; ruff + mypy clean. CanGeng
safetensors save_model refuses to persist aliased storage. Under ZeRO-2 (DeepSpeed flattens params into one buffer) tiny_lm's tied tok_emb/lm_head became non-covering slices, so save_model raised and aborted crash / frozen / nan-repro diagnostic snapshots. Add save_model_safe (lighttrain/observability/diagnostics/_save.py): try save_model, and on RuntimeError clone the state dict to break the storage sharing and save the raw tensors (values preserved; tying metadata dropped — fine for a diagnostic snapshot). Route frozen_step / crash_bundle / nan_repro through it; correct the stale nan_repro comment that claimed save_model already handled tied weights. CanGeng
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.