[Bug]: 'weight' must be 2-D in upsert.

[DavidYanAnDe]

What happened?

When I use DeepSpeed stage-3 and upsert in collection of chromadb, it will output error like below, but once I remove DeepSpeed it work correctly, so is there any conflict between DeepSpeed and chromadb?

Versions

chromadb 0.6.3, python 3.10.16

Relevant log output

Exception has occurred: RuntimeError
Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 90, in wrapper
    return func(self, *args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 406, in _validate_and_prepare_upsert_request
    upsert_embeddings = self._embed_record_set(record_set=upsert_records)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 526, in _embed_record_set
    return self._embed(input=record_set[field])  # type: ignore[literal-required]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 539, in _embed
    return self._embedding_function(input=input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/types.py", line 466, in __call__
    result = call(self, input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/utils/embedding_functions/sentence_transformer_embedding_function.py", line 48, in __call__
    for embedding in self._model.encode(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 623, in encode
    out_features = self.forward(features, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 690, in forward
    input = module(input, **module_kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/models/Transformer.py", line 393, in forward
    output_states = self.auto_model(**trans_features, **kwargs, return_dict=False)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 1078, in forward
    embedding_output = self.embeddings(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 211, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 164, in forward
    return F.embedding(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/functional.py", line 2267, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 309, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/data/cpfs_0/ydw/Ovis/ovis/train/dataset/conversation_dataset_memory_bank.py", line 103, in __getitem__
    collection.upsert(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/Collection.py", line 335, in upsert
    upsert_request = self._validate_and_prepare_upsert_request(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 93, in wrapper
    raise type(e)(msg).with_traceback(e.__traceback__)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 90, in wrapper
    return func(self, *args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 406, in _validate_and_prepare_upsert_request
    upsert_embeddings = self._embed_record_set(record_set=upsert_records)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 526, in _embed_record_set
    return self._embed(input=record_set[field])  # type: ignore[literal-required]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 539, in _embed
    return self._embedding_function(input=input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/types.py", line 466, in __call__
    result = call(self, input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/utils/embedding_functions/sentence_transformer_embedding_function.py", line 48, in __call__
    for embedding in self._model.encode(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 623, in encode
    out_features = self.forward(features, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 690, in forward
    input = module(input, **module_kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/models/Transformer.py", line 393, in forward
    output_states = self.auto_model(**trans_features, **kwargs, return_dict=False)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 1078, in forward
    embedding_output = self.embeddings(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 211, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 164, in forward
    return F.embedding(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/functional.py", line 2267, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D in upsert.
  File "/data/cpfs_0/ydw/Ovis/ovis/train/train_memory_bank.py", line 227, in train
    trainer.train()
  File "/data/cpfs_0/ydw/Ovis/ovis/train/train_memory_bank.py", line 237, in <module>
    train()
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 90, in wrapper
    return func(self, *args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 406, in _validate_and_prepare_upsert_request
    upsert_embeddings = self._embed_record_set(record_set=upsert_records)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 526, in _embed_record_set
    return self._embed(input=record_set[field])  # type: ignore[literal-required]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 539, in _embed
    return self._embedding_function(input=input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/types.py", line 466, in __call__
    result = call(self, input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/utils/embedding_functions/sentence_transformer_embedding_function.py", line 48, in __call__
    for embedding in self._model.encode(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 623, in encode
    out_features = self.forward(features, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 690, in forward
    input = module(input, **module_kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/models/Transformer.py", line 393, in forward
    output_states = self.auto_model(**trans_features, **kwargs, return_dict=False)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 1078, in forward
    embedding_output = self.embeddings(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 211, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 164, in forward
    return F.embedding(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/functional.py", line 2267, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 309, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/data/cpfs_0/ydw/Ovis/ovis/train/dataset/conversation_dataset_memory_bank.py", line 103, in __getitem__
    collection.upsert(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/Collection.py", line 335, in upsert
    upsert_request = self._validate_and_prepare_upsert_request(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 93, in wrapper
    raise type(e)(msg).with_traceback(e.__traceback__)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 90, in wrapper
    return func(self, *args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 406, in _validate_and_prepare_upsert_request
    upsert_embeddings = self._embed_record_set(record_set=upsert_records)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 526, in _embed_record_set
    return self._embed(input=record_set[field])  # type: ignore[literal-required]
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/models/CollectionCommon.py", line 539, in _embed
    return self._embedding_function(input=input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/api/types.py", line 466, in __call__
    result = call(self, input)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/chromadb/utils/embedding_functions/sentence_transformer_embedding_function.py", line 48, in __call__
    for embedding in self._model.encode(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 623, in encode
    out_features = self.forward(features, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/SentenceTransformer.py", line 690, in forward
    input = module(input, **module_kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/sentence_transformers/models/Transformer.py", line 393, in forward
    output_states = self.auto_model(**trans_features, **kwargs, return_dict=False)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 1078, in forward
    embedding_output = self.embeddings(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 211, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 164, in forward
    return F.embedding(
  File "/usr/local/conda/envs/ovis_test/lib/python3.10/site-packages/torch/nn/functional.py", line 2267, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D in upsert.

[tazarov]

@DavidYanAnDe, thanks for raising this. I have a feeling this issue is related to the embedding model and/or its config. I noticed you hare using SentenceTransformers. What model are you using with ST?

[DavidYanAnDe]

@tazarov Just all-MiniLM-L6-v2, because there were some error for default onnxrunningtime, so I change to use sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="all-MiniLM-L6-v2"
) and collection = self.chroma_client.get_or_create_collection(str(i), embedding_function=self.sentence_transformer_ef)

[tazarov]

interesting. Can you tell me which version of ST do you have pip show sentence-transfomers?

[DavidYanAnDe]

I know it's a bert transformer in all-MiniLM-L6-v2, so is it some settings in deepspeed influence it during training, even it doesn't participant in?
deepspeed zero_3 config:
{
"fp16": {
"enabled": false
},
"bf16": {
"enabled": "auto"
},
"train_micro_batch_size_per_gpu": "auto",
"train_batch_size": "auto",
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"zero_optimization": {
"stage": 3,
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
}
}

[DavidYanAnDe]

@tazarov sentence-transformers is 3.3.1

[DavidYanAnDe]

The reason why I said conflict between DeepSpeed and chromadb/ST is once I remove the "--deepspeed" in config it will work normally. Sadly, none of us can escape when training MLLM. -.-