Add Code World Model (CWM) #41199
Open
Add Code World Model (CWM) #41199
+1,656
−0
Conversation
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
Comment on lines
+144
to
+156
| config = self.model_tester.get_config() | ||
| model = CwmModel(config) | ||
| model.to(torch_device) | ||
| model.eval() | ||
|
|
||
| # input longer than sliding window | ||
| seq_length = config.sliding_window + 10 | ||
| input_ids = torch.randint(0, config.vocab_size, (1, seq_length), device=torch_device) | ||
|
|
||
| with torch.no_grad(): | ||
| outputs = model(input_ids) | ||
|
|
||
| self.assertEqual(outputs.last_hidden_state.shape, (1, seq_length, config.hidden_size)) |
There was a problem hiding this comment.
better if we can make an integration test and check that the generated ids are correct
Comment on lines
+184
to
+185
| # no errors | ||
| self.assertIsNotNone(outputs.last_hidden_state) |
There was a problem hiding this comment.
not clear what we are testing for here, last hidden state can never be None, no?
| def tearDown(self): | ||
| cleanup(torch_device, gc_collect=True) | ||
|
|
||
| def test_cwm_small_model_forward(self): |
There was a problem hiding this comment.
let's use actual model for all below tests and check generated token ids or logit values for a few important cases (sliding window, simple generation etc)
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, cwm |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Adds the Code World Model (CWM) - https://ai.meta.com/research/publications/cwm-an-open-weights-llm-for-research-on-code-generation-with-world-models/
High-level implementation details:
The model repos are:
Note that for VLLM compatibility, model
config.jsonstill refer toLlama3ForCausalLMand allamamodel_type— see example. vllm-project/vllm#25611 adds support mappingCwmForCausalLMto the Llama3 model class in VLLM since VLLM supportsLlama3+layer_typeswith local/global attention - see docs. The model type in theconfig.jsonwill be updated on HF (and the special automapping condition removed) once this PR is merged and a Transformers release has happened containing theCwmForCausalLMmodel class.@ArthurZucker, @zucchini-nlp
Supersedes #41188 due to some fork misery