Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

FSDP utils cleanup #854

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

FSDP utils cleanup #854

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5b9076a
use existing constant, clarify comment
kylesayrs Oct 19, 2024
d4c2119
Merge branch 'main' into kylesayrs/fsdp-wrapper-name
kylesayrs Oct 19, 2024
491496f
always import from accelerate since it is a dependency
kylesayrs Nov 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 20 additions & 19 deletions src/llmcompressor/utils/fsdp/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
from accelerate import Accelerator
try:
from accelerate import Accelerator
from torch.distributed.fsdp import FullyShardedDataParallel
from torch.distributed.fsdp._common_utils import TrainingState
from torch.distributed.fsdp._common_utils import FSDP_WRAPPED_MODULE, TrainingState
except ImportError:
FullyShardedDataParallel = None
Accelerator = None

from contextlib import nullcontext

Expand All @@ -14,22 +13,21 @@
"fix_fsdp_module_name",
]

FSDP_WRAPPER_NAME = "_fsdp_wrapped_module"


def summon_full_params_context(model, offload_to_cpu: bool = False):
if FullyShardedDataParallel is not None:
# avoid nested summon_full_param context
if (
hasattr(model, "training_state")
and model.training_state is TrainingState.SUMMON_FULL_PARAMS
):
return nullcontext()
return FullyShardedDataParallel.summon_full_params(
model, offload_to_cpu=offload_to_cpu
)
if FullyShardedDataParallel is None:
return nullcontext()

# do not call from within summon_full_param context
if (
hasattr(model, "training_state")
and model.training_state is TrainingState.SUMMON_FULL_PARAMS
):
return nullcontext()

return nullcontext()
return FullyShardedDataParallel.summon_full_params(
model, offload_to_cpu=offload_to_cpu
)


def main_process_first_context():
Expand All @@ -46,12 +44,15 @@ def main_process_first_context():
def fix_fsdp_module_name(name: str) -> str:
"""
Remove FSDP wrapper prefixes from a module name.
Accounts for scenario where FSDP_WRAPPER_NAME is
Accounts for scenario where FSDP_WRAPPED_MODULE is
at the end of the name, as well as in the middle.

:param name: name to strip
:return: stripped name
"""
return name.replace(FSDP_WRAPPER_NAME + ".", "").replace(
"." + FSDP_WRAPPER_NAME, ""
if FullyShardedDataParallel is None:
return name

return name.replace(FSDP_WRAPPED_MODULE + ".", "").replace(
"." + FSDP_WRAPPED_MODULE, ""
)
Loading