/usr/local/lib/python3.6/site-packages/transformers/utils
NameSizeModeActions
__pycache__/-0755rm
doc.py388360644editdlrm
dummy_detectron2_objects.py3910644editdlrm
dummy_flax_objects.py248550644editdlrm
dummy_pt_objects.py1114430644editdlrm
dummy_pytorch_quantization_and_torch_objects.py24760644editdlrm
dummy_scatter_objects.py11380644editdlrm
dummy_sentencepiece_and_speech_objects.py3420644editdlrm
dummy_sentencepiece_and_tokenizers_objects.py3010644editdlrm
dummy_sentencepiece_objects.py43710644editdlrm
dummy_speech_objects.py3150644editdlrm
dummy_tf_objects.py504560644editdlrm
dummy_timm_and_vision_objects.py9030644editdlrm
dummy_timm_objects.py8050644editdlrm
dummy_tokenizers_objects.py81730644editdlrm
dummy_vision_objects.py34860644editdlrm
fx.py231860644editdlrm
fx_transformations.py121780644editdlrm
generic.py98200644editdlrm
hp_naming.py49710644editdlrm
hub.py429320644editdlrm
import_utils.py292880644editdlrm
logging.py95160644editdlrm
model_parallel_utils.py22760644editdlrm
notebook.py145240644editdlrm
sentencepiece_model_pb2.py396070644editdlrm
versions.py43780644editdlrm
__init__.py50550644editdlrm
Edit: /usr/local/lib/python3.6/site-packages/transformers/utils/model_parallel_utils.py (2276B)
# coding=utf-8 # Copyright 2020 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from math import ceil def assert_device_map(device_map, num_blocks): blocks = list(range(0, num_blocks)) device_map_blocks = [item for sublist in list(device_map.values()) for item in sublist] # Duplicate check duplicate_blocks = [] for i in device_map_blocks: if device_map_blocks.count(i) > 1 and i not in duplicate_blocks: duplicate_blocks.append(i) # Missing blocks missing_blocks = [i for i in blocks if i not in device_map_blocks] extra_blocks = [i for i in device_map_blocks if i not in blocks] if len(duplicate_blocks) != 0: raise ValueError( "Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device. These " "attention blocks were specified more than once: " + str(duplicate_blocks) ) if len(missing_blocks) != 0: raise ValueError( "There are attention blocks for this model that are not specified in the device_map. Add these attention " "blocks to a device on the device_map: " + str(missing_blocks) ) if len(extra_blocks) != 0: raise ValueError( "The device_map contains more attention blocks than this model has. Remove these from the device_map:" + str(extra_blocks) ) def get_device_map(n_layers, devices): """Returns a dictionary of layers distributed evenly across all devices.""" layers = list(range(n_layers)) n_blocks = int(ceil(n_layers / len(devices))) layers_list = list(layers[i : i + n_blocks] for i in range(0, n_layers, n_blocks)) return dict(zip(devices, layers_list))