diff --git a/Gruntfile.js b/Gruntfile.js
index 86b9377..7425fcc 100644
--- a/Gruntfile.js
+++ b/Gruntfile.js
@@ -57,15 +57,6 @@ module.exports = function(grunt) {
       },
     },
 
-    umd: {
-      all: {
-        src: "dist/browser-polyfill.js",
-        template: "unit",
-        globalAlias: "browser",
-        amdModuleId: "webextension-polyfill",
-      },
-    },
-
     babel: {
       minify: {
         options: {
@@ -78,6 +69,25 @@ module.exports = function(grunt) {
           "dist/browser-polyfill.min.js": "dist/browser-polyfill.js",
         },
       },
+      umd: {
+        options: {
+          babelrc: false,
+          comments: true,
+          plugins: [
+            ["transform-es2015-modules-umd", {
+              globals: {
+                "webextension-polyfill": "browser",
+              },
+              exactGlobals: true,
+            }],
+          ],
+          sourceMap: true,
+          moduleId: "webextension-polyfill",
+        },
+        files: {
+          "dist/browser-polyfill.js": "dist/browser-polyfill.js",
+        },
+      },
     },
 
     concat: {
@@ -91,10 +101,9 @@ module.exports = function(grunt) {
 
   grunt.loadNpmTasks("gruntify-eslint");
   grunt.loadNpmTasks("grunt-replace");
-  grunt.loadNpmTasks("grunt-umd");
   grunt.loadNpmTasks("grunt-coveralls");
   grunt.loadNpmTasks("grunt-contrib-concat");
   grunt.loadNpmTasks("grunt-babel");
 
-  grunt.registerTask("default", ["eslint", "replace", "umd", "babel:minify", "concat:license"]);
+  grunt.registerTask("default", ["eslint", "replace", "babel:umd", "babel:minify", "concat:license"]);
 };
diff --git a/package.json b/package.json
index ad0e0ce..9bc58c0 100644
--- a/package.json
+++ b/package.json
@@ -17,6 +17,7 @@
   },
   "homepage": "https://github.com/mozilla/webextension-polyfill",
   "devDependencies": {
+    "babel-plugin-transform-es2015-modules-umd": "^6.24.1",
     "babel-preset-babili": "0.0.10",
     "chai": "^3.5.0",
     "eslint": "3.9.1",
@@ -25,7 +26,6 @@
     "grunt-contrib-concat": "^1.0.1",
     "grunt-coveralls": "^1.0.1",
     "grunt-replace": "*",
-    "grunt-umd": "^2.4.0",
     "gruntify-eslint": "*",
     "istanbul-lib-instrument": "^1.1.3",
     "jsdom": "^9.6.0",