Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
5 changes: 4 additions & 1 deletion deepspeed/runtime/bf16_optimizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from collections import OrderedDict
import torch
import sys
import os
from deepspeed import comm as dist
from deepspeed.runtime.constants import PIPE_REPLICATED
Expand All @@ -21,7 +22,7 @@
is_model_parallel_parameter,
see_memory_usage)

from deepspeed.utils import link_hp_params
from deepspeed.utils import link_hp_params, fragment_address
from deepspeed.checkpoint import enable_universal_checkpoint
from deepspeed.checkpoint.constants import (DS_VERSION,
PARTITION_COUNT,
Expand All @@ -31,6 +32,8 @@
GROUP_PADDINGS,
PARAM_SLICE_MAPPINGS)

setattr(sys.modules[__name__], 'fragment_address', fragment_address)


class BF16_Optimizer(ZeROOptimizer):
def __init__(self,
Expand Down
2 changes: 1 addition & 1 deletion deepspeed/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,6 @@
from .init_on_device import OnDevice
from .groups import *
from .nvtx import instrument_w_nvtx
from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping
from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address
from .mixed_precision_linkage import link_hp_params
from deepspeed.runtime.dataloader import RepeatingLoader